#!/usr/bin/env python3

# Create a 50 km gridded summary of Antarctic iceberg distribution.
#
# Bands:
#   1 = mean_count
#   2 = mean_area
#   3 = mean_mass
#   4 = SD_count
#   5 = SD_area
#   6 = SD_mass
#
# Means and SDs are calculated across annual values for 2018-2023.
# Years with no icebergs in a cell contribute zero.
# Cells with no iceberg records in any year are NoData (-999).
#
# Input:
#   cwd/output/Antarctic_icebergs_2018-2023.shp
#
# Output:
#   cwd/output/iceberg_metric_means_2018-2023.tif
#
# Install:
#   pip install GDAL numpy pandas geopandas pyproj pyogrio

from pathlib import Path
import math

import geopandas as gpd
import numpy as np
import pandas as pd
from osgeo import gdal, osr
from pyproj import Transformer

gdal.UseExceptions()


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

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

INPUT_FILE = OUTPUT_DIR / "Antarctic_icebergs_2018-2023.shp"
OUTPUT_FILE = OUTPUT_DIR / "iceberg_mean_metrics_50km_2018-2023.tif"

YEAR_FIELD = "year"
LAT_FIELD = "lat"
LON_FIELD = "lon"
AREA_FIELD = "area_km2"
MASS_FIELD = "mass_gt"

YEARS = list(range(2018, 2024))

EPSG = 3031
CELL_SIZE = 50_000
NODATA = -999.0


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

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

    if not INPUT_FILE.exists():
        raise FileNotFoundError(INPUT_FILE)

    print(f"[READ] {INPUT_FILE.name}")

    gdf = gpd.read_file(INPUT_FILE)

    required = [YEAR_FIELD, LAT_FIELD, LON_FIELD, AREA_FIELD, MASS_FIELD]
    missing = [f for f in required if f not in gdf.columns]

    if missing:
        raise ValueError(
            f"Missing field(s): {', '.join(missing)}\n"
            f"Available fields: {list(gdf.columns)}"
        )

    for field in required:
        gdf[field] = pd.to_numeric(gdf[field], errors="coerce")

    bad = gdf[required].isna().any(axis=1)

    if bad.any():
        raise ValueError(
            f"{bad.sum():,} record(s) contain invalid "
            "year/lat/lon/area/mass values."
        )

    if ((gdf[LAT_FIELD] < -90) | (gdf[LAT_FIELD] > 90)).any():
        raise ValueError("Invalid latitude values found.")

    if ((gdf[LON_FIELD] < -180) | (gdf[LON_FIELD] > 180)).any():
        raise ValueError("Invalid longitude values found.")

    if (gdf[AREA_FIELD] < 0).any():
        raise ValueError("Negative iceberg areas found.")

    if (gdf[MASS_FIELD] < 0).any():
        raise ValueError("Negative iceberg masses found.")

    gdf = gdf[gdf[YEAR_FIELD].isin(YEARS)].copy()

    present_years = sorted(gdf[YEAR_FIELD].astype(int).unique())

    if present_years != YEARS:
        raise ValueError(f"Expected years {YEARS}, found {present_years}")

    print(f"       {len(gdf):,} icebergs")

    # Transform supplied iceberg locations to EPSG:3031.
    transformer = Transformer.from_crs(
        "EPSG:4326",
        f"EPSG:{EPSG}",
        always_xy=True,
    )

    x, y = transformer.transform(
        gdf[LON_FIELD].to_numpy(float),
        gdf[LAT_FIELD].to_numpy(float),
    )

    x = np.asarray(x)
    y = np.asarray(y)

    # Snap output extent to the 50 km grid.
    xmin = math.floor(x.min() / CELL_SIZE) * CELL_SIZE
    xmax = math.ceil(x.max() / CELL_SIZE) * CELL_SIZE
    ymin = math.floor(y.min() / CELL_SIZE) * CELL_SIZE
    ymax = math.ceil(y.max() / CELL_SIZE) * CELL_SIZE

    width = int(round((xmax - xmin) / CELL_SIZE))
    height = int(round((ymax - ymin) / CELL_SIZE))

    print(f"[GRID] {width} × {height} cells")
    print(f"       {CELL_SIZE / 1000:.0f} km resolution")

    cols = np.floor((x - xmin) / CELL_SIZE).astype(np.int32)
    rows = np.floor((ymax - y) / CELL_SIZE).astype(np.int32)

    cols = np.clip(cols, 0, width - 1)
    rows = np.clip(rows, 0, height - 1)

    years = gdf[YEAR_FIELD].astype(int).to_numpy()
    areas = gdf[AREA_FIELD].to_numpy(float)
    masses = gdf[MASS_FIELD].to_numpy(float)

    # One raster per year for each metric.
    annual_count = np.zeros(
        (len(YEARS), height, width),
        dtype=np.float64,
    )
    annual_area = np.zeros_like(annual_count)
    annual_mass = np.zeros_like(annual_count)

    for i, year in enumerate(YEARS):
        select = years == year
        r = rows[select]
        c = cols[select]

        np.add.at(annual_count[i], (r, c), 1)
        np.add.at(annual_area[i], (r, c), areas[select])
        np.add.at(annual_mass[i], (r, c), masses[select])

        print(
            f"[YEAR] {year}: {select.sum():,} icebergs | "
            f"{areas[select].sum():,.1f} km² | "
            f"{masses[select].sum():,.1f} Gt"
        )

    # Mean and population SD across the six annual values.
    mean_count = annual_count.mean(axis=0)
    mean_area = annual_area.mean(axis=0)
    mean_mass = annual_mass.mean(axis=0)

    SD_count = annual_count.std(axis=0, ddof=0)
    SD_area = annual_area.std(axis=0, ddof=0)
    SD_mass = annual_mass.std(axis=0, ddof=0)

    has_data = annual_count.sum(axis=0) > 0

    arrays = [
        mean_count,
        mean_area,
        mean_mass,
        SD_count,
        SD_area,
        SD_mass,
    ]

    descriptions = [
        "mean_count",
        "mean_area",
        "mean_mass",
        "SD_count",
        "SD_area",
        "SD_mass",
    ]

    srs = osr.SpatialReference()
    srs.ImportFromEPSG(EPSG)

    mem = gdal.GetDriverByName("MEM").Create(
        "",
        width,
        height,
        len(arrays),
        gdal.GDT_Float32,
    )

    mem.SetGeoTransform(
        (xmin, CELL_SIZE, 0, ymax, 0, -CELL_SIZE)
    )
    mem.SetProjection(srs.ExportToWkt())

    for i, (array, description) in enumerate(
        zip(arrays, descriptions),
        start=1,
    ):
        output = np.full(
            (height, width),
            NODATA,
            dtype=np.float32,
        )
        output[has_data] = array[has_data]

        band = mem.GetRasterBand(i)
        band.WriteArray(output)
        band.SetNoDataValue(NODATA)
        band.SetDescription(description)

    # Remove old output and GDAL sidecars.
    for path in [
        OUTPUT_FILE,
        Path(str(OUTPUT_FILE) + ".aux.xml"),
        Path(str(OUTPUT_FILE) + ".ovr"),
    ]:
        if path.exists():
            path.unlink()

    result = gdal.GetDriverByName("COG").CreateCopy(
        str(OUTPUT_FILE),
        mem,
        strict=1,
        options=[
            "COMPRESS=DEFLATE",
            "PREDICTOR=3",
            "BIGTIFF=IF_SAFER",
        ],
    )

    if result is None:
        raise RuntimeError("Failed to create output COG.")

    result = None
    mem = None

    # Verify band descriptions.
    check = gdal.Open(str(OUTPUT_FILE))

    actual = [
        check.GetRasterBand(i).GetDescription()
        for i in range(1, check.RasterCount + 1)
    ]

    check = None

    if actual != descriptions:
        raise RuntimeError(f"Incorrect band descriptions: {actual}")

    print()
    print(f"[DATA] {has_data.sum():,} cells contain icebergs")
    print(f"[NODATA] {(~has_data).sum():,} empty cells")
    print("[BANDS]")

    for i, description in enumerate(descriptions, 1):
        print(f"       {i} = {description}")

    print(f"[OK] {OUTPUT_FILE.relative_to(BASE_DIR)}")


if __name__ == "__main__":
    main()