#!/usr/bin/env python3

# Merge Antarctic ice-shelf boundary shapefiles downloaded as per-shelf ZIPs.
#
# - Reads all ZIP files from cwd/input
# - Extracts all contained shapefiles
# - Derives FEATURE and YEAR from filenames (e.g. abbott_2018.shp)
# - Ensures source shapefiles are EPSG:3031
# - Merges all shelves and years into single shapefile in cwd/output
#
# Install:
#   pip install geopandas pandas shapely

from pathlib import Path
import re
import tempfile
import zipfile

import geopandas as gpd
import pandas as pd


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

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

OUTPUT_FILE = OUTPUT_DIR / "ice_shelf_boundaries_2009-2019.shp"

TARGET_CRS = "EPSG:3031"

NAME_RE = re.compile(r"^(.+)_(\d{4})$", re.IGNORECASE)


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

def parse_name(path: Path) -> tuple[str, int]:
    match = NAME_RE.match(path.stem)

    if not match:
        raise ValueError(
            f"Could not parse feature name/year from: {path.name}"
        )

    feature = match.group(1)
    year = int(match.group(2))

    return feature, year


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("*.zip"))

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

    frames = []

    for zip_path in zip_files:
        print(f"[ZIP] {zip_path.name}")

        with tempfile.TemporaryDirectory() as temp_dir:
            temp_dir = Path(temp_dir)

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

            shapefiles = sorted(
                path
                for path in temp_dir.rglob("*.shp")
                if "__MACOSX" not in path.parts
                and not path.name.startswith("._")
            )

            if not shapefiles:
                raise RuntimeError(
                    f"No shapefiles found in {zip_path.name}"
                )

            for shp_path in shapefiles:
                feature, year = parse_name(shp_path)

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

                gdf = gpd.read_file(shp_path)

                if gdf.crs is None:
                    raise RuntimeError(
                        f"CRS is undefined in {shp_path.name}"
                    )

                gdf = gdf.to_crs(TARGET_CRS)

                gdf["shelf"] = feature
                gdf["year"] = year

                gdf = gdf[["shelf", "year", "geometry"]]

                frames.append(gdf)

    merged = gpd.GeoDataFrame(
        pd.concat(frames, ignore_index=True),
        crs=TARGET_CRS,
    )

    remove_existing_shapefile(OUTPUT_FILE)

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

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


if __name__ == "__main__":
    main()