"""Read AlphaEarth Foundations annual satellite embeddings for the series grid from Source Cooperative.

The public COG mirror (s3://us-west-2.opendata.source.coop/tge-labs/aef/, HTTPS at data.source.coop) needs no
login. Each COG is 8192 x 8192 pixels in a UTM zone, 64 int8 bands (A00-A63), stored "bottom-up" (positive
y pixel size). We read just the window over the study area, flip it to the usual top-down order, reproject
it onto the shared 10 m British National Grid with nearest-neighbour resampling and cache it as one GeoTIFF.
"""
from pathlib import Path

import numpy as np
import pandas as pd
import pyarrow.parquet as pq
import rasterio
import requests
from affine import Affine
from pyproj import Transformer
from rasterio.warp import Resampling, reproject
from rasterio.windows import Window

BASE = "https://data.source.coop/tge-labs/aef/v1/annual"
INDEX_URL = f"{BASE}/aef_index.parquet"
NODATA = -128
GDAL_OPTS = {"GDAL_DISABLE_READDIR_ON_OPEN": "EMPTY_DIR", "GDAL_HTTP_MULTIRANGE": "YES",
             "GDAL_HTTP_MERGE_CONSECUTIVE_RANGES": "YES", "GDAL_HTTP_MAX_RETRY": "5", "GDAL_HTTP_RETRY_DELAY": "2",
             "VSI_CACHE": "TRUE", "CPL_VSIL_CURL_ALLOWED_EXTENSIONS": ".tiff"}


def dequantise(raw):
    """int8 -> float32 in [-1, 1]: (v / 127.5)^2 * sign(v); -128 (no data) -> NaN."""
    v = raw.astype(np.float32)
    out = (v / 127.5) ** 2 * np.sign(v)
    out[raw == NODATA] = np.nan
    return out


def find_tiles(bbox_wgs84, year, cache_dir):
    """Rows of the AEF index (one per COG) that intersect a lon/lat box for one year."""
    f = Path(cache_dir) / "aef_index.parquet"
    if not f.exists():
        f.parent.mkdir(parents=True, exist_ok=True)
        with requests.get(INDEX_URL, stream=True, timeout=300) as r:   # 78 MB, downloaded once
            r.raise_for_status()
            with open(f, "wb") as fh:
                for chunk in r.iter_content(1 << 20):
                    fh.write(chunk)
    cols = ["path", "crs", "year", "utm_zone", "utm_west", "utm_south", "utm_east", "utm_north",
            "wgs84_west", "wgs84_south", "wgs84_east", "wgs84_north"]
    df = pq.read_table(f, columns=cols, filters=[("year", "=", year)]).to_pandas()
    w, s, e, n = bbox_wgs84
    hit = (df.wgs84_west < e) & (df.wgs84_east > w) & (df.wgs84_south < n) & (df.wgs84_north > s)
    df = df[hit].copy()
    df["url"] = df.path.str.replace("s3://us-west-2.opendata.source.coop/tge-labs/aef/v1/annual",
                                    BASE, regex=False)
    return df.reset_index(drop=True)


def load_embeddings(tiles, dst_shape, dst_transform, dst_crs, margin=200):
    """Mosaic the AEF tiles onto a target grid. Returns int8 (64, rows, cols), NODATA where not covered."""
    rows, cols = dst_shape
    out = np.full((64, rows, cols), NODATA, dtype=np.int8)
    x0, y1 = dst_transform * (0, 0)
    x1, y0 = dst_transform * (cols, rows)
    with rasterio.Env(**GDAL_OPTS):
        for t in tiles.itertuples():
            to_utm = Transformer.from_crs(dst_crs, t.crs, always_xy=True)
            xs, ys = to_utm.transform([x0, x1, x0, x1], [y0, y0, y1, y1])
            west, east = max(min(xs) - margin, t.utm_west), min(max(xs) + margin, t.utm_east)
            south, north = max(min(ys) - margin, t.utm_south), min(max(ys) + margin, t.utm_north)
            if west >= east or south >= north:
                continue
            with rasterio.open(f"/vsicurl/{t.url}") as src:
                tr = src.transform                       # bottom-up: origin at the south-west corner, +10 m in y
                c0, c1 = int((west - tr.c) // tr.a), int(np.ceil((east - tr.c) / tr.a))
                r0, r1 = int((south - tr.f) // tr.e), int(np.ceil((north - tr.f) / tr.e))
                win = Window(c0, r0, c1 - c0, r1 - r0)
                raw = src.read(window=win)[:, ::-1, :]   # flip rows so that row 0 is the northern edge
            src_tr = Affine(tr.a, 0, tr.c + c0 * tr.a, 0, -tr.e, tr.f + r1 * tr.e)
            part = np.full_like(out, NODATA)
            reproject(raw, part, src_transform=src_tr, src_crs=t.crs, src_nodata=NODATA,
                      dst_transform=dst_transform, dst_crs=dst_crs, dst_nodata=NODATA,
                      resampling=Resampling.nearest)
            fill = (out[0] == NODATA) & (part[0] != NODATA)
            out[:, fill] = part[:, fill]
    return out


def tile_table(tiles):
    return pd.DataFrame({"file": tiles.url.str.split("/").str[-1], "zone": tiles.utm_zone,
                         "utm_west": tiles.utm_west, "utm_east": tiles.utm_east})