#!/usr/bin/env python3

# Prepare monthly Antarctic sea-ice extent polygons for mapping display.
#
# - Reads polygon ZIPs from cwd/input
# - Merges and cleans source geometries
# - Dissolves to one feature per year/month
# - Derives month names and end-of-month date
# - Reprojects from native source CRS to EPSG:3031
# - Writes master polygons to cwd/output
#
# Install:
#   pip install geopandas pandas shapely

from calendar import monthrange
from datetime import datetime
from pathlib import Path
import re
import tempfile
import zipfile

import geopandas as gpd
import pandas as pd
from shapely.geometry import GeometryCollection, MultiPolygon, Polygon


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

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

OUTPUT_FILE = OUTPUT_DIR / "seaice_extent_monthly_polygons.shp"

TARGET_CRS = "EPSG:3031"

DATE_RE = re.compile(r"_(\d{6})_polygon_")


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

def parse_filename(filename: str) -> tuple[int, int]:
    match = DATE_RE.search(filename)

    if not match:
        raise ValueError(f"Could not parse year/month from: {filename}")

    yyyymm = match.group(1)

    return int(yyyymm[:4]), int(yyyymm[4:6])


def read_zip(zip_path: Path) -> gpd.GeoDataFrame:
    with tempfile.TemporaryDirectory() as temp_dir:
        temp_dir = Path(temp_dir)

        with zipfile.ZipFile(zip_path) as archive:
            archive.extractall(temp_dir)

        shapefiles = list(temp_dir.glob("*.shp"))

        if len(shapefiles) != 1:
            raise RuntimeError(
                f"Expected one shapefile in {zip_path.name}, "
                f"found {len(shapefiles)}"
            )

        gdf = gpd.read_file(shapefiles[0])
        
        if gdf.crs is None:
            raise RuntimeError(
                f"Source CRS is undefined in {zip_path.name}"
            )

        gdf = gdf.to_crs(TARGET_CRS)

    year, month_num = parse_filename(zip_path.name)

    gdf["YEAR"] = year
    gdf["MONTH_NUM"] = month_num

    return gdf


def clean_polygon(geometry):
    if geometry is None or geometry.is_empty:
        return None

    fixed = geometry.buffer(0)

    if fixed.is_empty:
        return None

    if isinstance(fixed, Polygon):
        return MultiPolygon([fixed])

    if isinstance(fixed, MultiPolygon):
        return fixed

    if isinstance(fixed, GeometryCollection):
        polygons = [
            part
            for part in fixed.geoms
            if isinstance(part, Polygon)
        ]
        return MultiPolygon(polygons) if polygons else None

    return None


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

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


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

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

    zip_files = sorted(INPUT_DIR.glob("*_polygon_*.zip"))

    if not zip_files:
        raise FileNotFoundError(
            f"No polygon ZIP files found in:\n{INPUT_DIR}"
        )

    print(f"Found {len(zip_files)} polygon ZIP(s)")

    frames = []

    for zip_path in zip_files:
        print(f"[READ] {zip_path.name}")
        frames.append(read_zip(zip_path))

    merged = gpd.GeoDataFrame(
        pd.concat(frames, ignore_index=True),
        crs=TARGET_CRS,
    ).drop(
        columns=["FID"],
        errors="ignore",
    )

    if (
        not merged.is_valid.all()
        or "GeometryCollection" in set(merged.geometry.geom_type)
    ):
        print("[CLEAN] Repairing polygon geometries")

        cleaned = merged.geometry.apply(clean_polygon)

        if cleaned.isna().any():
            bad = list(cleaned[cleaned.isna()].index[:10])

            raise RuntimeError(
                f"Geometry cleaning would drop features at indices: {bad}"
            )

        merged = merged.set_geometry(cleaned)

    print("[DISSOLVE] One feature per year/month")

    dissolved = merged.dissolve(
        by=["YEAR", "MONTH_NUM"],
        as_index=False,
    )

    # Derive display and temporal fields from YEAR + MONTH_NUM.
    dissolved["MONTH"] = dissolved.apply(
        lambda row: datetime(
            int(row["YEAR"]),
            int(row["MONTH_NUM"]),
            1,
        ).strftime("%B"),
        axis=1,
    )

    dissolved["MONTH_SHOR"] = dissolved.apply(
        lambda row: datetime(
            int(row["YEAR"]),
            int(row["MONTH_NUM"]),
            1,
        ).strftime("%b"),
        axis=1,
    )

    dissolved["DATE"] = dissolved.apply(
        lambda row: datetime(
            int(row["YEAR"]),
            int(row["MONTH_NUM"]),
            monthrange(
                int(row["YEAR"]),
                int(row["MONTH_NUM"]),
            )[1],
        ).strftime("%Y-%m-%d"),
        axis=1,
    )

    remove_existing_shapefile(OUTPUT_FILE)

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

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


if __name__ == "__main__":
    main()