"""Extreme-heat days at western frontier towns, 2016-2025, from ERA5 via Open-Meteo.

One batched request to the Open-Meteo archive API (models=era5) for daily maximum
temperature at eight reference towns. Writes heat_days.json and manifest.json for
/research/data/, used by /research/heat-frontier-hazard.html. The raw response is cached
under pipeline/cache/website/ (gitignored), so --offline rebuilds without the network.

usage (from the repo root):
    python3 research/data/build_heat.py              # fetch, then write research/data/
    python3 research/data/build_heat.py --offline    # rebuild from the cached response
"""
import json
import sys
import time
import urllib.request
from datetime import datetime, timezone
from pathlib import Path

HERE = Path(__file__).resolve().parent
ROOT = HERE.parents[1]
CACHE = ROOT / "pipeline" / "cache" / "website" / "raw_era5_tmax.json"
UA = "TrinetraResearch/0.3 (+https://trinetrahq.com)"

# Reference towns (public town centroids, 2 dp) from dashboard/data/raksha/locations.json
TOWNS = [
    ("sriganganagar", "Sri Ganganagar", "Rajasthan", 29.90, 73.88),
    ("bikaner", "Bikaner", "Rajasthan", 28.02, 73.31),
    ("jaisalmer", "Jaisalmer", "Rajasthan", 26.91, 70.91),
    ("barmer", "Barmer", "Rajasthan", 25.75, 71.39),
    ("bhuj", "Bhuj", "Gujarat", 23.24, 69.67),
    ("fazilka", "Fazilka", "Punjab", 30.40, 74.03),
    ("ferozepur", "Ferozepur", "Punjab", 30.93, 74.61),
    ("attari", "Attari", "Punjab", 31.60, 74.57),
]
START, END = "2016-01-01", "2025-12-31"
YEARS = list(range(2016, 2026))


def url():
    lat = ",".join(f"{t[3]:.2f}" for t in TOWNS)
    lon = ",".join(f"{t[4]:.2f}" for t in TOWNS)
    return ("https://archive-api.open-meteo.com/v1/archive"
            f"?latitude={lat}&longitude={lon}&start_date={START}&end_date={END}"
            "&daily=temperature_2m_max&models=era5&timezone=Asia%2FKolkata")


def fetch():
    req = urllib.request.Request(url(), headers={"User-Agent": UA})
    for attempt in range(3):
        try:
            with urllib.request.urlopen(req, timeout=40) as r:
                body = r.read()
            data = json.loads(body)
            CACHE.parent.mkdir(parents=True, exist_ok=True)
            CACHE.write_bytes(body)
            return data, datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
        except Exception as e:  # network or HTTP error: back off and retry
            print("fetch failed:", e, file=sys.stderr)
            time.sleep(5 * (attempt + 1))
    raise SystemExit("could not fetch ERA5 data")


def main():
    args = [a for a in sys.argv[1:] if not a.startswith("--")]
    out = Path(args[0]) if args else HERE
    offline = "--offline" in sys.argv
    if offline:
        data = json.loads(CACHE.read_text())
        retrieved = datetime.fromtimestamp(CACHE.stat().st_mtime, timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    else:
        data, retrieved = fetch()
    if isinstance(data, dict):
        data = [data]
    towns = []
    for (tid, name, state, lat, lon), d in zip(TOWNS, data):
        days = d["daily"]["time"]
        tmax = d["daily"]["temperature_2m_max"]
        per_year = {y: {"ge40": 0, "ge45": 0, "max_c": None, "max_date": None, "n": 0} for y in YEARS}
        for day, t in zip(days, tmax):
            if t is None:
                continue
            y = int(day[:4])
            py = per_year[y]
            py["n"] += 1
            if t >= 40.0:
                py["ge40"] += 1
            if t >= 45.0:
                py["ge45"] += 1
            if py["max_c"] is None or t > py["max_c"]:
                py["max_c"], py["max_date"] = round(t, 1), day
        years = [{"year": y, "days_ge40": v["ge40"], "days_ge45": v["ge45"], "max_c": v["max_c"], "max_date": v["max_date"], "days_with_data": v["n"]} for y, v in per_year.items()]
        hottest = max(years, key=lambda r: r["max_c"] or -99)
        towns.append({
            "id": tid, "name": name, "state": state,
            "coords": [round(lat, 4), round(lon, 4)],
            "grid_coords": [round(d.get("latitude", lat), 4), round(d.get("longitude", lon), 4)],
            "elevation_m": d.get("elevation"),
            "mean_days_ge40": round(sum(r["days_ge40"] for r in years) / len(years), 1),
            "mean_days_ge45": round(sum(r["days_ge45"] for r in years) / len(years), 1),
            "hottest": {"c": hottest["max_c"], "date": hottest["max_date"]},
            "years": years,
        })
    payload = {
        "module": "website",
        "dataset": "heat_days",
        "generated_at": retrieved,
        "simulated": False,
        "period": {"start": START, "end": END, "timezone": "Asia/Kolkata"},
        "variable": "ERA5 daily maximum 2 m air temperature (°C), nearest 0.25° grid cell",
        "thresholds_c": [40.0, 45.0],
        "caveat": "Reanalysis grid cells (0.25 degrees, about 25 to 28 km) smooth local extremes, so station maxima are often higher and these counts are conservative. Air temperature is not WBGT.",
        "towns": towns,
    }
    out.mkdir(parents=True, exist_ok=True)
    body = json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
    (out / "heat_days.json").write_text(body)
    # Merge into the folder's manifest so other datasets here (rivers_year.json) keep their entries.
    mpath = out / "manifest.json"
    manifest = json.loads(mpath.read_text()) if mpath.exists() else {
        "module": "website", "version": "1.0.0", "simulated": False, "seed": None, "sources": [], "files": [],
    }
    source = {
        "name": "Open-Meteo Historical Weather API (ERA5, Copernicus Climate Change Service)",
        "url": "https://open-meteo.com/en/docs/historical-weather-api",
        "license": "CC BY 4.0 (Open-Meteo); ERA5: Copernicus licence",
        "retrieved_at": retrieved,
        "request": url(),
    }
    entry = {"path": "heat_days.json", "rows": len(towns), "bytes": len(body.encode())}
    manifest["sources"] = [s for s in manifest.get("sources", []) if s.get("name") != source["name"]] + [source]
    manifest["files"] = [f for f in manifest.get("files", []) if f.get("path") != "heat_days.json"] + [entry]
    manifest["generated_at"] = max(s for s in (manifest.get("generated_at") or "", retrieved) if s)
    note = "Research data for /research/heat-frontier-hazard.html, built by research/data/build_heat.py (python3 research/data/build_heat.py --offline rebuilds from cache). Advisory context only; follow IMD heat warnings and NDMA guidance."
    others = [n for n in (manifest.get("notes") or "").split(" | ") if n and not n.startswith("Research data for /research/heat")]
    manifest["notes"] = " | ".join([note] + others)
    mpath.write_text(json.dumps(manifest, ensure_ascii=False, indent=1))
    for t in towns:
        print(f"{t['name']:15s} ≥40: {t['mean_days_ge40']:5.1f}/yr  ≥45: {t['mean_days_ge45']:4.1f}/yr  hottest {t['hottest']}")


if __name__ == "__main__":
    main()
