#!/usr/bin/env python3
"""Observed Tx, EHD thresholds and EHD/EHDE labels on the 0.25 deg VnGAD grid.

The observation grid IS the target grid: VnGAD tasmax is already 0.25 deg
(31 lon x 61 lat, lon 102.0-109.5, lat 8.5-23.5) and mask7reg_VnGAD_0.25.nc is on
exactly that grid, giving 541 land points in 7 climate regions. Nothing about the
observations is regridded, so no interpolation error is introduced on the target
side; the CFS fields are the ones brought to this grid later.

OBSERVATIONS ARE A TARGET ONLY. Nothing built here may become a model feature,
including observations available at initialisation time.

EHD threshold (CLAUDE.md): thr(doy, point) = min(Tx90p, 35 C), where Tx90p is the
90th percentile of observed Tx in a 31-day window centred on the day of year.

Two baselines are produced and compared:
  * `hist`  1969-1998, thirty years entirely BEFORE the forecast sample. Free of
    any circularity, and the operational way a normal is fixed, but a warming
    trend then pushes the 1999-2011 exceedance rate above 10 %.
  * `train` the training years inside 1999-2011. Contemporaneous, so the base rate
    stays near 10 %, and leak-free because validation years are excluded.
The `hist` baseline is what selects the validation years (a decision that must not
see the labels it will be judged on); `train` is what labels the samples.

Output: outputs/Obs/obs_tx.npz, outputs/Obs/thr_<baseline>.npz
"""

from __future__ import annotations

import argparse
import datetime as dt
import os
import sys

import numpy as np
from netCDF4 import Dataset, num2date

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import cfs_common as C

OBSD = os.path.join(C.OUT, "Obs")
TX_NPZ = os.path.join(OBSD, "obs_tx.npz")


def thr_path(tag):
    return os.path.join(OBSD, "thr_%s.npz" % tag)


def load_mask():
    f = Dataset(C.MASK_VNGAD)
    lon = np.asarray(f.variables["lon"][:], dtype="f8")
    lat = np.asarray(f.variables["lat"][:], dtype="f8")
    reg = np.ma.filled(f.variables["vn_7regions"][:], np.nan)
    f.close()
    land = np.isfinite(reg)
    return lon, lat, np.asarray(reg), land


def build_tx():
    """Observed Tx at the 541 land points for every calendar day 1961-2023."""
    os.makedirs(OBSD, exist_ok=True)
    lon, lat, reg, land = load_mask()
    f = Dataset(C.OBS_TX)
    t = f.variables["time"]
    dates = num2date(t[:], t.units, t.calendar)
    tx = np.ma.filled(f.variables["tasmax"][:], np.nan).astype("float32")
    f.close()

    jj, ii = np.where(land)
    pts = tx[:, jj, ii]                                   # (ntime, 541)
    years = np.array([d.year for d in dates], dtype="i4")
    doys = np.array([C.doy_noleap(dt.date(d.year, d.month, d.day)) for d in dates],
                    dtype="i4")
    ords = np.array([dt.date(d.year, d.month, d.day).toordinal() for d in dates],
                    dtype="i8")
    np.savez(TX_NPZ, tx=pts, year=years, doy=doys, ordinal=ords,
             lon=lon, lat=lat, reg=reg[jj, ii].astype("i4"), jj=jj, ii=ii)
    print("obs_tx: %s  points=%d  %s..%s  missing=%d"
          % (pts.shape, len(jj), dates[0], dates[-1], int(np.isnan(pts).sum())))


def build_thr(tag, years):
    """Tx90p on a 31-day circular window from `years`, then thr = min(Tx90p, 35)."""
    z = np.load(TX_NPZ)
    tx, yr, doy = z["tx"], z["year"], z["doy"]
    sel = np.isin(yr, years)
    tx, doy = tx[sel], doy[sel]
    npt = tx.shape[1]

    # bucket sample indices by day of year once, then pool the 31-day window
    bydoy = [np.where(doy == d)[0] for d in range(1, 366)]
    qp = np.empty((365, npt), dtype="float32")
    for d in range(365):
        idx = np.concatenate([bydoy[(d + o) % 365]
                              for o in range(-C.WIN_HALF, C.WIN_HALF + 1)])
        qp[d] = np.nanpercentile(tx[idx], C.PCTL, axis=0)
    thr = C.tx_thres(qp)
    np.savez(thr_path(tag), qp=qp, thr=thr, years=np.array(sorted(years)))
    cap = float((thr >= C.HOT_DAY_C - 1e-6).mean())
    print("thr_%s: years=%d  Tx90p mean=%.2f  thr mean=%.2f  capped by 35C: %.1f%%"
          % (tag, len(years), float(qp.mean()), float(thr.mean()), 100 * cap))


def ehd_of(tag):
    """Daily EHD label for every observed day, using threshold set `tag`."""
    z = np.load(TX_NPZ)
    thr = np.load(thr_path(tag))["thr"]
    return z["tx"] >= thr[z["doy"] - 1], z


def base_rates(tag):
    """Observed EHD frequency per year -- the statistic the split is chosen on."""
    ehd, z = ehd_of(tag)
    yr = z["year"]
    print("\nEHD base rate by year (threshold %s):" % tag)
    tot = []
    for y in C.YEARS:
        s = yr == y
        r = float(np.nanmean(ehd[s]))
        tot.append(r)
        print("  %d  %.4f" % (y, r))
    print("  1999-2011 pooled: %.4f" % float(np.nanmean(ehd[np.isin(yr, C.YEARS)])))
    return np.array(tot)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--stage", default="all",
                    choices=["tx", "thr_hist", "rates", "all"])
    ap.add_argument("--train-years", type=int, nargs="*", default=None)
    a = ap.parse_args()

    if a.stage in ("tx", "all"):
        build_tx()
    if a.stage in ("thr_hist", "all"):
        build_thr("hist", list(range(1969, 1999)))
    if a.stage in ("rates", "all"):
        base_rates("hist")
    if a.train_years:
        build_thr("train", a.train_years)


if __name__ == "__main__":
    main()