#!/usr/bin/env python3
"""Build the observed EHD threshold field from gridded Tx observations.

Tx90p_obs(doy,lat,lon): 90th percentile of Tx inside a 31-day window centred on
doy (doy-15 .. doy+15, wrapping around the year), pooled over the base years.
Tx_thres_obs = min(Tx90p_obs, 35) -- the union of HD (Tx>=35) and HWD (Tx>=Tx90p),
i.e. figure 1 of HW_HPC_1_cla.docx.

The base period is an argument on purpose: with k-fold-by-year validation the
climatology must be recomputed from the training years only, otherwise the
threshold leaks information from the test year.

Usage:
  make_obs_thresholds.py [--years 2003-2022] [--exclude 2010,2011] [--tag name]
"""
import argparse
import os
import numpy as np
import xarray as xr

OBS = "/home/tanpv/Data/QG.25.196/OBS/Tx_EC_0.25_2002-2023.nc"
MASK = "/home/tanpv/Data/QG.25.196/OBS/Mask_VN_S2S_EC_0.25.nc"
OUTDIR = "/home/tanpv/Progs/claude/HW_EC_10/outputs/Cli"

HALF_WINDOW = 15          # -> 31-day window
PCTL = 90
HD_THRESHOLD = 35.0       # Vietnamese hot-day convention (deg C)
NDOY = 365                # source uses a 365_day calendar, no leap days


def parse_years(spec, exclude):
    a, b = spec.split("-")
    yrs = list(range(int(a), int(b) + 1))
    if exclude:
        drop = {int(y) for y in exclude.split(",")}
        yrs = [y for y in yrs if y not in drop]
    return yrs


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--years", default="2003-2022", help="base period, inclusive")
    ap.add_argument("--exclude", default=None, help="years to leave out (fold test years)")
    ap.add_argument("--tag", default=None, help="output filename tag")
    args = ap.parse_args()

    years = parse_years(args.years, args.exclude)
    tag = args.tag or (f"{args.years}" + (f"_ex{args.exclude.replace(',', '-')}"
                                          if args.exclude else ""))
    os.makedirs(OUTDIR, exist_ok=True)

    ds = xr.open_dataset(OBS, decode_times=xr.coders.CFDatetimeCoder(use_cftime=True))
    tx = ds.tasmax
    yr = np.array([t.year for t in ds.time.values])
    doy = np.array([t.dayofyr for t in ds.time.values])

    sel = np.isin(yr, years)
    print(f"base years: {years[0]}-{years[-1]} (n={len(years)}), "
          f"{sel.sum()} days, grid {tx.shape[1]}x{tx.shape[2]}")

    data = tx.values[sel].astype("float32")          # (ndays, lat, lon)
    d_doy = doy[sel]

    nlat, nlon = data.shape[1], data.shape[2]
    q = np.full((NDOY, nlat, nlon), np.nan, dtype="float32")

    # index days by day-of-year once, then pool the 31-day window per doy
    by_doy = [np.where(d_doy == d + 1)[0] for d in range(NDOY)]
    for d in range(NDOY):
        offs = [(d + k) % NDOY for k in range(-HALF_WINDOW, HALF_WINDOW + 1)]
        idx = np.concatenate([by_doy[o] for o in offs])
        with np.errstate(all="ignore"):
            q[d] = np.nanpercentile(data[idx], PCTL, axis=0)
        if (d + 1) % 73 == 0:
            print(f"  doy {d+1}/{NDOY}  (n samples per point = {len(idx)})", flush=True)

    thres = np.minimum(q, HD_THRESHOLD)

    mask = xr.open_dataset(MASK)
    out = xr.Dataset(
        {
            "tx90p_obs": (("doy", "lat", "lon"), q,
                          {"long_name": f"Tx {PCTL}th percentile, {2*HALF_WINDOW+1}-day window",
                           "units": "C", "base_period": f"{years[0]}-{years[-1]}",
                           "excluded_years": args.exclude or "none"}),
            "tx_thres_obs": (("doy", "lat", "lon"), thres.astype("float32"),
                             {"long_name": "EHD threshold = min(Tx90p, 35C)",
                              "units": "C"}),
            "vn": (("lat", "lon"), mask.vn.values),
            "reg": (("lat", "lon"), mask.reg.values),
        },
        coords={"doy": np.arange(1, NDOY + 1), "lat": ds.lat.values, "lon": ds.lon.values},
    )
    out.attrs["source"] = OBS
    out.attrs["note"] = ("Tx_thres = min(Tx90p, 35C); EHD = 1 where Tx >= Tx_thres, "
                         "i.e. the union of HD (Tx>=35C) and HWD (Tx>=Tx90p).")

    path = os.path.join(OUTDIR, f"Qp_tx_obs_{tag}.nc")
    out.to_netcdf(path, encoding={v: {"zlib": True, "complevel": 4}
                                  for v in ["tx90p_obs", "tx_thres_obs"]})
    print(f"wrote {path}")

    land = np.isfinite(mask.vn.values) & (mask.vn.values > 0)
    print(f"land points: {land.sum()}")
    print(f"  Tx90p  over land: min {np.nanmin(q[:, land]):.1f}  "
          f"mean {np.nanmean(q[:, land]):.1f}  max {np.nanmax(q[:, land]):.1f} C")
    frac35 = np.nanmean(q[:, land] >= HD_THRESHOLD)
    print(f"  fraction of (doy,point) where Tx90p >= 35C (so 35C binds): {frac35:.3f}")


if __name__ == "__main__":
    main()