# /// script
# requires-python = ">=3.11"
# dependencies = ["shapely>=2.0"]
# ///
"""Build src/static/flood-zones-p1.geojson from CMU Water Center's flood map.

CMU's Chiang Mai flood hazard map (watercenter.cmu.ac.th, CMFlood, built from
the October 2024 flood) publishes the flooded area at five P.1 levels as KML:
/data/KML5Layer/L1.kml (P.1 = 4.30 m) to L5.kml (5.30 m, the 2024 record).
Each file is the whole extent at that level. The dashboard wants one ring per
stage instead: zone k is what floods first when P.1 reaches stage k's level
(the extent at level k minus everything flooded at a lower level).

    uv run scripts/build_flood_zones.py              # fetch the KML from CMU
    uv run scripts/build_flood_zones.py --kml-dir D  # or read L1..L5.kml from D

The levels must match P1_FLOOD_STAGES in src/ml/features.py.
"""

import argparse
import json
import re
import urllib.request
from pathlib import Path

import shapely
from shapely import make_valid
from shapely.geometry import MultiPolygon, Polygon, mapping
from shapely.ops import unary_union

REPO = Path(__file__).resolve().parents[1]
OUT = REPO / "src" / "static" / "flood-zones-p1.geojson"
BASE = "https://watercenter.cmu.ac.th/data/KML5Layer"
LEVELS = [4.30, 4.50, 4.70, 5.00, 5.30]  # L1..L5, P.1 gauge level in m
# ~3 m at Chiang Mai: invisible at the zoom levels the map allows, and it
# keeps the file small. Slivers below MIN_AREA (about 300 m²) come from the
# layers' borders not lining up exactly and are dropped.
SIMPLIFY_DEG = 3e-5
MIN_AREA_DEG2 = 2.5e-8
GRID = 1e-5  # coordinates snapped to 5 decimals (~1 m) with valid geometry
SOURCE = (
    "CMU Water Center (ศูนย์วิชาการสนับสนุนการบริหารจัดการน้ำ มหาวิทยาลัยเชียงใหม่), "
    "Chiang Mai flood hazard map from the October 2024 flood, "
    "https://watercenter.cmu.ac.th/chiangmai/cmflood/warning/floodmap"
)


def fetch(i: int, kml_dir: Path | None) -> str:
    if kml_dir:
        return (kml_dir / f"L{i}.kml").read_text(encoding="utf-8")
    req = urllib.request.Request(
        f"{BASE}/L{i}.kml", headers={"User-Agent": "Mozilla/5.0"}
    )
    with urllib.request.urlopen(req, timeout=60) as r:
        return r.read().decode("utf-8")


def ring(text: str) -> list:
    pts = []
    for tok in text.split():
        x, y, *_ = tok.split(",")
        pts.append((float(x), float(y)))
    return pts


def kml_extent(text: str):
    polys = []
    for body in re.findall(r"<Polygon>(.*?)</Polygon>", text, re.S):
        outer = re.search(
            r"<outerBoundaryIs>.*?<coordinates>(.*?)</coordinates>", body, re.S
        )
        inner = re.findall(
            r"<innerBoundaryIs>.*?<coordinates>(.*?)</coordinates>", body, re.S
        )
        polys.append(
            make_valid(Polygon(ring(outer.group(1)), [ring(s) for s in inner]))
        )
    return unary_union(polys)


def polygons(geom) -> list:
    if isinstance(geom, Polygon):
        return [geom]
    return [
        g for g in getattr(geom, "geoms", []) if isinstance(g, (Polygon, MultiPolygon))
    ]


def clean(geom):
    parts = []
    for g in polygons(geom):
        for p in polygons(g):
            if p.area >= MIN_AREA_DEG2:
                parts.append(p)
    return unary_union(parts)


def rounded(coords):
    # the geometry is already on the grid; this only drops float noise
    if isinstance(coords, (list, tuple)) and coords and isinstance(coords[0], float):
        return [round(coords[0], 5), round(coords[1], 5)]
    return [rounded(c) for c in coords]


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--kml-dir", type=Path, help="read L1..L5.kml from here")
    args = ap.parse_args()

    flooded = None
    features = []
    for i, level in enumerate(LEVELS, 1):
        extent = kml_extent(fetch(i, args.kml_dir))
        extent = shapely.set_precision(
            make_valid(extent.simplify(SIMPLIFY_DEG, preserve_topology=True)), GRID
        )
        if flooded is None:
            new, flooded = extent, extent
        else:
            new = shapely.difference(extent, flooded, grid_size=GRID)
            flooded = shapely.union(flooded, extent, grid_size=GRID)
        zone = shapely.set_precision(clean(new), GRID)
        assert zone.is_valid, f"zone {i} is not a valid geometry"
        geo = mapping(zone)
        features.append(
            {
                "type": "Feature",
                "properties": {"zone": i, "level": level},
                "geometry": {
                    "type": geo["type"],
                    "coordinates": rounded(geo["coordinates"]),
                },
            }
        )
        print(f"zone {i} (P.1 {level:.2f} m): {len(polygons(zone))} polygons")
    doc = {
        "type": "FeatureCollection",
        "source": SOURCE,
        "note": "Zone k floods first when P.1 reaches its level; zones are disjoint.",
        # highest zone first, so the zones that flood first are drawn on top
        "features": features[::-1],
    }
    OUT.write_text(
        json.dumps(doc, ensure_ascii=False, separators=(",", ":")) + "\n",
        encoding="utf-8",
    )
    print(f"wrote {OUT.relative_to(REPO)}: {OUT.stat().st_size / 1024:.0f} KB")


if __name__ == "__main__":
    main()