#!/usr/bin/env python3

# Create a 100 km gridded iceberg-observation product from the merged
# NPI + AAD SCAR iceberg database.
#
# Bands:
#   1 = obs_count          Number of ship observation events
#   2 = iceberg_count      Total number of icebergs observed
#   3 = icebergs_per_obs   Mean icebergs per observation
#
# Cells with no observations are NoData (-999).
# icebergs_per_obs is also NoData where fewer than 3 observations occur.
#
# Input:
#   cwd/output/NPI_AAD_iceberg_database_1977-2011_merged.csv
#
# Output:
#   cwd/output/NPI_AAD_iceberg_observation_density_100km.tif
#
# Install:
#   pip install pandas numpy GDAL

from pathlib import Path

import numpy as np
import pandas as pd
from osgeo import gdal, osr


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

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

INPUT_FILE = OUTPUT_DIR / "NPI_AAD_iceberg_database_1977-2011_merged.csv"
OUTPUT_FILE = OUTPUT_DIR / "NPI_AAD_iceberg_observation_density_150km.tif"

CELL_SIZE = 150_000
TARGET_EPSG = 3031
MIN_OBS_PER_CELL = 3
NODATA = -999.0

COG_OPTIONS = [
    "COMPRESS=DEFLATE",
    "PREDICTOR=3",
    "BIGTIFF=IF_SAFER",
    "NUM_THREADS=ALL_CPUS",
]


# =============================================================================
# HELPERS
# =============================================================================

def transform_points(longitude, latitude):
    source = osr.SpatialReference()
    source.ImportFromEPSG(4326)
    source.SetAxisMappingStrategy(osr.OAMS_TRADITIONAL_GIS_ORDER)

    target = osr.SpatialReference()
    target.ImportFromEPSG(TARGET_EPSG)
    target.SetAxisMappingStrategy(osr.OAMS_TRADITIONAL_GIS_ORDER)

    transform = osr.CoordinateTransformation(source, target)

    points = transform.TransformPoints(
        list(zip(longitude, latitude))
    )

    x = np.array([p[0] for p in points], dtype=np.float64)
    y = np.array([p[1] for p in points], dtype=np.float64)

    return x, y


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

def main():
    gdal.UseExceptions()

    if not INPUT_FILE.exists():
        raise FileNotFoundError(f"Input not found:\n{INPUT_FILE}")

    df = pd.read_csv(INPUT_FILE)

    required = {"Latitude", "Longitude", "Total_Count"}
    missing = required - set(df.columns)

    if missing:
        raise RuntimeError(
            f"Missing required column(s): {', '.join(sorted(missing))}"
        )

    df["Latitude"] = pd.to_numeric(df["Latitude"], errors="coerce")
    df["Longitude"] = pd.to_numeric(df["Longitude"], errors="coerce")
    df["Total_Count"] = pd.to_numeric(df["Total_Count"], errors="coerce")

    valid = (
        df["Latitude"].between(-90, 90)
        & df["Longitude"].between(-180, 180)
        & df["Total_Count"].notna()
    )

    skipped = len(df) - valid.sum()
    df = df.loc[valid].copy()

    if df.empty:
        raise RuntimeError("No valid observations found")

    if (df["Total_Count"] < 0).any():
        raise RuntimeError("Negative Total_Count values found")

    print(f"[READ] {len(df):,} valid observations")

    if skipped:
        print(f"[SKIP] {skipped:,} rows with missing/invalid coordinates or counts")

    # -------------------------------------------------------------------------
    # Reproject observation positions to EPSG:3031
    # -------------------------------------------------------------------------

    x, y = transform_points(
        df["Longitude"].to_numpy(),
        df["Latitude"].to_numpy(),
    )

    # Snap the grid origin to the 100 km EPSG:3031 grid.
    xmin = np.floor(x.min() / CELL_SIZE) * CELL_SIZE
    ymax = np.ceil(y.max() / CELL_SIZE) * CELL_SIZE

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

    width = int(cols.max()) + 1
    height = int(rows.max()) + 1

    print(
        f"[GRID] {width} x {height} cells | "
        f"{CELL_SIZE / 1000:g} km | EPSG:{TARGET_EPSG}"
    )

    # -------------------------------------------------------------------------
    # Aggregate observations
    # -------------------------------------------------------------------------

    obs_count = np.zeros(
        (height, width),
        dtype=np.float64,
    )

    iceberg_count = np.zeros(
        (height, width),
        dtype=np.float64,
    )

    np.add.at(
        obs_count,
        (rows, cols),
        1,
    )

    np.add.at(
        iceberg_count,
        (rows, cols),
        df["Total_Count"].to_numpy(dtype=np.float64),
    )

    observed = obs_count > 0

    icebergs_per_obs = np.full(
        (height, width),
        NODATA,
        dtype=np.float64,
    )

    density_valid = obs_count >= MIN_OBS_PER_CELL

    icebergs_per_obs[density_valid] = (
        iceberg_count[density_valid]
        / obs_count[density_valid]
    )

    # NoData only where no observations exist for the count bands.
    obs_out = np.where(
        observed,
        obs_count,
        NODATA,
    ).astype(np.float32)

    iceberg_out = np.where(
        observed,
        iceberg_count,
        NODATA,
    ).astype(np.float32)

    density_out = icebergs_per_obs.astype(np.float32)

    print(
        f"[CELLS] {observed.sum():,} with observations | "
        f"{density_valid.sum():,} with >= {MIN_OBS_PER_CELL} observations"
    )

    # -------------------------------------------------------------------------
    # Write COG
    # -------------------------------------------------------------------------

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

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

    mem.SetGeoTransform(
        (
            xmin,
            CELL_SIZE,
            0,
            ymax,
            0,
            -CELL_SIZE,
        )
    )

    mem.SetProjection(srs.ExportToWkt())

    arrays = [
        obs_out,
        iceberg_out,
        density_out,
    ]

    descriptions = [
        "obs_count",
        "iceberg_count",
        "icebergs_per_obs",
    ]

    for band_num, (array, description) in enumerate(
        zip(arrays, descriptions),
        start=1,
    ):
        band = mem.GetRasterBand(band_num)
        band.WriteArray(array)
        band.SetNoDataValue(NODATA)
        band.SetDescription(description)

    mem.SetMetadata({
        "source": "Merged NPI and AAD SCAR iceberg observation database",
        "cell_size_m": str(CELL_SIZE),
        "min_obs_for_icebergs_per_obs": str(MIN_OBS_PER_CELL),
    })

    output = gdal.Translate(
        str(OUTPUT_FILE),
        mem,
        options=gdal.TranslateOptions(
            format="COG",
            creationOptions=COG_OPTIONS,
        ),
    )

    if output is None:
        raise RuntimeError(f"Failed to create:\n{OUTPUT_FILE}")

    output = None
    mem = None

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


if __name__ == "__main__":
    main()