#!/usr/bin/env python3

# Create annual Antarctic GRACE mass-anomaly bands from NetCDF.
#
# Processing:
#   - Read NetCDF from cwd/input
#   - Convert lwe_thickness to metres water equivalent where required
#   - Mask to Antarctic land/ice using the source land_mask
#   - Calculate anomalies relative to the mean of all 2002 observations
#   - Create cumulative or annual mean anomaly bands by year
#   - Reproject to a fixed 5 km EPSG:3031 grid using bilinear resampling
#   - Clip to ADD_land_mask.shp and apply masked Gaussian smoothing
#   - Write a multiband COG to cwd/output
#
# Install:
#   pip install numpy xarray rasterio scipy fiona netCDF4

from __future__ import annotations

from pathlib import Path

import fiona
import numpy as np
import rasterio
import xarray as xr
from rasterio.crs import CRS
from rasterio.features import rasterize
from rasterio.transform import from_origin
from rasterio.warp import Resampling, reproject
from scipy.ndimage import gaussian_filter


# =============================================================================
# CONFIGURATION
# =============================================================================

BASE_DIR = Path.cwd()
INPUT_DIR = BASE_DIR / "input"
OUTPUT_DIR = BASE_DIR / "output"

VAR_NAME = "lwe_thickness"
MASK_VAR = "land_mask"

BASELINE_YEAR = 2002
CROP_SOUTH_OF_LAT = -60.0

DISPLAY_MASK_VECTOR = "ADD_land_mask.shp"

OUT_FILENAME = "GRACE_Antarctica_mass_anomaly_mwe_2002-2025.tif"
NODATA = -9999.0

# "cumulative" = mean anomaly from 2002 through each output year
# "annual"     = mean anomaly within each individual year
YEAR_MODE = "cumulative"

# Fixed Antarctic EPSG:3031 output grid.
DST_CRS = CRS.from_epsg(3031)
DST_RES_M = 5000.0
DST_XMIN = -3330000.0
DST_YMIN = -3330000.0
DST_XMAX = 3330000.0
DST_YMAX = 3330000.0

# 20 pixels × 5 km = ~100 km Gaussian sigma.
GAUSSIAN_SIGMA_PIXELS = 20.0

COG_PROFILE = {
    "driver": "COG",
    "compress": "DEFLATE",
    "predictor": 2,
    "zlevel": 9,
    "blocksize": 512,
    "overview_resampling": "average",
    "bigtiff": "if_safer",
}


# =============================================================================
# INPUT / COORDINATES
# =============================================================================

def find_first_nc() -> Path:
    files = sorted(INPUT_DIR.glob("*.nc"))
    if not files:
        raise FileNotFoundError(f"No NetCDF files found in: {INPUT_DIR}")
    return files[0]


def pick_lat_lon_names(ds: xr.Dataset) -> tuple[str, str]:
    lat_name = next(
        (n for n in ("lat", "latitude", "y") if n in ds.coords or n in ds.variables),
        None,
    )
    lon_name = next(
        (n for n in ("lon", "longitude", "x") if n in ds.coords or n in ds.variables),
        None,
    )

    if lat_name is None or lon_name is None:
        raise KeyError(
            "Could not identify latitude/longitude coordinates. "
            f"Coords: {list(ds.coords)} | Variables: {list(ds.variables)}"
        )

    return lat_name, lon_name


def normalize_dataset_coords(
    ds: xr.Dataset,
    lat_name: str,
    lon_name: str,
) -> xr.Dataset:
    """Shift longitude to -180..180 where required and sort latitude north-to-south."""

    lon = ds[lon_name].values
    if np.nanmax(lon) > 180.0:
        lon = ((lon + 180.0) % 360.0) - 180.0
        ds = ds.assign_coords({lon_name: lon}).sortby(lon_name)

    lat = ds[lat_name].values
    if lat[0] < lat[-1]:
        ds = ds.sortby(lat_name, ascending=False)

    return ds


def crop_dataset_south(
    ds: xr.Dataset,
    lat_name: str,
) -> xr.Dataset:
    return ds.where(ds[lat_name] <= CROP_SOUTH_OF_LAT, drop=True)


# =============================================================================
# UNITS / MASKS
# =============================================================================

def convert_to_mwe_meters(
    da: xr.DataArray,
) -> tuple[xr.DataArray, str]:
    """Convert source values to metres water equivalent where units permit."""

    units = str(da.attrs.get("units", "")).strip().lower()
    da = da.astype(np.float32)

    if "cent" in units or units.startswith("cm") or "cm" in units:
        return da / 100.0, "m"

    if "mill" in units or units.startswith("mm") or "mm" in units:
        return da / 1000.0, "m"

    if units in {"m", "meter", "meters", "metre", "metres"}:
        return da, "m"

    sample = da.isel(time=0) if "time" in da.dims else da
    vmax = float(np.nanpercentile(np.abs(sample.values), 99))

    if 20.0 < vmax < 2000.0:
        print(
            f"[WARNING] Units '{units}' not recognised; "
            "magnitude suggests cm, converting /100"
        )
        return da / 100.0, "m"

    if vmax >= 2000.0:
        print(
            f"[WARNING] Units '{units}' not recognised; "
            "magnitude suggests mm, converting /1000"
        )
        return da / 1000.0, "m"

    print(f"[WARNING] Units '{units}' not recognised; values left unchanged")
    return da, units or "unknown"


def build_keep_mask(
    ds: xr.Dataset,
    lat_name: str,
    lon_name: str,
) -> np.ndarray:

    if MASK_VAR not in ds.variables:
        raise KeyError(
            f"Mask variable '{MASK_VAR}' not found. "
            f"Variables: {list(ds.variables)}"
        )

    mask = ds[MASK_VAR].transpose(lat_name, lon_name).values
    return np.isfinite(mask) & (mask != 0)


# =============================================================================
# GEOREFERENCING
# =============================================================================

def build_source_georef(
    lat_vals: np.ndarray,
    lon_vals: np.ndarray,
):
    if lat_vals.size < 2 or lon_vals.size < 2:
        raise ValueError("Need at least two latitude and longitude coordinates")

    xres = float(abs(lon_vals[1] - lon_vals[0]))
    yres = float(abs(lat_vals[1] - lat_vals[0]))

    west = float(lon_vals.min() - xres / 2.0)
    north = float(lat_vals.max() + yres / 2.0)

    return (
        CRS.from_epsg(4326),
        from_origin(west, north, xres, yres),
        xres,
        yres,
    )


def build_target_georef():
    width = int(np.ceil((DST_XMAX - DST_XMIN) / DST_RES_M))
    height = int(np.ceil((DST_YMAX - DST_YMIN) / DST_RES_M))

    transform = from_origin(
        DST_XMIN,
        DST_YMAX,
        DST_RES_M,
        DST_RES_M,
    )

    return transform, width, height


# =============================================================================
# DISPLAY MASK / SMOOTHING
# =============================================================================

def rasterize_display_mask(
    vector_path: Path,
    transform,
    width: int,
    height: int,
) -> np.ndarray:

    if not vector_path.exists():
        raise FileNotFoundError(f"Display mask not found: {vector_path}")

    with fiona.open(vector_path) as src:
        if len(src) == 0:
            raise ValueError(f"Display mask contains no features: {vector_path}")

        if src.crs_wkt:
            vector_crs = CRS.from_wkt(src.crs_wkt)
        elif src.crs:
            vector_crs = CRS(src.crs)
        else:
            vector_crs = None

        if vector_crs is not None and vector_crs != DST_CRS:
            raise ValueError(
                f"Display mask CRS ({vector_crs}) does not match "
                f"output CRS ({DST_CRS}); provide it in EPSG:3031"
            )

        shapes = [
            (feature["geometry"], 1)
            for feature in src
            if feature["geometry"] is not None
        ]

    return rasterize(
        shapes,
        out_shape=(height, width),
        transform=transform,
        fill=0,
        default_value=1,
        dtype="uint8",
        all_touched=False,
    ).astype(bool)


def gaussian_smooth_masked(
    band: np.ndarray,
    keep_mask: np.ndarray,
) -> np.ndarray:

    valid = keep_mask & np.isfinite(band) & (band != NODATA)

    if not np.any(valid):
        return np.full(band.shape, np.float32(NODATA), dtype=np.float32)

    values = np.where(valid, band, 0.0).astype(np.float32)
    weights = valid.astype(np.float32)

    smooth_values = gaussian_filter(
        values,
        sigma=GAUSSIAN_SIGMA_PIXELS,
        mode="nearest",
    )

    smooth_weights = gaussian_filter(
        weights,
        sigma=GAUSSIAN_SIGMA_PIXELS,
        mode="nearest",
    )

    with np.errstate(invalid="ignore", divide="ignore"):
        smooth = smooth_values / smooth_weights

    smooth = np.where(
        (smooth_weights > 0) & keep_mask,
        smooth,
        np.nan,
    )

    return np.where(
        np.isfinite(smooth),
        smooth,
        np.float32(NODATA),
    ).astype(np.float32)


# =============================================================================
# MAIN
# =============================================================================

def main() -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    nc_path = find_first_nc()
    print(f"[INPUT] {nc_path.relative_to(BASE_DIR)}")

    ds = xr.open_dataset(
        nc_path,
        decode_cf=True,
        mask_and_scale=True,
    )

    if VAR_NAME not in ds.variables:
        raise KeyError(
            f"Variable '{VAR_NAME}' not found. "
            f"Variables: {list(ds.variables)}"
        )

    lat_name, lon_name = pick_lat_lon_names(ds)

    ds = normalize_dataset_coords(ds, lat_name, lon_name)
    ds = crop_dataset_south(ds, lat_name)

    if ds.sizes.get(lat_name, 0) == 0 or ds.sizes.get(lon_name, 0) == 0:
        raise ValueError(
            f"Crop south of {CROP_SOUTH_OF_LAT}° produced an empty grid"
        )

    da = ds[VAR_NAME].transpose("time", lat_name, lon_name)

    if not np.issubdtype(da["time"].dtype, np.datetime64):
        ds = xr.decode_cf(ds)
        ds = normalize_dataset_coords(ds, lat_name, lon_name)
        ds = crop_dataset_south(ds, lat_name)
        da = ds[VAR_NAME].transpose("time", lat_name, lon_name)

        if not np.issubdtype(da["time"].dtype, np.datetime64):
            raise ValueError(
                f"Could not decode time coordinate: {da['time'].dtype}"
            )

    # Convert to metres water equivalent.
    da_mwe, out_units = convert_to_mwe_meters(da)

    if out_units != "m":
        print(f"[WARNING] Output units will be '{out_units}', not metres")

    # Apply source land/ice mask before baseline and anomaly calculations.
    keep_mask = build_keep_mask(ds, lat_name, lon_name)
    da_mwe = da_mwe.where(keep_mask)

    # Baseline = mean of all observations within BASELINE_YEAR.
    years_da = da_mwe["time"].dt.year
    baseline_count = int((years_da == BASELINE_YEAR).sum().values)

    print(f"[BASELINE] {BASELINE_YEAR}: {baseline_count} timestep(s)")

    if baseline_count == 0:
        raise ValueError(f"No observations found for baseline year {BASELINE_YEAR}")

    baseline = da_mwe.where(
        years_da == BASELINE_YEAR
    ).mean(
        dim="time",
        skipna=True,
    )

    if baseline.isnull().all():
        raise ValueError(
            f"Baseline year {BASELINE_YEAR} produced only NoData"
        )

    anomaly = da_mwe - baseline

    # Build annual output bands.
    years = np.sort(
        np.unique(da_mwe["time"].dt.year.values.astype(int))
    )

    if years.size == 0:
        raise ValueError("No valid years found")

    bands = []
    out_years = []

    for year in years:
        if YEAR_MODE == "cumulative":
            band = anomaly.where(
                anomaly["time"].dt.year <= year
            ).mean(
                dim="time",
                skipna=True,
            )

        elif YEAR_MODE == "annual":
            band = anomaly.where(
                anomaly["time"].dt.year == year
            ).mean(
                dim="time",
                skipna=True,
            )

        else:
            raise ValueError(
                f"YEAR_MODE must be 'cumulative' or 'annual': {YEAR_MODE}"
            )

        bands.append(band)
        out_years.append(int(year))

    stack = xr.concat(
        bands,
        dim="year",
    ).assign_coords(
        year=("year", out_years)
    )

    src_data = stack.values.astype(np.float32)
    src_data = np.where(
        np.isfinite(src_data),
        src_data,
        np.float32(NODATA),
    )

    # Georeferencing.
    lat_vals = ds[lat_name].values
    lon_vals = ds[lon_name].values

    src_crs, src_transform, xres_deg, yres_deg = build_source_georef(
        lat_vals,
        lon_vals,
    )

    dst_transform, dst_width, dst_height = build_target_georef()

    display_mask = rasterize_display_mask(
        INPUT_DIR / DISPLAY_MASK_VECTOR,
        dst_transform,
        dst_width,
        dst_height,
    )

    # Reproject and smooth each annual band.
    dst_data = np.full(
        (len(out_years), dst_height, dst_width),
        np.float32(NODATA),
        dtype=np.float32,
    )

    for i, year in enumerate(out_years):
        print(f"[PROCESS] {year} ({i + 1}/{len(out_years)})")

        band = np.full(
            (dst_height, dst_width),
            np.float32(NODATA),
            dtype=np.float32,
        )

        reproject(
            source=src_data[i],
            destination=band,
            src_transform=src_transform,
            src_crs=src_crs,
            dst_transform=dst_transform,
            dst_crs=DST_CRS,
            src_nodata=np.float32(NODATA),
            dst_nodata=np.float32(NODATA),
            resampling=Resampling.bilinear,
        )

        band = np.where(
            display_mask,
            band,
            np.float32(NODATA),
        )

        band = gaussian_smooth_masked(
            band,
            display_mask,
        )

        dst_data[i] = np.where(
            display_mask,
            band,
            np.float32(NODATA),
        )

    # Write multiband COG.
    output_path = OUTPUT_DIR / OUT_FILENAME

    profile = {
        "dtype": "float32",
        "count": len(out_years),
        "height": dst_height,
        "width": dst_width,
        "crs": DST_CRS,
        "transform": dst_transform,
        "nodata": np.float32(NODATA),
        **COG_PROFILE,
    }

    sigma_km = GAUSSIAN_SIGMA_PIXELS * DST_RES_M / 1000.0

    print(f"[GRID] EPSG:3031 | {DST_RES_M / 1000:g} km")
    print(f"[SOURCE] EPSG:4326 | {xres_deg:g}° × {yres_deg:g}°")
    print(f"[MODE] {YEAR_MODE}")
    print(f"[SMOOTH] sigma={GAUSSIAN_SIGMA_PIXELS:g} px (~{sigma_km:g} km)")
    print(f"[YEARS] {out_years[0]}–{out_years[-1]} ({len(out_years)} bands)")
    print(f"[OUTPUT] {output_path.relative_to(BASE_DIR)}")

    with rasterio.open(output_path, "w", **profile) as dst:
        for band_num, year in enumerate(out_years, start=1):
            dst.write(dst_data[band_num - 1], band_num)
            dst.set_band_description(
                band_num,
                f"mwe_change_{year}",
            )

        dst.update_tags(
            SOURCE_VAR=VAR_NAME,
            ANALYSIS_MASK_VAR=MASK_VAR,
            DISPLAY_MASK_VECTOR=DISPLAY_MASK_VECTOR,
            CROP_SOUTH_OF_LAT=str(CROP_SOUTH_OF_LAT),
            BASELINE=f"mean({BASELINE_YEAR} monthly {VAR_NAME})",
            PRODUCT=(
                f"{YEAR_MODE}_mean_anomaly_per_year_to_3031_"
                "bilinear_gaussian_vector_displaymask"
            ),
            UNITS="m",
            NODATA=str(NODATA),
            OUTPUT_RES_M=str(DST_RES_M),
            GAUSSIAN_SIGMA_PIXELS=str(GAUSSIAN_SIGMA_PIXELS),
            GAUSSIAN_SIGMA_KM=f"{sigma_km:.1f}",
        )

    ds.close()

    print("[OK] Done")


if __name__ == "__main__":
    main()