#!/usr/bin/env python3
"""Operational EHD forecast from one CFS initial time, using the certified chain.

The chain is exactly the one the evaluation certified, with nothing re-fitted on
the forecast itself:
  1. 6-hourly TMAX -> daily Tx by alignment A (buckets 06/12/18Z of day D and
     00Z of D+1), the alignment picked against observations.
  2. weekly max Tx at each of the 541 observation points, nearest coarse cell.
  3. BCSD: quantile-map the coarse weekly max onto the coarse observed
     distribution using ONLY the training years within +-30 days of this
     forecast's day of year, then add the corrected anomaly to the fine observed
     climatology of the same window.
  4. an EHD is declared at a point when the downscaled Tx reaches the local
     threshold plus DELTA = +0.40 C, the offset fitted on the training years.
  5. EXP1: a region is declared EHDE when more than 20% of its points are.

Only W1 is a certified binary product. W2-W6 are produced and printed, but the
evaluation found no threshold at any lead beyond W1 that meets the target box,
so they are labelled as severity margin, not as a yes/no forecast.
"""

from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import cfs_common as C
import build_obs as B
import grib_read as G
from bcsd import doy_window_mask, qmap_fit, qmap_apply

DELTA = 0.40                 # certified W1 offset, fitted on training years
REG_FRAC = 0.20              # EXP1
BOX = ("95:120", "0:30")
REG_NAME = {1: "S1 Tây Bắc", 2: "S2 Đông Bắc", 3: "S3 ĐB Bắc Bộ",
            4: "S4 Bắc Trung Bộ", 5: "S5 Nam Trung Bộ", 6: "S6 Tây Nguyên",
            7: "S7 Nam Bộ"}


def read_rt(indir, init):
    d0 = dt.datetime(int(init[:4]), int(init[4:6]), int(init[6:8]), int(init[8:10]))
    vts = [(d0 + dt.timedelta(hours=6 * (k + 1))).strftime("%Y%m%d%H")
           for k in range(C.NSTEP)]
    files = [os.path.join(indir, "%s.grb2" % v) for v in vts]
    miss = [f for f in files if not os.path.exists(f)]
    if miss:
        raise RuntimeError("thieu %d buoc 6 gio" % len(miss))
    g = C.flux_grid()
    ny, nx = len(g["wlat"]), len(g["wlon"])
    arr, nfill, _ = G.read_batch(files, vts, ":TMAX:2 m above ground:",
                                 [("TMAX", "2 m above ground")], BOX, ny, nx)
    if nfill < C.NSTEP:
        print("  canh bao: thieu %d/%d ban ghi" % (C.NSTEP - nfill, C.NSTEP))
    return arr[:, 0]                                     # (184, ny, nx), Kelvin


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--init", default="2026082900")
    ap.add_argument("--indir", default=None)
    ap.add_argument("--out", default=None)
    a = ap.parse_args()
    indir = a.indir or os.path.join(C.OUT, "RT", "cfs_%s" % a.init)
    out = a.out or os.path.join(C.OUT, "RT", "fcst_%s.json" % a.init)

    step = read_rt(indir, a.init)                        # (184, ny, nx) K
    daily = np.stack([np.nanmax(step[4 * k:4 * k + 4], 0)
                      for k in range(C.NLT)]) - 273.15   # (46, ny, nx) degC

    z = np.load(B.TX_NPZ)
    reg = z["reg"]
    olat, olon = z["lat"][z["jj"]], z["lon"][z["ii"]]
    g = C.flux_grid()
    pj = np.abs(olat[:, None] - g["wlat"][None, :]).argmin(1)
    pi = np.abs(olon[:, None] - g["wlon"][None, :]).argmin(1)
    _, cell = np.unique(pj * 1000 + pi, return_inverse=True)
    ncell = cell.max() + 1
    txp = daily[:, pj, pi]                               # (46, npt)

    wk = np.load(os.path.join(C.OUT, "Cube", "weekly.npz"))
    yrs, M, O, doy = wk["yrs"], wk["M"], wk["O"], wk["doy"]
    train = np.isin(yrs, C.TRAIN_YEARS)
    thr_all = np.load(B.thr_path("train"))["thr"]        # (365, npt)

    d0 = dt.date(int(a.init[:4]), int(a.init[4:6]), int(a.init[6:8]))
    res = {"init": a.init, "delta": DELTA, "reg_frac": REG_FRAC, "weeks": {}}
    print("\nDU BAO EHD — initial %s, thanh vien 01\n" % a.init)
    for w, (wname, idx) in enumerate(C.WEEKS):
        days = [d0 + dt.timedelta(days=int(k)) for k in idx]
        q = days[len(days) // 2].timetuple().tm_yday    # doy of window centre
        m_op = np.nanmax(txp[idx], 0)                    # (npt) model weekly max

        sel = doy_window_mask(doy[:, w], q) & train
        Mc = np.stack([np.nanmean(M[:, w][:, cell == c], 1) for c in range(ncell)], 1)
        Oc = np.stack([np.nanmean(O[:, w][:, cell == c], 1) for c in range(ncell)], 1)
        mc_op = np.array([np.nanmean(m_op[cell == c]) for c in range(ncell)])
        anom = np.empty(ncell)
        for c in range(ncell):
            mq, oq = qmap_fit(Mc[sel, c], Oc[sel, c])
            anom[c] = qmap_apply(mc_op[c], mq, oq) - Oc[sel, c].mean()
        fine_clim = np.nanmean(O[sel][:, w], axis=0)
        d1 = fine_clim + anom[cell]                      # (npt) downscaled Tx

        thr = thr_all[q - 1]
        margin = d1 - thr
        ev = margin >= DELTA
        P = {int(r): float(ev[reg == r].mean()) for r in np.unique(reg)}
        res["weeks"][wname] = {
            "valid_from": days[0].isoformat(), "valid_to": days[-1].isoformat(),
            "doy": int(q), "certified": wname == "W1",
            "tx": d1.round(2).tolist(), "thr": thr.round(2).tolist(),
            "margin": margin.round(2).tolist(), "event": ev.tolist(),
            "P": P, "EHDE": {str(r): bool(P[r] > REG_FRAC) for r in P},
            "n_point_event": int(ev.sum()), "frac_point": float(ev.mean())}
        tag = "" if wname == "W1" else "   (khong duoc chung nhan)"
        print("%s  %s .. %s   %d/%d o luoi vuot nguong (%.1f%%)%s"
              % (wname, days[0], days[-1], ev.sum(), len(ev), 100 * ev.mean(), tag))
        for r in sorted(P):
            print("     %-18s P=%5.1f%%   %s"
                  % (REG_NAME[r], 100 * P[r], "EHDE" if P[r] > REG_FRAC else "-"))
    os.makedirs(os.path.dirname(out), exist_ok=True)
    json.dump(res, open(out, "w"))
    print("\nsaved", out)


if __name__ == "__main__":
    main()