#!/usr/bin/env python3
"""Score downscaled Tx forecasts as EHD / EHDE event forecasts.

RMSE is not the objective: the project targets POD > 0.6, POFD < 0.2, FAR < 0.2,
FBI ~ 1, CSI > 0.4 for extreme-heat events. A method can win on RMSE and still lose
on the event scores, so every downscaling method is scored both ways.

EHD is defined per grid point by  Tx >= Tx_thres  with Tx_thres = min(Tx90p, 35C).
The threshold of each fold is rebuilt from that fold's training years only
(Qp_tx_obs_fold_<firstyear>.nc), so a test year never sees its own climatology.

EHDE (EXP1) aggregates to region level: for region i, P_i is the fraction of its land
points with EHD = 1 that day; the region is flagged when P_i exceeds --cover (0.20).

Usage:
  evaluate_ehd.py --pred outputs/downscale/pred_bcsd.npy [--pred ...] [--cover 0.2]
"""
import argparse
import json
import os
import sys

import numpy as np
import xarray as xr

ROOT = "/home/tanpv/Progs/claude/HW_EC_10"
OBS = f"{ROOT}/outputs/OBS_4EC15/tx_obs_aligned.nc"
STATIC = f"{ROOT}/outputs/Cli/static_fields_0.25.nc"
CLI = f"{ROOT}/outputs/Cli"

WEEKS = [("W1", 0, 11), ("W2", 11, 18), ("W3", 18, 25),
         ("W4", 25, 32), ("W5", 32, 39), ("W6", 39, 46)]
FOLD_START = [2003, 2007, 2011, 2015, 2019]
YEARS = list(range(2003, 2023))


def contingency(pred, obs):
    """POD / POFD / FAR / FBI / CSI / ETS / EDI from two boolean arrays."""
    ok = np.isfinite(pred.astype("float32")) & np.isfinite(obs.astype("float32"))
    p, o = pred[ok].astype(bool), obs[ok].astype(bool)
    H = int(np.sum(p & o))
    M = int(np.sum(~p & o))
    F = int(np.sum(p & ~o))
    C = int(np.sum(~p & ~o))
    n = H + M + F + C
    if n == 0 or (H + M) == 0:
        return dict(H=H, M=M, F=F, C=C, n=n, base=np.nan, POD=np.nan, POFD=np.nan,
                    FAR=np.nan, FBI=np.nan, CSI=np.nan, ETS=np.nan, EDI=np.nan)
    pod = H / (H + M)
    pofd = F / (F + C) if (F + C) else np.nan
    far = F / (H + F) if (H + F) else np.nan
    fbi = (H + F) / (H + M)
    csi = H / (H + M + F) if (H + M + F) else np.nan
    hr = (H + M) * (H + F) / n
    ets = (H - hr) / (H + M + F - hr) if (H + M + F - hr) != 0 else np.nan
    edi = np.nan
    if 0 < pod < 1 and 0 < pofd < 1:
        edi = (np.log(pofd) - np.log(pod)) / (np.log(pofd) + np.log(pod))
    return dict(H=H, M=M, F=F, C=C, n=n, base=(H + M) / n, POD=pod, POFD=pofd,
                FAR=far, FBI=fbi, CSI=csi, ETS=ets, EDI=edi)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--pred", action="append", required=True,
                    help="prediction .npy, shaped (nyear, ninit, nlead, nland)")
    ap.add_argument("--cover", type=float, default=0.20,
                    help="EXP1 areal coverage threshold for a regional EHDE")
    ap.add_argument("--out", default=f"{ROOT}/outputs/downscale/ehd_scores.json")
    args = ap.parse_args()

    o = xr.open_dataset(OBS)
    st = xr.open_dataset(STATIC)
    land = (np.isfinite(st.vn.values) & (st.vn.values > 0)).ravel()
    land_idx = np.where(land)[0]
    reg = st.reg.values.ravel()[land_idx]
    obs = o.tx_obs.values.reshape(len(YEARS), o.sizes["init"], o.sizes["lead"],
                                  -1)[..., land_idx]
    doy = o.doy.values
    nyear, ninit, nlead, nland = obs.shape

    # per-fold threshold, assembled into one (nyear, ninit, nlead, nland) field
    thr = np.empty_like(obs)
    for k, y0 in enumerate(FOLD_START):
        f = np.arange(k * 4, k * 4 + 4)
        path = f"{CLI}/Qp_tx_obs_fold_{y0}.nc"
        if not os.path.exists(path):
            print(f"missing {path}", file=sys.stderr)
            return 1
        t = xr.open_dataset(path).tx_thres_obs.values.reshape(365, -1)[:, land_idx]
        thr[f] = t[doy[f] - 1]
    ehd_obs = obs >= thr

    results = {}
    for path in args.pred:
        name = os.path.basename(path).replace("pred_", "").replace(".npy", "")
        pred = np.load(path)
        if pred.shape != obs.shape:
            print(f"{name}: shape {pred.shape} != {obs.shape}", file=sys.stderr)
            return 1
        ehd_pred = pred >= thr
        valid = np.isfinite(pred) & np.isfinite(obs)

        res = {"gridpoint": {}, "regional_EHDE": {}}
        for wname, l0, l1 in WEEKS:
            v = valid[:, :, l0:l1]
            res["gridpoint"][wname] = contingency(
                ehd_pred[:, :, l0:l1][v], ehd_obs[:, :, l0:l1][v])
        res["gridpoint"]["ALL"] = contingency(ehd_pred[valid], ehd_obs[valid])

        # EXP1: regional coverage -> regional event
        for wname, l0, l1 in WEEKS + [("ALL", 0, nlead)]:
            po, pp = [], []
            for r in range(1, 8):
                m = reg == r
                if m.sum() == 0:
                    continue
                vo = ehd_obs[:, :, l0:l1][..., m].mean(axis=-1)
                vp = ehd_pred[:, :, l0:l1][..., m].mean(axis=-1)
                po.append(vo.ravel())
                pp.append(vp.ravel())
            res["regional_EHDE"][wname] = contingency(
                np.concatenate(pp) > args.cover, np.concatenate(po) > args.cover)

        results[name] = res
        g = res["gridpoint"]["ALL"]
        e = res["regional_EHDE"]["ALL"]
        print(f"{name:10s} gridpoint: base={g['base']:.3f} POD={g['POD']:.3f} "
              f"POFD={g['POFD']:.3f} FAR={g['FAR']:.3f} FBI={g['FBI']:.3f} "
              f"CSI={g['CSI']:.3f}", flush=True)
        print(f"{'':10s} EXP1 reg : base={e['base']:.3f} POD={e['POD']:.3f} "
              f"POFD={e['POFD']:.3f} FAR={e['FAR']:.3f} FBI={e['FBI']:.3f} "
              f"CSI={e['CSI']:.3f}", flush=True)

    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    with open(args.out, "w") as fh:
        json.dump(results, fh, indent=2)
    print(f"wrote {args.out}")
    return 0


if __name__ == "__main__":
    sys.exit(main())