#!/usr/bin/env python3

# Create pooled and annual iceberg-density surfaces for October 2018-2023.
#
# Input:
#   cwd/output/Antarctic_icebergs_2018-2023.shp
#   cwd/input/ice_mask.shp
#
# Output:
#   cwd/output/iceberg_interpolated_density_2018-2023.tif
#
# Bands:
#   1 = combined_2018-2023
#   2-7 = density_2018 ... density_2023
#
# Density uses iceberg centroid locations, a fixed 100 km Gaussian bandwidth
# and 10 km cells. Land/ice and areas north of 55°S are NoData (-999).
#
# Install:
#   pip install GDAL numpy pandas geopandas pyproj scipy pyogrio

from pathlib import Path
import math

import geopandas as gpd
import numpy as np
import pandas as pd
from osgeo import gdal, ogr, osr
from pyproj import Transformer
from scipy.ndimage import gaussian_filter

gdal.UseExceptions()


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

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

ICEBERG_FILE = OUTPUT_DIR / "Antarctic_icebergs_2018-2023.shp"
ICE_MASK_FILE = INPUT_DIR / "ice_mask.shp"
OUTPUT_FILE = OUTPUT_DIR / "iceberg_interpolated_density_2018-2023.tif"

YEAR_FIELD = "year"
LAT_FIELD = "lat"
LON_FIELD = "lon"

YEARS = range(2018, 2024)

EPSG = 3031
CELL_SIZE = 10_000
BANDWIDTH_KM = 100
BOUNDARY_LATITUDE = -55.0
NODATA = -999.0


# =============================================================================
# FUNCTIONS
# =============================================================================

def rasterize_mask(width, height, geotransform, projection):
    ds = ogr.Open(str(ICE_MASK_FILE))
    if ds is None:
        raise RuntimeError(f"Could not open {ICE_MASK_FILE}")

    layer = ds.GetLayer(0)
    mem = gdal.GetDriverByName("MEM").Create("", width, height, 1, gdal.GDT_Byte)
    mem.SetGeoTransform(geotransform)
    mem.SetProjection(projection)
    mem.GetRasterBand(1).Fill(0)

    gdal.RasterizeLayer(
        mem,
        [1],
        layer,
        burn_values=[1],
        options=["ALL_TOUCHED=TRUE"],
    )

    mask = mem.GetRasterBand(1).ReadAsArray().astype(bool)
    mem = None
    ds = None
    return mask


def make_density(x, y, xmin, ymax, width, height, sigma):
    cols = np.floor((x - xmin) / CELL_SIZE).astype(int)
    rows = np.floor((ymax - y) / CELL_SIZE).astype(int)

    valid = (
        (cols >= 0) & (cols < width)
        & (rows >= 0) & (rows < height)
    )

    counts = np.zeros((height, width), dtype=np.float64)
    np.add.at(counts, (rows[valid], cols[valid]), 1)

    smoothed = gaussian_filter(
        counts,
        sigma=sigma,
        mode="constant",
        cval=0,
        truncate=4.0,
    )

    cell_area_km2 = (CELL_SIZE / 1000) ** 2
    return (smoothed / cell_area_km2).astype(np.float32)


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

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

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

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

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

    gdf = gpd.read_file(ICEBERG_FILE)

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

    if missing:
        raise ValueError(f"Missing fields: {', '.join(missing)}")

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

    if gdf[required].isna().any(axis=None):
        raise ValueError("Invalid year/lat/lon values found.")

    gdf = gdf[
        gdf[YEAR_FIELD].isin(YEARS)
        & (gdf[LAT_FIELD] <= BOUNDARY_LATITUDE)
    ].copy()

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

    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)
    years = gdf[YEAR_FIELD].astype(int).to_numpy()

    # 55°S circular domain.
    bx, by = transformer.transform(0, BOUNDARY_LATITUDE)
    radius = math.hypot(bx, by)
    grid_radius = math.ceil(radius / CELL_SIZE) * CELL_SIZE

    xmin = ymin = -grid_radius
    xmax = ymax = grid_radius

    width = int((xmax - xmin) / CELL_SIZE)
    height = width

    geotransform = (xmin, CELL_SIZE, 0, ymax, 0, -CELL_SIZE)

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

    centres_x = xmin + (np.arange(width) + 0.5) * CELL_SIZE
    centres_y = ymax - (np.arange(height) + 0.5) * CELL_SIZE
    xx, yy = np.meshgrid(centres_x, centres_y)

    south_of_55 = np.hypot(xx, yy) <= radius

    print(f"[MASK] {ICE_MASK_FILE.name}")
    ice_mask = rasterize_mask(width, height, geotransform, projection)
    ocean = south_of_55 & ~ice_mask

    sigma = BANDWIDTH_KM * 1000 / CELL_SIZE

    descriptions = ["combined_2018-2023"] + [
        f"density_{year}" for year in YEARS
    ]

    densities = [
        make_density(x, y, xmin, ymax, width, height, sigma)
    ]

    print(
        f"[DENSITY] combined: {len(x):,} icebergs | "
        f"max={densities[0][ocean].max():.6f}"
    )

    for year in YEARS:
        select = years == year

        density = make_density(
            x[select],
            y[select],
            xmin,
            ymax,
            width,
            height,
            sigma,
        )

        densities.append(density)

        print(
            f"[DENSITY] {year}: {select.sum():,} icebergs | "
            f"max={density[ocean].max():.6f}"
        )

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

    mem.SetGeoTransform(geotransform)
    mem.SetProjection(projection)

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

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

    if OUTPUT_FILE.exists():
        OUTPUT_FILE.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

    check = gdal.Open(str(OUTPUT_FILE))

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

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

    check = None

    print("\n[BANDS]")
    for i, description in enumerate(descriptions, 1):
        print(f"       {i} = {description}")

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


if __name__ == "__main__":
    main()