#!/usr/bin/env python3

# Smooth monthly Antarctic sea-ice extent polygons for cartographic display.
#
# Approximates ArcGIS PAEK smoothing using:
#   - regularly sampled polygon boundaries
#   - local second-degree polynomial fitting
#   - exponentially decreasing weights along the boundary
#
# Install:
#   pip install geopandas numpy shapely

from pathlib import Path

import geopandas as gpd
import numpy as np
from shapely.geometry import (
    GeometryCollection,
    MultiPolygon,
    Polygon,
)
from shapely.ops import unary_union


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

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

INPUT_FILE = OUTPUT_DIR / "seaice_extent_monthly_polygons.shp"
OUTPUT_FILE = OUTPUT_DIR / "seaice_extent_monthly_polygons_smoothed.shp"

SMOOTH_TOLERANCE_M = 180_000
SAMPLE_SPACING_M = 5_000
SIMPLIFY_TOLERANCE_M = 200

TARGET_EPSG = 3031


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

def sample_ring(coords, spacing):
    """Resample a closed ring at approximately regular spacing."""

    xy = np.asarray(coords, dtype=float)

    if np.allclose(xy[0], xy[-1]):
        xy = xy[:-1]

    closed = np.vstack([xy, xy[0]])

    segments = np.diff(closed, axis=0)
    lengths = np.hypot(segments[:, 0], segments[:, 1])
    cumulative = np.concatenate(([0.0], np.cumsum(lengths)))
    perimeter = cumulative[-1]

    if perimeter == 0:
        return xy

    n = max(8, int(np.ceil(perimeter / spacing)))
    distances = np.arange(n) * perimeter / n

    idx = np.searchsorted(
        cumulative,
        distances,
        side="right",
    ) - 1

    idx = np.clip(idx, 0, len(lengths) - 1)

    fraction = (
        (distances - cumulative[idx])
        / lengths[idx]
    )

    return (
        closed[idx]
        + segments[idx] * fraction[:, None]
    )


def smoothing_coefficients(spacing, tolerance, point_count):

    half_window = tolerance / 2

    radius = max(
        1,
        int(round(half_window / spacing)),
    )

    radius = min(
        radius,
        max(1, (point_count - 1) // 2),
    )

    offsets = np.arange(
        -radius,
        radius + 1,
        dtype=float,
    ) * spacing

    # Normalise distances to avoid numerical instability in the
    # quadratic polynomial fit.
    u = offsets / half_window

    # Exponential weighting: nearer vertices have greater influence.
    weights = np.exp(-2.0 * np.abs(u))

    design = np.column_stack([
        np.ones_like(u),
        u,
        u ** 2,
    ])

    weighted_design = weights[:, None] * design

    fit = (
        np.linalg.pinv(design.T @ weighted_design)
        @ design.T
        @ np.diag(weights)
    )

    coeffs = fit[0]

    # Preserve absolute coordinate position exactly.
    coeffs /= coeffs.sum()

    return radius, coeffs


def smooth_ring(coords):
    sampled = sample_ring(
        coords,
        SAMPLE_SPACING_M,
    )

    n = len(sampled)

    if n < 5:
        return coords

    perimeter = np.sum(
        np.hypot(
            np.diff(
                np.vstack([sampled, sampled[0]]),
                axis=0,
            )[:, 0],
            np.diff(
                np.vstack([sampled, sampled[0]]),
                axis=0,
            )[:, 1],
        )
    )

    spacing = perimeter / n

    radius, coeffs = smoothing_coefficients(
        spacing,
        SMOOTH_TOLERANCE_M,
        n,
    )

    x = np.zeros(n)
    y = np.zeros(n)

    for offset, coefficient in zip(
        range(-radius, radius + 1),
        coeffs,
    ):
        rolled = np.roll(sampled, -offset, axis=0)

        x += coefficient * rolled[:, 0]
        y += coefficient * rolled[:, 1]

    smoothed = np.column_stack([x, y])

    return np.vstack([
        smoothed,
        smoothed[0],
    ])


def smooth_polygon(polygon):
    exterior = smooth_ring(
        polygon.exterior.coords
    )

    interiors = [
        smooth_ring(ring.coords)
        for ring in polygon.interiors
    ]

    result = Polygon(
        exterior,
        interiors,
    )

    if not result.is_valid:
        result = result.buffer(0)

    return result


def polygon_parts(geometry):
    if geometry is None or geometry.is_empty:
        return []

    if isinstance(geometry, Polygon):
        return [geometry]

    if isinstance(geometry, MultiPolygon):
        return list(geometry.geoms)

    if isinstance(geometry, GeometryCollection):
        parts = []

        for part in geometry.geoms:
            parts.extend(polygon_parts(part))

        return parts

    return []


def smooth_geometry(geometry):
    parts = polygon_parts(geometry)

    if not parts:
        return None

    smoothed = [
        smooth_polygon(part)
        for part in parts
    ]

    result = unary_union(smoothed)

    if not result.is_valid:
        result = result.buffer(0)

    return result

def vertex_count(geometry):
    if isinstance(geometry, Polygon):
        return (
            len(geometry.exterior.coords)
            + sum(len(ring.coords) for ring in geometry.interiors)
        )

    if isinstance(geometry, MultiPolygon):
        return sum(vertex_count(part) for part in geometry.geoms)

    return 0

def remove_existing_shapefile(path):
    for ext in (
        ".shp",
        ".shx",
        ".dbf",
        ".prj",
        ".cpg",
        ".qix",
        ".fix",
    ):
        sidecar = path.with_suffix(ext)

        if sidecar.exists():
            sidecar.unlink()


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

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

    gdf = gpd.read_file(INPUT_FILE)

    if gdf.crs is None:
        raise RuntimeError(
            "Input has no defined CRS"
        )

    if gdf.crs.to_epsg() != TARGET_EPSG:
        raise RuntimeError(
            f"Expected EPSG:{TARGET_EPSG}, "
            f"found {gdf.crs}"
        )

    print(
        f"Smoothing {len(gdf)} polygon(s) "
        f"with {SMOOTH_TOLERANCE_M / 1000:.0f} km tolerance"
    )

    smoothed = []

    for i, geometry in enumerate(
        gdf.geometry,
        start=1,
    ):
        print(
            f"[SMOOTH] {i}/{len(gdf)}"
        )

        result = smooth_geometry(geometry)

        if result is None or result.is_empty:
            raise RuntimeError(
                f"Smoothing produced empty geometry "
                f"for feature {i - 1}"
            )

        smoothed.append(result)

    gdf = gdf.set_geometry(smoothed)

    before = sum(vertex_count(g) for g in gdf.geometry)

    gdf["geometry"] = gdf.geometry.simplify(
        SIMPLIFY_TOLERANCE_M,
        preserve_topology=True,
    )

    after = sum(vertex_count(g) for g in gdf.geometry)

    print(
        f"[SIMPLIFY] {SIMPLIFY_TOLERANCE_M / 1000:g} km: "
        f"{before:,} -> {after:,} vertices "
        f"({100 * (1 - after / before):.1f}% reduction)"
    )

    if not gdf.is_valid.all():
        bad = list(
            gdf.index[~gdf.is_valid][:10]
        )

        raise RuntimeError(
            f"Invalid output geometries at indices: {bad}"
        )

    remove_existing_shapefile(
        OUTPUT_FILE
    )

    gdf.to_file(
        OUTPUT_FILE,
        driver="ESRI Shapefile",
        index=False,
    )

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


if __name__ == "__main__":
    main()