"""Township-resolution ShakeMap for all of Taiwan.

Predicts CWA intensity at every township/city-district centre (the 371 centres in
data/tw_town_centers.json) and paints each township polygon
(data/twtown.geojson, twtown2010.4) with that value.

Two known differences between the two files are reconciled by normalising names:

- the boundary file is the 2010 vintage, so it still says 桃園縣 and uses
  市/鎮/鄉 where the centre file uses 區 (桃園市), plus 員林鎮/員林市 and
  頭份鎮/頭份市;
- some coastal districts appear as an extra "(海)" polygon, and a few centres
  (釣魚臺列嶼, 烏坵鄉, 東沙群島, 南沙群島) have no polygon at all.

Usage:
  python shakemap.py                                   # 2024-04-03 M7.1 Hualien
  python shakemap.py --from-event "2024/04/03 07:58:09"
  python shakemap.py --M 6.5 --depth 10 --ev-lat 23.9 --ev-lon 121.6
"""
import argparse
import json
import math
import os
import sys

import numpy as np

from palert.features import INTENSITY_LEVELS
from palert.paths import DATA_DIR
from palert.predict import IntensityGMM

CENTERS_PATH = os.path.join(DATA_DIR, "tw_town_centers.json")
BOUNDARY_PATH = os.path.join(DATA_DIR, "twtown.geojson")
DATASET_PATH = os.path.join(DATA_DIR, "cwa_dataset.csv")

ASCII_LABELS = ["0", "1", "2", "3", "4", "5-", "5+", "6-", "6+", "7"]
LEVEL_COLORS = ["#ffffff", "#c8e6f5", "#a5d8f3", "#7fc9ee", "#57b6e5",
                "#3b9ede", "#f5d000", "#f5a000", "#f56000", "#d00000"]
SUFFIXES = ("市", "鎮", "鄉", "區")

# The boundary file predates the 2014 Taoyuan upgrade. Only this one county was
# renamed; do NOT strip 縣/市 from county names or 新竹縣/新竹市 (and the
# Chiayi pair) would collide.
COUNTY_ALIAS = {"桃園縣": "桃園市"}


def norm(name):
    """Normalise a township name: drop the (海) marker and the 市/鎮/鄉/區 suffix."""
    name = (name or "").strip()
    name = name.replace("(海)", "").replace("（海）", "")
    while name and name[-1] in SUFFIXES:
        name = name[:-1]
    return name


def load_towns(path=CENTERS_PATH):
    """Return [{'county','town','lat','lon'}] from the nested centres JSON."""
    with open(path, encoding="utf-8") as fh:
        raw = json.load(fh)
    root = raw.get("台灣") or next(iter(raw.values()))
    towns = []
    for county, items in root.items():
        for town, info in items.items():
            towns.append({
                "county": county,
                "town": town,
                "lat": float(info["latitude"]),
                "lon": float(info["longitude"]),
            })
    return towns


def load_boundaries(path=BOUNDARY_PATH):
    """Return a list of (key, county, town, [exterior_ring, ...]) entries."""
    with open(path, encoding="utf-8") as fh:
        gj = json.load(fh)
    out = []
    for feat in gj["features"]:
        props = feat.get("properties") or {}
        county, town = props.get("county", ""), props.get("town", "")
        county = COUNTY_ALIAS.get(county, county)
        geom = feat.get("geometry") or {}
        coords = geom.get("coordinates")
        if not coords:
            continue
        polys = coords if geom.get("type") == "MultiPolygon" else [coords]
        rings = []
        for poly in polys:
            if poly and len(poly[0]) >= 3:
                rings.append(np.asarray(poly[0], dtype=np.float64)[:, :2])
        if rings:
            out.append((norm(county) + "/" + norm(town), county, town, rings))
    return out


def event_from_dataset(when, path=DATASET_PATH):
    import pandas as pd

    if not os.path.exists(path):
        raise SystemExit(
            f"{path} not found. --from-event needs the training dataset, which is not\n"
            "shipped with the published model. Build it first (fetch_all.py /\n"
            "fetch_cwa_scweb.py / merge_datasets.py), or pass the event directly:\n"
            '  python shakemap.py --M 7.1 --depth 19.7 --ev-lat 23.88 --ev-lon 121.57'
        )
    df = pd.read_csv(path)
    sel = df[df["event"] == when]
    if sel.empty:
        raise SystemExit(f"event {when!r} not in {path}")
    r = sel.iloc[0]
    return float(r["M"]), float(r["depth"]), float(r["ev_lat"]), float(r["ev_lon"])


def main():
    try:
        sys.stdout.reconfigure(encoding="utf-8")
    except Exception:
        pass

    ap = argparse.ArgumentParser()
    ap.add_argument("--M", type=float, default=7.1)
    ap.add_argument("--depth", type=float, default=19.7)
    ap.add_argument("--ev-lat", type=float, default=23.88, dest="ev_lat")
    ap.add_argument("--ev-lon", type=float, default=121.57, dest="ev_lon")
    ap.add_argument("--from-event", help='pull M/depth/epicentre from cwa_dataset.csv, e.g. "2024/04/03 07:58:09"')
    ap.add_argument("--out", default=os.path.join(DATA_DIR, "shakemap.png"))
    ap.add_argument("--model", default=None, help="ONNX model (default: TDN-IntensityGMM-TW-524k-v1.onnx)")
    args = ap.parse_args()

    if args.from_event:
        M, depth, ev_lat, ev_lon = event_from_dataset(args.from_event)
    else:
        M, depth, ev_lat, ev_lon = args.M, args.depth, args.ev_lat, args.ev_lon
    print(f"event: M{M} depth {depth} km @ ({ev_lat}, {ev_lon})")

    towns = load_towns()
    boundaries = load_boundaries()
    print(f"towns: {len(towns)} centres, {len(boundaries)} boundary features")

    # --- predict at every township centre (single batched call) ---
    gmm = IntensityGMM(args.model) if args.model else IntensityGMM()
    lats = np.array([t["lat"] for t in towns])
    lons = np.array([t["lon"] for t in towns])
    pga, pgv, level = gmm.predict_batch(M, depth, ev_lat, ev_lon, lats, lons)

    by_key = {}
    for t, p, v, lv in zip(towns, pga, pgv, level):
        by_key[f"{norm(t['county'])}/{norm(t['town'])}"] = (p, v, int(lv), t)

    # --- reconcile boundary features against the centres ---
    matched, unmatched_feat, unmatched_town = {}, [], set(by_key)
    for key, county, town, rings in boundaries:
        if key in by_key:
            matched[key] = by_key[key]
            unmatched_town.discard(key)
        else:
            unmatched_feat.append(f"{county}{town}")
    missing_ring = [f"{by_key[k][3]['county']}{by_key[k][3]['town']}" for k in sorted(unmatched_town)]
    if unmatched_feat:
        print(f"[warn] {len(unmatched_feat)} boundary features with no centre: {unmatched_feat}")
    if missing_ring:
        print(f"[warn] {len(missing_ring)} centres with no polygon (not drawn): {missing_ring}")

    # --- render ---
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.collections import PolyCollection
    from matplotlib.patches import Patch

    polys, face = [], []
    minx = miny = float("inf")
    maxx = maxy = float("-inf")
    for key, county, town, rings in boundaries:
        info = matched.get(key)
        color = LEVEL_COLORS[info[2]] if info else "#dddddd"
        for ring in rings:
            polys.append(ring)
            face.append(color)
            minx, maxx = min(minx, ring[:, 0].min()), max(maxx, ring[:, 0].max())
            miny, maxy = min(miny, ring[:, 1].min()), max(maxy, ring[:, 1].max())

    fig, ax = plt.subplots(figsize=(7.5, 10), dpi=130)
    ax.add_collection(PolyCollection(
        polys, facecolors=face, edgecolors="black", linewidths=0.25, zorder=2
    ))
    ax.plot(ev_lon, ev_lat, marker="*", color="k", markersize=17,
            markeredgecolor="white", markeredgewidth=0.8, zorder=5, label="epicenter")

    pad = 0.08
    ax.set_xlim(minx - pad, maxx + pad)
    ax.set_ylim(miny - pad, maxy + pad)
    ax.set_aspect(1.0 / math.cos(math.radians((miny + maxy) / 2)))
    ax.set_xlabel("longitude")
    ax.set_ylabel("latitude")
    ax.set_title(f"Predicted CWA intensity by township  |  M{M} depth {depth} km "
                 f"@ ({ev_lat:.2f}, {ev_lon:.2f})")

    handles = [Patch(facecolor=LEVEL_COLORS[i], edgecolor="black", linewidth=0.3,
                     label=ASCII_LABELS[i]) for i in range(10)]
    ax.legend(handles=handles, title="intensity", ncol=2, loc="lower left",
              fontsize=8, title_fontsize=8, framealpha=0.9)
    fig.tight_layout()
    fig.savefig(args.out)
    print(f"[save] {args.out}")

    # --- summary ---
    order = np.argsort(-pga)
    print(f"\ntop 10 towns by PGA (max {pga.max():.1f} gal, {INTENSITY_LEVELS[int(level.max())]}):")
    print(f"    {'town':<12}{'county':<10}{'PGA(gal)':>10}{'PGV(cm/s)':>11}  intensity")
    for i in order[:10]:
        t = towns[i]
        print(f"    {t['town']:<12}{t['county']:<10}{pga[i]:10.1f}{pgv[i]:11.3f}  "
              f"{INTENSITY_LEVELS[int(level[i])]}")
    counts = np.bincount(level, minlength=10)
    print("\ntownship count by intensity:",
          {ASCII_LABELS[i]: int(c) for i, c in enumerate(counts) if c})


if __name__ == "__main__":
    main()