#!/usr/bin/env python
"""Stage 04 (development smoke-test) — one leave-one-year-out fold, primary target.

Proves the single-member MOS pipeline end to end for ONE held-out year: fit the Tx90p
threshold and the calibration models WITHOUT the held-out year, predict P(>=1 EHD in block)
on it, and score against the climatological and raw-forecast references. The decision
threshold is chosen on TRAIN predictions only (leakage rule). This is NOT the full
cross-validated result (that is stage 06 over all 20 folds); it is a correctness+sanity run.

Usage: python scripts/04_calibrate_singlefold.py --holdout 2022 --train-years 2013..2021
Outputs: results/tables/singlefold_<holdout>_scores.csv
         results/figures/singlefold_<holdout>_reliability.png
"""
from __future__ import annotations
import os, sys, logging, argparse, time
import numpy as np
import pandas as pd
import xarray as xr
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "src"))
from ehd.config import load_config, lead_blocks
from ehd import climatology as clim, features as feat, calibration as cal, verify as vf

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("stage04")


def load_land(cfg, ref):
    reg = xr.open_dataset(cfg["paths"]["mask"])["reg"].reindex(
        lat=ref.lat, lon=ref.lon, method="nearest")
    iy, ix = np.where(np.isfinite(reg.values))
    return iy, ix, reg.values


def assemble_years(cfg, years, thr, blocks, obs_vals, obs_lut, land_iy, land_ix):
    acc = {b: {"X": [], "y": []} for b in blocks}
    for y in years:
        fc = xr.open_dataarray(os.path.join(cfg["paths"]["outputs"],
                                            "forecast_daily_tx", f"fc_daily_tx_{y}.nc"))
        thr_i = thr.reindex(lat=fc.lat, lon=fc.lon, method="nearest")
        out = feat.assemble_year(fc, obs_vals, obs_lut, thr_i, blocks, land_iy, land_ix)
        for b in blocks:
            acc[b]["X"].append(out[b]["X"]); acc[b]["y"].append(out[b]["y_primary"])
    return {b: (np.concatenate(v["X"]), np.concatenate(v["y"])) for b, v in acc.items()}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--holdout", type=int, default=2022)
    ap.add_argument("--train-years", type=int, nargs="*",
                    default=list(range(2013, 2022)))
    ap.add_argument("--families", nargs="*", default=["raw", "climat", "logistic", "hgb"])
    args = ap.parse_args()
    cfg = load_config(); blocks = lead_blocks(cfg)
    res = cfg["paths"]["results"]

    da = clim.load_gridded_obs(cfg["paths"]["gridded_obs"])
    base = [y for y in range(cfg["obs"]["gridded_years"][0], cfg["obs"]["gridded_years"][1] + 1)
            if y != args.holdout]
    log.info("fitting Tx90p on %d obs years (holdout %d excluded)", len(base), args.holdout)
    thr = clim.ehd_threshold(clim.tx90p_climatology(da, base, 90, 15), 35.0, "min")
    obs_vals, obs_lut = feat.build_obs_index(da)

    ref = xr.open_dataarray(os.path.join(cfg["paths"]["outputs"], "forecast_daily_tx",
                                         f"fc_daily_tx_{args.holdout}.nc"))
    land_iy, land_ix, _ = load_land(cfg, ref)
    log.info("land cells: %d", land_iy.size)

    t0 = time.time()
    train = assemble_years(cfg, args.train_years, thr, blocks, obs_vals, obs_lut, land_iy, land_ix)
    test = assemble_years(cfg, [args.holdout], thr, blocks, obs_vals, obs_lut, land_iy, land_ix)
    log.info("assembled train(%d yr)+test in %.1fs", len(args.train_years), time.time() - t0)

    rows = []
    rel = {}
    for b in blocks:
        Xtr, ytr = train[b]; Xte, yte = test[b]
        clim_rate = np.nanmean(ytr)
        p_ref_clim = np.full(yte.shape, clim_rate)
        for fam in args.families:
            mdl = cal.make_model(fam)
            mdl.fit(Xtr, ytr)
            pte = mdl.predict_proba(Xte)
            ptr = mdl.predict_proba(Xtr)
            thr_star, _ = vf.best_threshold(ytr, ptr, "Peirce")   # threshold from TRAIN only
            bs = vf.binary_scores(yte, pte, thr_star)
            rows.append(dict(block=b, family=fam, n_test=int(np.isfinite(yte).sum()),
                             base_rate=round(float(np.nanmean(yte)), 3),
                             brier=round(vf.brier(yte, pte), 4),
                             bss_vs_clim=round(vf.brier_skill_score(yte, pte, p_ref_clim), 3),
                             auc=round(vf.auc(yte, pte), 3),
                             dec_thr=round(thr_star, 2),
                             POD=round(bs["POD"], 3), FAR=round(bs["FAR"], 3),
                             CSI=round(bs["CSI"], 3), Peirce=round(bs["Peirce"], 3),
                             freq_bias=round(bs["freq_bias"], 2)))
            if b in ("W1", "W4"):
                rel[(b, fam)] = vf.reliability_curve(yte, pte, 10)
    df = pd.DataFrame(rows)
    # BSS of calibrated families vs RAW forecast, per block
    raw_bs = {b: df[(df.block == b) & (df.family == "raw")].brier.values[0] for b in blocks}
    df["bss_vs_raw"] = df.apply(lambda r: round(1 - r.brier / raw_bs[r.block], 3), axis=1)
    out_csv = os.path.join(res, "tables", f"singlefold_{args.holdout}_scores.csv")
    df.to_csv(out_csv, index=False)
    log.info("scores:\n%s", df.to_string(index=False))

    # reliability figure for W1 & W4
    fig, axes = plt.subplots(1, 2, figsize=(11, 5))
    colors = {"raw": "#999999", "logistic": "#0072B2", "hgb": "#D55E00"}
    for ax, b in zip(axes, ("W1", "W4")):
        ax.plot([0, 1], [0, 1], "k--", lw=0.7)
        for fam in ("raw", "logistic", "hgb"):
            if (b, fam) in rel:
                conf, obs, cnt = rel[(b, fam)]
                ax.plot(conf, obs, "-o", color=colors[fam], label=fam, ms=4)
        ax.set_title(f"Reliability — {b} (holdout {args.holdout})")
        ax.set_xlabel("forecast probability"); ax.set_ylabel("observed frequency")
        ax.set_xlim(0, 1); ax.set_ylim(0, 1); ax.grid(alpha=0.3); ax.legend(fontsize=8)
    fig.tight_layout()
    fig.savefig(os.path.join(res, "figures", f"singlefold_{args.holdout}_reliability.png"), dpi=130)
    log.info("wrote reliability figure")


if __name__ == "__main__":
    main()