#!/usr/bin/env python3
"""Where to put the decision threshold, and whether the target box is reachable.

The BCSD field is scored as `d1 - thr`: how far the downscaled Tx sits above the
local EHD threshold. Declaring an event at `score >= delta` traces one ROC curve,
and delta is the only free parameter. It is fitted on TRAINING years and applied
unchanged to the validation years, because a threshold picked on the data it is
scored against is not a forecast.

Three fitting rules, because they disagree and the disagreement is the point:
  csi   maximise CSI on the training years
  fbi   set FBI = 1 (the natural operating point, exc_ratio ~ 1)
  far   the sharpest threshold whose training FAR is under 0.20
  box   among thresholds meeting the whole target box ON TRAINING years, the one
        with the best training CSI -- the only rule that can honestly claim the
        target, since the box is checked where the threshold is fitted

Scored twice: per grid point, and per region under EXP1 (a region has EHDE=1 when
more than 20% of its points do), which is the operational unit and a far easier
target -- regional events are commoner and less noisy than point events.
"""

from __future__ import annotations

import argparse
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

FLD = os.environ.get("BCSD_RES",
                     os.path.join(C.OUT, "AI_Results", "bcsd_fluxonly.json"))
RES = os.environ.get("THR_RES",
                     os.path.join(C.OUT, "AI_Results", "threshold.json"))
REG_FRAC = 0.20                      # EXP1: a region fires when >20% of it fires
DELTAS = np.round(np.arange(-4.0, 6.01, 0.05), 2)


def contingency(fc, ob):
    """fc, ob: boolean arrays of the same shape. Returns the usual scores."""
    fc, ob = fc.ravel(), ob.ravel()
    a = float(np.sum(fc & ob))                       # hits
    b = float(np.sum(fc & ~ob))                      # false alarms
    c = float(np.sum(~fc & ob))                      # misses
    d = float(np.sum(~fc & ~ob))                     # correct negatives
    pod = a / (a + c) if a + c else np.nan
    pofd = b / (b + d) if b + d else np.nan
    far = b / (a + b) if a + b else np.nan
    fbi = (a + b) / (a + c) if a + c else np.nan
    csi = a / (a + b + c) if a + b + c else np.nan
    out = dict(n=a + b + c + d, base=(a + c) / (a + b + c + d),
               POD=pod, POFD=pofd, FAR=far, FBI=fbi, CSI=csi)
    # EDI/SEDI are undefined at H or F of exactly 0 or 1; clip rather than crash.
    H = min(max(pod, 1e-6), 1 - 1e-6)
    F = min(max(pofd, 1e-6), 1 - 1e-6)
    lH, lF = np.log(H), np.log(F)
    out["EDI"] = (lF - lH) / (lF + lH)
    out["SEDI"] = ((lF - lH - np.log(1 - F) + np.log(1 - H)) /
                   (lF + lH + np.log(1 - F) + np.log(1 - H)))
    return out


def regional(ev, reg):
    """(ninit, npt) point events -> (ninit, nreg) regional EHDE under EXP1."""
    ids = np.unique(reg)
    out = np.zeros((ev.shape[0], len(ids)), dtype=bool)
    for k, r in enumerate(ids):
        m = reg == r
        out[:, k] = ev[:, m].mean(1) > REG_FRAC
    return out


def curve(score, obs, mask, reg=None):
    """Metrics at every delta, over the rows selected by `mask`."""
    s, o = score[mask], obs[mask]
    ok = np.isfinite(s)
    rows = []
    for dl in DELTAS:
        fc = np.zeros_like(o, dtype=bool)
        fc[ok] = s[ok] >= dl
        if reg is None:
            rows.append(contingency(fc, o))
        else:
            rows.append(contingency(regional(fc, reg), regional(o, reg)))
    return rows


BOX = dict(POD=0.60, POFD=0.20, FAR=0.20, CSI=0.40, FBI=(0.7, 1.3))


def in_box(r):
    return (r["POD"] > BOX["POD"] and r["POFD"] < BOX["POFD"]
            and r["FAR"] < BOX["FAR"] and r["CSI"] > BOX["CSI"]
            and BOX["FBI"][0] < r["FBI"] < BOX["FBI"][1])


def pick(rows, rule):
    """Choose delta on the training curve. Returns an index into DELTAS."""
    if rule == "csi":
        v = np.array([r["CSI"] if np.isfinite(r["CSI"]) else -1 for r in rows])
        return int(np.argmax(v))
    if rule == "fbi":
        v = np.array([abs(r["FBI"] - 1) if np.isfinite(r["FBI"]) else 9e9
                      for r in rows])
        return int(np.argmin(v))
    if rule == "box":
        okk = [i for i, r in enumerate(rows) if in_box(r)]
        if not okk:
            return -1
        return int(max(okk, key=lambda i: rows[i]["CSI"]))
    if rule == "far":
        # sharpest (lowest delta => highest POD) threshold still under FAR 0.20
        okk = [i for i, r in enumerate(rows)
               if np.isfinite(r["FAR"]) and r["FAR"] < 0.20 and r["POD"] > 0]
        return int(min(okk)) if okk else int(len(DELTAS) - 1)
    raise ValueError(rule)


def fmt(tag, dl, r):
    return ("%-6s d=%+5.2f  POD %.3f  POFD %.3f  FAR %.3f  FBI %.3f  "
            "CSI %.3f  EDI %.3f  SEDI %.3f  base %.4f"
            % (tag, dl, r["POD"], r["POFD"], r["FAR"], r["FBI"], r["CSI"],
               r["EDI"], r["SEDI"], r["base"]))


_SCALE = ["point"]


def out_base():
    return RES if _SCALE[0] == "point" else RES.replace(".json", "_region.json")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--scale", choices=["point", "region"], default="point")
    a = ap.parse_args()
    _SCALE[0] = a.scale

    res = {}
    for w, wname in enumerate(C.WEEK_NAMES):
        f = "%s.w%d.fields.npz" % (FLD, w)
        z = np.load(f)
        d1, thr, E, yrs, reg = z["d1"], z["thr"], z["E"].astype(bool), z["yrs"], z["reg"]
        score = d1 - thr
        rg = reg if a.scale == "region" else None

        train = np.isin(yrs, C.TRAIN_YEARS) & np.isfinite(score).any(1)
        valid = np.isin(yrs, C.VALID_YEARS) & np.isfinite(score).any(1)

        ctr = curve(score, E, train, rg)
        cva = curve(score, E, valid, rg)

        row = {"n_train_init": int(train.sum()), "n_valid_init": int(valid.sum())}
        print("\n=== %s (%s, train %d init / valid %d init) ==="
              % (wname, a.scale, train.sum(), valid.sum()))
        for rule in ("csi", "fbi", "far", "box"):
            i = pick(ctr, rule)
            if i < 0:
                row[rule] = None
                print("%-6s KHONG co nguong nao dat hop chi tieu tren nam HUAN LUYEN"
                      % (rule + " tr"))
                continue
            row[rule] = {"delta": float(DELTAS[i]),
                         "train": {k: float(v) for k, v in ctr[i].items()},
                         "valid": {k: float(v) for k, v in cva[i].items()}}
            print(fmt(rule + " tr", DELTAS[i], ctr[i]))
            print(fmt(rule + " VA", DELTAS[i], cva[i]))

        # is the target box reachable AT ALL on the validation years, at any delta?
        hit = [i for i, r in enumerate(cva) if in_box(r)]
        if hit:
            print("  dai delta dat hop tren nam KIEM DINH: %+.2f .. %+.2f"
                  % (DELTAS[hit[0]], DELTAS[hit[-1]]))
        np.savez("%s.w%d.curve.npz" % (os.path.splitext(out_base())[0], w),
                 delta=DELTAS,
                 **{"tr_" + k: np.array([r[k] for r in ctr]) for k in ctr[0]},
                 **{"va_" + k: np.array([r[k] for r in cva]) for k in cva[0]})
        row["box_reachable"] = bool(hit)
        row["best_valid_CSI"] = float(max(r["CSI"] for r in cva))
        print("  hop chi tieu dat duoc tren nam kiem dinh: %s   CSI ton tai lon nhat %.3f"
              % ("CO" if hit else "KHONG", row["best_valid_CSI"]))
        res[wname] = row

    os.makedirs(os.path.dirname(RES), exist_ok=True)
    out = out_base()
    json.dump(res, open(out, "w"), indent=1)
    print("\nsaved", out)


if __name__ == "__main__":
    main()