#!/usr/bin/env python3
"""Stage 1, section 3.3: Tx90p and Tx_thres for model (EC) and observations.

Tx90p_obs(lon,lat,doy)      from the continuous VnGAD-on-EC-grid Tx series
Tx90p_mod(lon,lat,lt,doy)   from EC_Combined mx2t6, one percentile per lead time

Both use the 31-day window (doy-15, doy+15) pooled over all 20 hindcast years,
then Tx_thres = min(Tx90p, 35 C), i.e. the union of the hot-day and heat-wave-day
definitions of section 2 of HW_HPC_1, as drawn in its figure 1.

Only the 541 land points of Mask_VN_S2S_EC_0.25.nc are computed; the rest of the
grid is written as NaN.
"""

from __future__ import annotations

import argparse
import os
import sys
import time

import numpy as np
from joblib import Parallel, delayed

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


def _pctl(block, q):
    """Percentile along axis 0, skipping the nan branch when it is not needed."""
    if not np.isnan(block).any():
        return np.percentile(block, q, axis=0)
    return np.nanpercentile(block, q, axis=0)


def obs_percentile(grid, years, pctl=C.PCTL):
    """Tx90p_obs as (365, npt)."""
    tx, doy, _ = C.read_obs_continuous(grid, years)
    out = np.full((365, grid.npt), np.nan, dtype="f4")
    for d in range(1, 366):
        sel = np.isin(doy, C.window_doys(d))
        out[d - 1] = _pctl(tx[sel], pctl)
    return out


def _load_year(year, grid):
    """mx2t6 (ninit, 46, npt) and the valid-day doy (ninit, 46) for one year."""
    inits = C.init_dates(year)
    tx = np.empty((len(inits), C.NLT, grid.npt), dtype="f4")
    doy = np.empty((len(inits), C.NLT), dtype="i4")
    for i, init in enumerate(inits):
        tx[i] = C.read_ec_points(init, "mx2t6", grid)
        doy[i] = [C.doy_noleap(d) for d in C.valid_days(init)]
    return tx, doy


def _mod_percentile_one_lt(tx_lt, doy_lt, pctl):
    """Tx90p for a single lead time: (nsample, npt) + (nsample,) -> (365, npt)."""
    npt = tx_lt.shape[1]
    out = np.full((365, npt), np.nan, dtype="f4")
    for d in range(1, 366):
        sel = np.isin(doy_lt, C.window_doys(d))
        if sel.sum() == 0:
            continue
        out[d - 1] = _pctl(tx_lt[sel], pctl)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--jobs", type=int, default=20)
    ap.add_argument("--years", type=int, nargs="*", default=C.YEARS)
    ap.add_argument("--outdir", default=os.path.join(C.OUT, "Cli"))
    ap.add_argument("--tag", default="")
    args = ap.parse_args()

    grid = C.Grid()
    t0 = time.time()

    print("[obs] percentile over %d years ..." % len(args.years), flush=True)
    qp_obs = obs_percentile(grid, args.years)
    thr_obs = C.tx_thres(qp_obs)
    print("[obs] done in %.1fs  Tx90p range %.1f..%.1f" %
          (time.time() - t0, np.nanmin(qp_obs), np.nanmax(qp_obs)), flush=True)

    print("[mod] reading EC_Combined mx2t6 ...", flush=True)
    t1 = time.time()
    parts = Parallel(n_jobs=min(args.jobs, len(args.years)), verbose=0)(
        delayed(_load_year)(y, grid) for y in args.years)
    tx = np.concatenate([p[0] for p in parts], axis=0)
    doy = np.concatenate([p[1] for p in parts], axis=0)
    del parts
    print("[mod] loaded %s in %.1fs (%.0f MB)" %
          (tx.shape, time.time() - t1, tx.nbytes / 1e6), flush=True)

    print("[mod] percentile per lead time ...", flush=True)
    t2 = time.time()
    res = Parallel(n_jobs=args.jobs, verbose=1)(
        delayed(_mod_percentile_one_lt)(tx[:, k], doy[:, k], C.PCTL)
        for k in range(C.NLT))
    qp_mod = np.stack(res, axis=0)                    # (46, 365, npt)
    thr_mod = C.tx_thres(qp_mod)
    print("[mod] done in %.1fs  Tx90p range %.1f..%.1f" %
          (time.time() - t2, np.nanmin(qp_mod), np.nanmax(qp_mod)), flush=True)

    os.makedirs(args.outdir, exist_ok=True)
    doys = np.arange(1, 366, dtype="i4")
    lts = np.arange(C.NLT, dtype="i4")

    C.write_nc(
        os.path.join(args.outdir, "Qp_tx_obs%s.nc" % args.tag),
        [("doy", doys), ("lat", grid.lat), ("lon", grid.lon)],
        [("Tx90p_obs", ("doy", "lat", "lon"), grid.to_grid(qp_obs),
          {"long_name": "90th percentile of observed Tx, 31-day window", "units": "C"}),
         ("Tx_thres_obs", ("doy", "lat", "lon"), grid.to_grid(thr_obs),
          {"long_name": "min(Tx90p_obs, 35 C)", "units": "C"})],
        {"source": "OBS/Tx_EC_0.25_2002-2023.nc", "years": str(args.years),
         "window_days": 2 * C.WIN_HALF + 1, "percentile": C.PCTL})

    C.write_nc(
        os.path.join(args.outdir, "Qp_tx_mod%s.nc" % args.tag),
        [("lt", lts), ("doy", doys), ("lat", grid.lat), ("lon", grid.lon)],
        [("Tx90p_mod", ("lt", "doy", "lat", "lon"), grid.to_grid(qp_mod),
          {"long_name": "90th percentile of EC mx2t6, 31-day window, per lead time",
           "units": "C"}),
         ("Tx_thres_mod", ("lt", "doy", "lat", "lon"), grid.to_grid(thr_mod),
          {"long_name": "min(Tx90p_mod, 35 C)", "units": "C"})],
        {"source": "EC_Combined mx2t6", "years": str(args.years),
         "window_days": 2 * C.WIN_HALF + 1, "percentile": C.PCTL})

    print("[done] total %.1fs -> %s" % (time.time() - t0, args.outdir), flush=True)


if __name__ == "__main__":
    main()