#!/usr/bin/env python3
"""Re-score a finished national run under a frequency-unbiased decision rule.

Maximising CSI is the right rule when the only thing that matters is the overlap
of forecast and observed events, but for a rare event it always over-forecasts:
the national logreg run reached POD 0.67 at W4 with FBI 3.3.  The brief asks for
FBI close to 1, so this script adds a second decision rule in which the threshold
is set, per forecast week, so that the number of forecast events matches the
number of observed events.  The threshold still comes from the out-of-fold
predictions inside the training years, never from the held-out years.

It also produces the reliability diagram and Brier skill score, which stage 1
could not provide because the raw EC control run is deterministic.

Reads the shards a national run left behind, so it costs seconds and touches
nothing that is still running.
"""

from __future__ import annotations

import argparse
import glob
import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ehd_common as C
import ehd_metrics as M
import train_ai as T
from train_national import RES_DIR, shard_dir


# These now live in train_ai so that train_national.py uses exactly the same
# threshold logic on its main path; kept here as aliases so this script, which
# re-scores runs finished before that move, keeps working unchanged.
population_weights = T.population_weights
weighted_quantile = T.weighted_quantile
weighted_scores = T.weighted_scores
bias_one_thresholds = T.bias_one_thresholds
csi_thresholds = T.csi_thresholds
platt_calibrate = T.platt_calibrate


def load_run(tag):
    d = shard_dir(tag)
    oof = [np.load(f) for f in sorted(glob.glob(os.path.join(d, "oof_*.npz")))]
    if not oof:
        raise SystemExit("no out-of-fold files in %s" % d)
    o = {"score": np.concatenate([x["score"] for x in oof]),
         "evt": np.concatenate([x["evt"] for x in oof]),
         "week": np.concatenate([x["week"] for x in oof]).astype("i4")}
    pred = [np.load(f) for f in sorted(glob.glob(os.path.join(d, "pred_*.npz")))]
    if not pred:
        raise SystemExit("no prediction files in %s" % d)
    p = {k: np.concatenate([x[k] for x in pred])
         for k in ("score", "evt", "week", "lt", "point", "capable", "init")}
    p["score"] = p["score"].astype("f8")
    p["week"] = p["week"].astype("i4")
    p["init"] = p["init"].astype("i8")
    return o, p


def report(score, pred, evt, info, label, out):
    rows = {"daily_all": T.evaluate(score, pred, evt, info),
            "daily_capable": T.evaluate(score, pred, evt, info,
                                        info["capable"].astype(bool))}
    ws, wp, wo, wi = T.weekly_event(score, pred, evt, info)
    rows["weekly_all"] = T.evaluate(ws, wp, wo, wi)
    print("\n########## %s ##########" % label, flush=True)
    for name, tab in rows.items():
        print(T.fmt_table(name, tab), flush=True)
    out[label] = rows
    return rows


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--tag", default="")
    ap.add_argument("--name", default="")
    ap.add_argument("--neg-keep", type=float, default=0.15)
    args = ap.parse_args()

    oof, p = load_run(args.tag)
    ok = np.isfinite(p["score"]) & np.isfinite(p["evt"])
    score = p["score"][ok]
    evt = p["evt"][ok]
    info = {k: p[k][ok] for k in ("week", "lt", "point", "capable", "init")}
    out = {}

    w_oof = population_weights(oof["evt"], args.neg_keep)
    true_rate = float((w_oof * oof["evt"]).sum() / w_oof.sum())
    print("[weights] out-of-fold base rate %.4f raw -> %.4f after undoing the thinning"
          % (float(oof["evt"].mean()), true_rate), flush=True)
    th_csi = csi_thresholds(oof["score"], oof["evt"], oof["week"], w_oof)
    th_fbi = bias_one_thresholds(oof["score"], oof["evt"], oof["week"], w_oof)
    print("[thresholds] CSI-optimal : %s"
          % {C.WEEK_NAMES[k]: round(v, 4) for k, v in th_csi.items()}, flush=True)
    print("[thresholds] FBI ~ 1     : %s"
          % {C.WEEK_NAMES[k]: round(v, 4) for k, v in th_fbi.items()}, flush=True)

    report(score, T.apply_thresholds(score, info["week"], th_csi), evt, info,
           "rule A - threshold maximises CSI", out)
    report(score, T.apply_thresholds(score, info["week"], th_fbi), evt, info,
           "rule B - threshold gives FBI ~ 1", out)

    # probabilistic scores, which the deterministic EC run could not provide
    prob = platt_calibrate(oof["score"], oof["evt"], w_oof, score)
    print("\n[calibration] mean forecast probability %.4f vs observed rate %.4f"
          % (float(prob.mean()), float((evt > 0.5).mean())), flush=True)
    rel = {}
    for w in range(len(C.WEEKS)):
        m = info["week"] == w
        rel[C.WEEK_NAMES[w]] = {
            "reliability": M.reliability(prob[m], evt[m] > 0.5),
            "bss": M.brier_skill_score(prob[m], evt[m] > 0.5),
            "auc": M.auc_score(score[m], evt[m] > 0.5)}
    print("\n[probabilistic] Brier skill score against sample climatology")
    for w, v in rel.items():
        print("    %-4s BSS = %7.4f   AUC = %.3f" % (w, v["bss"], v["auc"]))
    out["probabilistic"] = rel

    name = args.name or ("run%s" % (args.tag or "_default"))
    path = os.path.join(RES_DIR, "recalibrated_%s.json" % name)
    with open(path, "w") as fh:
        json.dump({"thresholds_csi": {str(k): float(v) for k, v in th_csi.items()},
                   "thresholds_fbi1": {str(k): float(v) for k, v in th_fbi.items()},
                   "results": out}, fh, indent=1, default=float)
    print("\n[done] -> %s" % path, flush=True)


if __name__ == "__main__":
    main()