"""
Samples ESA WorldCover's Sentinel-2 NDVI yearly percentile composite
(p10/p50/p90) at station points, via the public S3 COGs.

Unlike fetch_worldcover_landcover.py, this does NOT hand-compute the
tile ID and guess the S3 key pattern -- VITO provides an authoritative
grid file for exactly this purpose (found directly, not guessed):
    https://esa-worldcover.s3.eu-central-1.amazonaws.com/esa_worldcover_grid_composites.fgb
This is a FlatGeobuf containing one feature per 1x1 degree composite
tile, with (per VITO's own documentation) the real S3 path/URL for
each tile as an attribute -- reading it means never guessing a file
naming convention, which is exactly the class of guess that failed for
the IDPR WMS layer name earlier in this session.

The composites use a 1x1 degree grid, NOT the 3x3 degree grid the
landcover classification map uses -- so unlike fetch_worldcover_
landcover.py's single hardcoded tile, station coordinates spanning
more than 1 degree in either direction genuinely need more than one
tile here. This script looks up each station's tile individually
rather than assuming one tile covers everyone.

NEEDS: geopandas (to read the grid FlatGeobuf) and rasterio (to sample
the COGs). CAVEAT: the grid file's actual column names/schema are not
independently verified here (no network access to inspect it directly)
-- the script prints all columns and searches for a plausible URL/path
column rather than assuming an exact name, and reports clearly if it
can't find one.

Usage:
    python -m scripts.fetch_worldcover_ndvi --check
    python -m scripts.fetch_worldcover_ndvi --stations datasets/station_elevations.csv
"""
import argparse
from pathlib import Path

import pandas as pd

GRID_URL = "https://esa-worldcover.s3.eu-central-1.amazonaws.com/esa_worldcover_grid_composites.fgb"

# The grid file's tile URLs use the s3:// scheme (confirmed against real
# output: s3://esa-worldcover-s2/ndvi/2020/N48/...) -- GDAL's S3 driver
# tries to SIGN requests with real AWS credentials by default even
# though this bucket is fully public, unlike the HTTPS URL
# fetch_worldcover_landcover.py used (a plain HTTPS GET needs no
# signing at all, which is why that one worked without this). This
# environment variable is the standard fix for reading a public S3
# bucket without needing real credentials.
import os
os.environ.setdefault("AWS_NO_SIGN_REQUEST", "YES")


def load_grid():
    import geopandas as gpd
    print(f"Loading tile grid from {GRID_URL} ...")
    grid = gpd.read_file(GRID_URL)
    print(f"  {len(grid)} tile(s), columns: {list(grid.columns)}")
    return grid


def find_url_column(grid) -> str:
    """
    The grid's real column name for the tile's S3 path isn't
    independently confirmed -- search for a plausible one rather than
    assume, and fail loudly (not silently) if nothing matches.
    """
    candidates = [c for c in grid.columns if any(
        kw in c.lower() for kw in ("url", "href", "path", "s3", "ndvi", "product", "file")
    )]
    if not candidates:
        raise ValueError(
            f"No column looks like a tile URL/path among {list(grid.columns)}. "
            f"Inspect the grid file's real schema directly (e.g. print(grid.head()) "
            f"in a notebook) and adjust find_url_column accordingly."
        )
    print(f"  candidate URL/path column(s): {candidates} -- using {candidates[0]!r}")
    return candidates[0]


def check_access() -> bool:
    try:
        import geopandas  # noqa: F401
        import rasterio  # noqa: F401
    except ImportError as e:
        print(f"Missing dependency: {e}. pip install geopandas rasterio --break-system-packages")
        return False

    try:
        grid = load_grid()
        url_col = find_url_column(grid)
        print(f"  Sample values from {url_col!r}: {grid[url_col].head(3).tolist()}")
        return True
    except Exception as e:
        print(f"  FAILED: {e}")
        return False


def fetch_for_stations(stations_path: Path, output_path: Path) -> None:
    import geopandas as gpd
    import rasterio
    from shapely.geometry import Point

    stations = pd.read_csv(stations_path)
    lat_col = "latitude" if "latitude" in stations.columns else "lat"
    lon_col = "longitude" if "longitude" in stations.columns else "lon"

    grid = load_grid()
    url_col = find_url_column(grid)

    station_points = gpd.GeoDataFrame(
        stations,
        geometry=[Point(lon, lat) for lon, lat in zip(stations[lon_col], stations[lat_col])],
        crs=grid.crs,
    )
    joined = gpd.sjoin(station_points, grid[[url_col, "geometry"]], how="left", predicate="within")

    rows = []
    for tile_url, group in joined.groupby(url_col):
        if pd.isna(tile_url):
            for _, row in group.iterrows():
                rows.append({"station_code": row["station_code"], "ndvi_p10": None,
                             "ndvi_p50": None, "ndvi_p90": None})
            continue
        print(f"  opening {tile_url} for {len(group)} station(s)...")
        with rasterio.Env(AWS_NO_SIGN_REQUEST="YES"), rasterio.open(tile_url) as src:
            coords = list(zip(group[lon_col], group[lat_col]))
            values = list(src.sample(coords))
        for (_, row), val in zip(group.iterrows(), values):
            # NDVI percentile composite is documented as 3 bands: p90, p50, p10
            rows.append({"station_code": row["station_code"],
                         "ndvi_p90": float(val[0]) if len(val) > 0 else None,
                         "ndvi_p50": float(val[1]) if len(val) > 1 else None,
                         "ndvi_p10": float(val[2]) if len(val) > 2 else None})

    out = pd.DataFrame(rows)
    out.to_csv(output_path, index=False)
    print(f"Saved {len(out)} stations' NDVI values to {output_path}")
    print(out.describe())


def main() -> None:
    parser = argparse.ArgumentParser(description="Sample ESA WorldCover NDVI at station points")
    parser.add_argument("--check", action="store_true")
    parser.add_argument("--stations", type=Path, default=Path("datasets/station_elevations.csv"))
    parser.add_argument("--output", type=Path, default=Path("datasets/worldcover_ndvi.csv"))
    args = parser.parse_args()

    if args.check:
        check_access()
        return

    fetch_for_stations(args.stations, args.output)


if __name__ == "__main__":
    main()