#!/usr/bin/env python3

# Merge future ice-free area shapefiles and add scenario/estimate attributes.
#
# Install requirements:
#   pip install geopandas pyogrio

from pathlib import Path
import geopandas as gpd
import pandas as pd


# =========================
# CONFIG
# =========================

BASE_DIR = Path.cwd()

INPUT_DIR = BASE_DIR / "input"
OUTPUT_DIR = BASE_DIR / "output"

OUTPUT_FILE = OUTPUT_DIR / "future_ice-free_areas.shp"


# =========================
# PROCESSING
# =========================

def get_attributes(filename):
    """Get scenario and estimate from filename."""

    name = filename.lower()

    if "rcp45" in name:
        scenario = "RCP4.5"
    elif "rcp85" in name:
        scenario = "RCP8.5"
    else:
        raise ValueError(f"Could not determine scenario from: {filename}")

    if "_lower_" in name:
        estimate = "lower"
    elif "_upper_" in name:
        estimate = "upper"
    elif "_best_" in name:
        estimate = "mean"
    else:
        raise ValueError(f"Could not determine estimate from: {filename}")

    return scenario, estimate


def main():

    OUTPUT_DIR.mkdir(exist_ok=True)

    shapefiles = sorted(INPUT_DIR.glob("*.shp"))

    if not shapefiles:
        raise FileNotFoundError(f"No shapefiles found in: {INPUT_DIR}")

    layers = []

    for shp in shapefiles:

        scenario, estimate = get_attributes(shp.name)

        print(f"Reading: {shp.name}")

        gdf = gpd.read_file(shp)

        # Drop all existing attributes
        gdf = gdf[["geometry"]].copy()

        # Add new attributes
        gdf["scenario"] = scenario
        gdf["estimate"] = estimate

        layers.append(gdf)

    merged = gpd.GeoDataFrame(
        pd.concat(layers, ignore_index=True),
        geometry="geometry",
        crs=layers[0].crs,
    )

    merged["geometry"] = merged.geometry.make_valid()
    merged.to_file(OUTPUT_FILE)

    print(f"Created: {OUTPUT_FILE}")


if __name__ == "__main__":
    main()