#!/usr/bin/env python3
"""Downscaling ladder for weekly max Tx: control, BCSD, and a conditional ML model.

    D0  nearest coarse cell, no correction                      (control)
    D1  BCSD: quantile-map the coarse forecast onto the coarse
        observed distribution, then add the corrected ANOMALY to
        the fine observed climatology                            (Wood et al. 2002/2004)
    D2  LightGBM regression on the multivariate coarse state      (conditional)

Each rung is scored twice, because the two answers routinely disagree:
  * as a value    -- ME / MAE / RMSE / correlation, and
  * as an event   -- the frequency with which the downscaled field crosses the
                     EHD threshold, relative to observations.
A method can win on RMSE and still be useless here: regression shrinks toward
climatology, and a shrunk field thresholded at Tx90p produces almost no events.
The exceedance ratio is what exposes that, so it is reported for every rung.

Fitting uses training years only, including every climatology and every quantile map.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import warnings

import numpy as np

warnings.filterwarnings("ignore")
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import cfs_common as C
import build_obs as B
from weekly_skill import auc, climz

RES = os.environ.get("BCSD_RES",
                     os.path.join(C.OUT, "AI_Results", "bcsd.json"))
DOY_HALF = 30          # +/- 30 days pooled for the quantile map and climatologies
NQ = 41                # quantile grid


def doy_window_mask(doy, q, half=DOY_HALF):
    d = np.abs(doy - q)
    return np.minimum(d, 365 - d) <= half


def qmap_fit(mod, obs):
    """Empirical quantile map: model value -> observed value of the same rank."""
    qs = np.linspace(0, 100, NQ)
    return np.percentile(mod, qs), np.percentile(obs, qs)


def qmap_apply(x, mq, oq):
    return np.interp(x, mq, oq)


def metrics(pred, obs, thr, ehde):
    """Value scores, then how often the field crosses the EHD threshold."""
    ok = np.isfinite(pred) & np.isfinite(obs)
    p, o = pred[ok], obs[ok]
    out = dict(ME=float((p - o).mean()), MAE=float(np.abs(p - o).mean()),
               RMSE=float(np.sqrt(((p - o) ** 2).mean())),
               CORR=float(np.corrcoef(p, o)[0, 1]))
    fc = pred >= thr
    out["exc_fcst"] = float(np.nanmean(fc))
    out["exc_obs"] = float(np.nanmean(ehde))
    out["exc_ratio"] = out["exc_fcst"] / max(out["exc_obs"], 1e-9)
    a = float((fc & ehde).sum()); b = float((fc & ~ehde).sum())
    c = float((~fc & ehde).sum()); d = float((~fc & ~ehde).sum())
    out["POD"] = a / max(a + c, 1); out["POFD"] = b / max(b + d, 1)
    out["FAR"] = b / max(a + b, 1); out["FBI"] = (a + b) / max(a + c, 1)
    out["CSI"] = a / max(a + b + c, 1)
    H, F = out["POD"], max(out["POFD"], 1e-9)
    with np.errstate(divide="ignore", invalid="ignore"):
        out["SEDI"] = float((np.log(F) - np.log(max(H, 1e-9)) - np.log(1 - F)
                             + np.log(max(1 - H, 1e-9)))
                            / (np.log(F) + np.log(max(H, 1e-9)) + np.log(1 - F)
                               + np.log(max(1 - H, 1e-9))))
    return out


def part_path(w):
    return "%s.w%d.json" % (RES, w)


def field_path(w):
    return "%s.w%d.fields.npz" % (RES, w)


def merge():
    """Combine the per-week part files written by the array tasks."""
    res = {}
    for w, wname in enumerate(C.WEEK_NAMES):
        f = part_path(w)
        if not os.path.exists(f):
            print("THIEU %s" % f)
            continue
        res[wname] = json.load(open(f))
    json.dump(res, open(RES, "w"), indent=1)
    print("%-3s %-3s %7s %7s %7s %7s %9s %6s %6s %6s %6s"
          % ("wk", "", "ME", "MAE", "RMSE", "CORR", "exc_ratio",
             "POD", "FAR", "CSI", "SEDI"))
    for wname, row in res.items():
        for tag in ("D0", "D1", "D2"):
            r = row[tag]
            print("%-3s %-3s %7.3f %7.3f %7.3f %7.4f %9.3f %6.3f %6.3f %6.3f %6.3f"
                  % (wname, tag, r["ME"], r["MAE"], r["RMSE"], r["CORR"],
                     r["exc_ratio"], r["POD"], r["FAR"], r["CSI"], r["SEDI"]))
    print("\nsaved", RES)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--week", type=int, default=-1,
                    help="compute only this week (0-5) into its own part file")
    ap.add_argument("--merge", action="store_true")
    ap.add_argument("--save-fields", action="store_true",
                    help="also dump the downscaled fields, for threshold tuning")
    a = ap.parse_args()
    if a.merge:
        return merge()
    ncore = int(os.environ.get("SLURM_CPUS_PER_TASK", 8))

    wk = np.load(os.path.join(C.OUT, "Cube", "weekly.npz"))
    yrs, M, O, E, doy = wk["yrs"], wk["M"], wk["O"], wk["E"], wk["doy"]
    fx = np.load(os.path.join(C.OUT, "Cube", "weekly_feats.npz"))["X"]
    pgf = os.environ.get("PGB_WEEKLY",
                         os.path.join(C.OUT, "Cube", "pgb_weekly.npz"))
    pg = np.load(pgf) if os.path.exists(pgf) else None

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

    thr_all = np.load(B.thr_path("train"))["thr"]          # (365, npt)
    train = np.isin(yrs, C.TRAIN_YEARS)
    valid = np.isin(yrs, C.VALID_YEARS)

    import lightgbm as lgb
    res = {}
    for w, wname in enumerate(C.WEEK_NAMES):
        if a.week >= 0 and w != a.week:
            continue
        d = doy[:, w]
        thr = thr_all[d - 1]                                # (ninit, npt)
        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)

        d0 = M[:, w].copy()                                 # D0: raw nearest cell
        d1 = np.full_like(d0, np.nan)                       # D1: BCSD

        for q in range(1, 366):
            sel = doy_window_mask(d, q) & train
            tgt = (d == q)
            if sel.sum() < 50 or tgt.sum() == 0:
                continue
            # BC on the coarse cell, then SD onto the fine observed climatology
            anom = np.empty((int(tgt.sum()), ncell))
            for c in range(ncell):
                mq, oq = qmap_fit(Mc[sel, c], Oc[sel, c])
                anom[:, c] = qmap_apply(Mc[tgt, c], mq, oq) - Oc[sel, c].mean()
            # fine-scale observed climatology for THIS week, per point (npt,).
            # O is (ninit, nweek, npt): indexing axis 1 by p walked the WEEK
            # axis and collapsed every point into one scalar.
            fine_clim = np.nanmean(O[sel][:, w], axis=0)
            d1[tgt] = fine_clim[None, :] + anom[:, cell]

        # D2: conditional regression on the multivariate coarse state
        Zf = np.stack([climz(fx[:, [w], v], doy[:, [w]], yrs, train)[:, 0]
                       for v in range(fx.shape[2])], -1)
        blocks = [Zf]
        if pg is not None:
            blocks.append(np.stack([climz(pg["X"][:, [w], v], doy[:, [w]], yrs, train)[:, 0]
                                    for v in range(pg["X"].shape[2])], -1))
        dd = d.astype(float)
        h = np.stack([np.cos(2 * np.pi * dd / 365), np.sin(2 * np.pi * dd / 365)], -1)
        blocks.append(np.repeat(h[:, None, :], Zf.shape[1], 1))
        blocks.append(np.repeat(np.arange(Zf.shape[1])[None, :, None],
                                Zf.shape[0], 0).astype("float32"))   # point identity
        F = np.concatenate(blocks, -1)

        def mats(mask):
            f = F[mask].reshape(-1, F.shape[-1])
            t = O[mask, w].ravel()
            ok = np.isfinite(f).all(1) & np.isfinite(t)
            return f[ok], t[ok], ok
        ftr, ytr, _ = mats(train)
        fva, yva, okva = mats(valid)
        gb = lgb.LGBMRegressor(n_estimators=600, learning_rate=0.05, num_leaves=127,
                               min_child_samples=100, subsample=0.8,
                               colsample_bytree=0.8, verbose=-1, n_jobs=ncore)
        gb.fit(ftr, ytr)
        d2 = np.full(F.shape[:2], np.nan)
        tmp = np.full(int(valid.sum()) * F.shape[1], np.nan)
        tmp[okva] = gb.predict(fva)
        d2[valid] = tmp.reshape(int(valid.sum()), F.shape[1])

        if a.save_fields:
            # d1 is filled for every init whose doy window had enough training
            # data, valid years included, so the threshold can be fitted on the
            # training years and applied out of sample from this one file.
            np.savez(field_path(w), d1=d1.astype("float32"),
                     d0=d0.astype("float32"), thr=thr.astype("float32"),
                     O=O[:, w].astype("float32"), E=E[:, w], yrs=yrs,
                     doy=doy[:, w], reg=reg)

        row = {}
        for tag, pred in (("D0", d0), ("D1", d1), ("D2", d2)):
            row[tag] = metrics(pred[valid], O[valid, w], thr[valid], E[valid, w])
        res[wname] = row
        print("\n=== %s ===" % wname, flush=True)
        print("%-3s %7s %7s %7s %7s %9s %6s %6s %6s %6s"
              % ("", "ME", "MAE", "RMSE", "CORR", "exc_ratio", "POD", "FAR", "CSI", "SEDI"))
        for tag in ("D0", "D1", "D2"):
            r = row[tag]
            print("%-3s %7.3f %7.3f %7.3f %7.4f %9.3f %6.3f %6.3f %6.3f %6.3f"
                  % (tag, r["ME"], r["MAE"], r["RMSE"], r["CORR"], r["exc_ratio"],
                     r["POD"], r["FAR"], r["CSI"], r["SEDI"]), flush=True)

    os.makedirs(os.path.dirname(RES), exist_ok=True)
    if a.week >= 0:
        out = part_path(a.week)
        json.dump(res[C.WEEK_NAMES[a.week]], open(out, "w"), indent=1)
    else:
        out = RES
        json.dump(res, open(out, "w"), indent=1)
    print("\nsaved", out)


if __name__ == "__main__":
    main()