#!/usr/bin/env python3

# Merge annual iceberg-distribution GeoPackages into one shapefile.
#
# Expected input filenames:
#   201810_distribution.gpkg
#   201910_distribution.gpkg
#   ...
#
# - Adds year from the first four digits of each filename
# - Validates EPSG:3031
# - Retains polygonal iceberg geometries only
# - Reports any non-polygon geometries excluded
#
# Input:
#   cwd/input/*.gpkg
#
# Output:
#   cwd/output/Antarctic_icebergs_2018-2023.shp
#
# Install:
#   pip install geopandas pandas shapely pyogrio

from pathlib import Path
import re

import geopandas as gpd
import pandas as pd
from shapely import make_valid
from shapely.geometry import MultiPolygon, Polygon
from shapely.ops import unary_union


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

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

OUTPUT_FILE = OUTPUT_DIR / "Antarctic_icebergs_2018-2023.shp"

EXPECTED_EPSG = 3031

FIELD_RENAMES = {
    "area_uncertainty_km2": "A_unc_km2",
    "perimeter_km": "P_km",
    "long_axis_km": "long_ax_km",
    "short_axis_km": "shortax_km",
    "mass_uncertainty_gt": "mass_unc",
}


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

def polygonal_geometry(geom):
    """Return only polygonal content from a geometry."""
    if geom is None or geom.is_empty:
        return None

    if not geom.is_valid:
        geom = make_valid(geom)

    if isinstance(geom, (Polygon, MultiPolygon)):
        return geom

    if geom.geom_type == "GeometryCollection":
        polygons = [
            part
            for part in geom.geoms
            if isinstance(part, (Polygon, MultiPolygon))
        ]

        if polygons:
            return unary_union(polygons)

    return None


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

def main():
    gpkg_files = sorted(INPUT_DIR.glob("*.gpkg"))

    if not gpkg_files:
        raise FileNotFoundError(f"No GeoPackages found in:\n{INPUT_DIR}")

    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    layers = []
    total_dropped = 0

    for path in gpkg_files:
        match = re.match(r"^(\d{4})\d{2}_distribution\.gpkg$", path.name)

        if not match:
            raise ValueError(
                f"Unexpected filename:\n{path.name}\n\n"
                "Expected format: YYYYMM_distribution.gpkg"
            )

        year = int(match.group(1))

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

        gdf = gpd.read_file(path)

        if gdf.crs is None:
            raise ValueError(f"CRS is undefined: {path.name}")

        if gdf.crs.to_epsg() != EXPECTED_EPSG:
            raise ValueError(
                f"Unexpected CRS in {path.name}: {gdf.crs}\n"
                f"Expected EPSG:{EXPECTED_EPSG}"
            )

        # Show source geometry types before cleaning.
        geometry_counts = gdf.geometry.geom_type.value_counts()

        for geom_type, count in geometry_counts.items():
            if geom_type not in ("Polygon", "MultiPolygon"):
                print(f"       [GEOMETRY] {count:,} {geom_type}")

        original_count = len(gdf)

        gdf["geometry"] = gdf.geometry.apply(polygonal_geometry)
        gdf = gdf[gdf.geometry.notna() & ~gdf.geometry.is_empty].copy()

        dropped = original_count - len(gdf)
        total_dropped += dropped

        if dropped:
            print(f"       [SKIP] {dropped:,} non-polygon feature(s)")

        gdf["year"] = year

        print(f"       {len(gdf):,} icebergs retained | year={year}")

        layers.append(gdf)

    merged = gpd.GeoDataFrame(
        pd.concat(layers, ignore_index=True),
        crs=f"EPSG:{EXPECTED_EPSG}",
    )

    merged = merged.rename(columns=FIELD_RENAMES)

    # Shapefile field names cannot exceed 10 characters.
    long_fields = [
        field
        for field in merged.columns
        if field != "geometry" and len(field) > 10
    ]

    if long_fields:
        raise ValueError(
            "These field names exceed the 10-character Shapefile limit:\n"
            + "\n".join(f"  {field}" for field in long_fields)
        )

    # Remove an existing shapefile and sidecars.
    for existing in OUTPUT_DIR.glob(f"{OUTPUT_FILE.stem}.*"):
        existing.unlink()

    merged.to_file(
        OUTPUT_FILE,
        driver="ESRI Shapefile",
        encoding="UTF-8",
    )

    print()
    print(f"[TOTAL] {len(merged):,} iceberg polygons")
    print(f"[SKIPPED] {total_dropped:,} non-polygon feature(s)")
    print(f"[CRS] EPSG:{EXPECTED_EPSG}")
    print(f"[OK] {OUTPUT_FILE.relative_to(BASE_DIR)}")


if __name__ == "__main__":
    main()