#!/usr/bin/env python3
"""Stage 2 section 6: AI models for EHD, from the simplest option upwards.

Model families, in the order the docx suggests trying them:
  ec_raw    the raw EC forecast, carried through unchanged as the benchmark
  logreg    logistic regression on the standardised features
  lgbm_clf  gradient-boosted trees predicting P(EHD)
  lgbm_reg  gradient-boosted trees predicting the observed severity, the event
            then being read off the predicted severity

Everything is fitted on the sixteen training initialisation years and scored on
the four held-out years.  Decision thresholds are never tuned on the held-out
years: they come from out-of-fold predictions inside the training set, chosen
per forecast week so the frequency bias can be controlled at every lead time.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time

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 ai_data as D

RES_DIR = os.path.join(C.OUT, "AI_Results")
NFOLD = 4


# --------------------------------------------------------------------- models
def make_model(kind, pos_weight):
    if kind == "logreg":
        from sklearn.linear_model import LogisticRegression
        from sklearn.pipeline import make_pipeline
        from sklearn.preprocessing import StandardScaler
        return make_pipeline(
            StandardScaler(),
            LogisticRegression(max_iter=400, C=1.0, class_weight="balanced",
                               solver="lbfgs"))
    if kind == "lgbm_clf":
        import lightgbm as lgb
        return lgb.LGBMClassifier(
            n_estimators=600, learning_rate=0.05, num_leaves=63,
            min_child_samples=200, subsample=0.8, subsample_freq=1,
            colsample_bytree=0.8, reg_lambda=1.0,
            scale_pos_weight=pos_weight, n_jobs=-1, verbose=-1)
    if kind == "mlp":
        from nn_models import TorchMLP
        return TorchMLP(pos_weight=pos_weight, epochs=12, verbose=False)
    if kind == "lgbm_reg":
        import lightgbm as lgb
        return lgb.LGBMRegressor(
            n_estimators=600, learning_rate=0.05, num_leaves=63,
            min_child_samples=200, subsample=0.8, subsample_freq=1,
            colsample_bytree=0.8, reg_lambda=1.0, n_jobs=-1, verbose=-1)
    raise ValueError(kind)


def fit_predict(kind, Xtr, ytr, Xte, pos_weight):
    """Fit one model and return a continuous score on Xte (higher = more likely)."""
    if kind == "lgbm_reg":
        m = make_model(kind, pos_weight)
        m.fit(Xtr, ytr["sev"])
        return m.predict(Xte), m
    m = make_model(kind, pos_weight)
    m.fit(Xtr, ytr["evt"])
    return m.predict_proba(Xte)[:, 1], m


# --------------------------------------------------------------- thresholding
def oof_scores(kind, X, y, info, pos_weight, nfold=NFOLD, seed=0):
    """Out-of-fold scores over the training years, split by whole years."""
    years = np.unique(info["year"])
    rng = np.random.default_rng(seed)
    shuffled = rng.permutation(years)
    folds = np.array_split(shuffled, nfold)
    out = np.empty(X.shape[0], dtype="f8")
    for f in folds:
        te = np.isin(info["year"], f)
        tr = ~te
        s, _ = fit_predict(kind, X[tr], {k: v[tr] for k, v in y.items()},
                           X[te], pos_weight)
        out[te] = s
    return out


def weekly_thresholds(score, y_evt, week, metric="CSI"):
    """One decision threshold per forecast week, chosen on the training data."""
    th = {}
    for w in range(len(C.WEEKS)):
        m = week == w
        if m.sum() == 0 or y_evt[m].sum() == 0:
            th[w] = np.inf
            continue
        th[w], _ = M.best_threshold(score[m], y_evt[m].astype(bool), metric)
    return th


# ------------------------------------------- thresholds on the real population
# Training thins the non-events (keeping `neg_keep` of them) so the trees see a
# workable class balance, but that leaves the out-of-fold score distribution with
# a base rate several times the real one.  Choosing a threshold directly on that
# distribution sets it far too low and over-forecasts by the same factor: the
# first national run reached FBI 3.3.  Every threshold below is therefore chosen
# on statistics reweighted back to the original balance.

def population_weights(evt, neg_keep):
    """Undo the negative thinning: each surviving non-event stands for 1/neg_keep."""
    return np.where(evt > 0.5, 1.0, 1.0 / neg_keep)


def weighted_quantile(x, w, q):
    """Quantile q of x under sample weights w."""
    order = np.argsort(x)
    x, w = x[order], w[order]
    cw = np.cumsum(w)
    return float(np.interp(q * cw[-1], cw, x))


def weighted_scores(pred, obs, w):
    """Contingency scores with sample weights."""
    a = float(w[pred & obs].sum()); b = float(w[pred & ~obs].sum())
    c = float(w[~pred & obs].sum()); d = float(w[~pred & ~obs].sum())
    return M.scores(a, b, c, d)


def bias_one_thresholds(score, obs, week, w):
    """Per week, the threshold whose forecast frequency equals the observed one."""
    th = {}
    for k in range(len(C.WEEKS)):
        m = week == k
        if m.sum() == 0:
            th[k] = np.inf
            continue
        rate = float((w[m] * obs[m]).sum() / w[m].sum())
        th[k] = (np.inf if rate <= 0
                 else weighted_quantile(score[m], w[m], 1.0 - rate))
    return th


def csi_thresholds(score, obs, week, w, ngrid=300):
    """Per week, the threshold maximising the weighted CSI."""
    th = {}
    ob = obs > 0.5
    for k in range(len(C.WEEKS)):
        m = week == k
        if m.sum() == 0 or ob[m].sum() == 0:
            th[k] = np.inf
            continue
        s, o, ww = score[m], ob[m], w[m]
        grid = np.unique(np.quantile(s, np.linspace(0.5, 0.9999, ngrid)))
        best, bt = -np.inf, float(np.median(s))
        for t in grid:
            v = weighted_scores(s >= t, o, ww)["CSI"]
            if np.isfinite(v) and v > best:
                best, bt = v, float(t)
        th[k] = bt
    return th


def platt_calibrate(score_oof, evt_oof, w_oof, score_new):
    """Weighted Platt scaling fitted on out-of-fold data only.

    The models are trained with balanced class weights on thinned data, so their
    raw probabilities are far too high; a one-dimensional logistic fit on the
    out-of-fold scores maps them back onto the real event frequency, which is what
    a reliability diagram and the Brier score need.
    """
    from sklearn.linear_model import LogisticRegression
    x = np.asarray(score_oof, dtype="f8")
    xn = np.asarray(score_new, dtype="f8")
    if x.min() >= 0.0 and x.max() <= 1.0:
        eps = 1e-6
        def _logit(v):
            v = np.clip(v, eps, 1 - eps)
            return np.log(v / (1 - v))
        x, xn = _logit(x), _logit(xn)
    lr = LogisticRegression(max_iter=1000)
    lr.fit(x[:, None], (evt_oof > 0.5).astype(int), sample_weight=w_oof)
    return lr.predict_proba(xn[:, None])[:, 1]


def stratified_thresholds(score, obs, week, strat, w, rule="fbi1", min_events=50):
    """One threshold per (week, stratum) instead of one per week.

    A single national threshold has to serve climatic zones whose event rates
    differ by an order of magnitude, so the frequency bias can only be right on
    average.  Where a stratum holds too few events to read a threshold off
    reliably it falls back to the week-wide one rather than to a noisy fit.
    """
    pick = bias_one_thresholds if rule == "fbi1" else csi_thresholds
    base = pick(score, obs, week, w)
    th, ob = {}, obs > 0.5
    for sv in np.unique(strat):
        m = strat == sv
        sub = pick(score[m], obs[m], week[m], w[m])
        for k in range(len(C.WEEKS)):
            enough = int(ob[m & (week == k)].sum()) >= min_events
            th[(k, int(sv))] = sub[k] if enough else base[k]
    return th


def apply_stratified(score, week, strat, th):
    """Vectorised version of apply_thresholds for (week, stratum) thresholds."""
    ns = max(int(np.max(strat)) + 1, 1)
    tab = np.full(len(C.WEEKS) * ns, np.inf)
    for (k, sv), t in th.items():
        if sv < ns:
            tab[k * ns + sv] = t
    return score >= tab[week.astype("i8") * ns + strat.astype("i8")]


def apply_thresholds(score, week, th):
    pred = np.zeros(score.size, dtype=bool)
    for w, t in th.items():
        m = week == w
        pred[m] = score[m] >= t
    return pred


# ---------------------------------------------------------------- evaluation
def evaluate(score, pred, y_evt, info, subset=None):
    """Per-week scores, plus the all-week total, on an optional row subset."""
    obs = y_evt.astype(bool)
    rows = {}
    sel_all = np.ones(obs.size, bool) if subset is None else subset
    for w in range(len(C.WEEKS)):
        m = sel_all & (info["week"] == w)
        if m.sum() == 0:
            continue
        s = M.score_binary(pred[m], obs[m])
        s["AUC"] = M.auc_score(score[m], obs[m])
        rows[C.WEEK_NAMES[w]] = s
    m = sel_all
    s = M.score_binary(pred[m], obs[m])
    s["AUC"] = M.auc_score(score[m], obs[m])
    rows["ALL"] = s
    return rows


def weekly_event(score, pred, y_evt, info):
    """Collapse each (initialisation, week, point) to a single event/no-event."""
    key = (info["init"].astype("i8") * 16 + info["week"]) * 100000 + info["point"]
    uniq, inv = np.unique(key, return_inverse=True)
    n = uniq.size
    o = np.zeros(n, bool); p = np.zeros(n, bool)
    s = np.full(n, -np.inf); wk = np.zeros(n, "i4")
    np.logical_or.at(o, inv, y_evt.astype(bool))
    np.logical_or.at(p, inv, pred)
    np.maximum.at(s, inv, score)
    wk[inv] = info["week"]
    return s, p, o, {"week": wk}


def fmt_table(title, rows):
    keys = ["base_rate", "POD", "POFD", "FAR", "FBI", "CSI", "ETS", "AUC"]
    out = ["  " + title,
           "    %-6s" % "" + "".join("%9s" % k for k in keys)]
    for name, s in rows.items():
        out.append("    %-6s" % name + "".join("%9.3f" % s.get(k, np.nan) for k in keys))
    return "\n".join(out)


# --------------------------------------------------------------------- driver
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--per-region", type=int, default=3)
    ap.add_argument("--models", nargs="*",
                    default=["ec_raw", "logreg", "lgbm_clf", "lgbm_reg"],
                    choices=["ec_raw", "logreg", "lgbm_clf", "lgbm_reg", "mlp"])
    ap.add_argument("--fsets", nargs="*", default=["F0", "F1", "F2", "F3"])
    ap.add_argument("--out", default=os.path.join(RES_DIR, "model_selection.json"))
    args = ap.parse_args()

    store = D.Store()
    pts = D.sample_points(store.grid, args.per_region)
    print("[points] %d test grid points: %s" % (pts.size, pts.tolist()), flush=True)
    os.makedirs(RES_DIR, exist_ok=True)

    results = {}
    cache = {}
    for fset in args.fsets:
        t0 = time.time()
        Xtr, ytr, itr = store.build(pts, C.TRAIN_YEARS, fset)
        Xte, yte, ite = store.build(pts, C.VALID_YEARS, fset)
        names = store.feature_names(fset)
        ok_tr = np.isfinite(Xtr).all(axis=1) & np.isfinite(ytr["evt"])
        ok_te = np.isfinite(Xte).all(axis=1) & np.isfinite(yte["evt"])
        Xtr, Xte = Xtr[ok_tr], Xte[ok_te]
        ytr = {k: v[ok_tr] for k, v in ytr.items()}
        yte = {k: v[ok_te] for k, v in yte.items()}
        itr = {k: v[ok_tr] for k, v in itr.items()}
        ite = {k: v[ok_te] for k, v in ite.items()}
        pos = float((1 - ytr["evt"].mean()) / max(ytr["evt"].mean(), 1e-9))
        print("\n=== %s: %d features | train %d rows | valid %d rows | base rate %.4f "
              "| pos_weight %.1f | %.1fs" %
              (fset, Xtr.shape[1], Xtr.shape[0], Xte.shape[0],
               ytr["evt"].mean(), pos, time.time() - t0), flush=True)
        cache[fset] = (Xtr, ytr, itr, Xte, yte, ite, names, pos)

        for kind in args.models:
            if kind == "ec_raw" and fset != args.fsets[0]:
                continue
            t1 = time.time()
            if kind == "ec_raw":
                j = names.index("sv_mod")
                s_tr, s_te = Xtr[:, j].astype("f8"), Xte[:, j].astype("f8")
                model = None
            else:
                s_tr = oof_scores(kind, Xtr, ytr, itr, pos)
                s_te, model = fit_predict(kind, Xtr, ytr, Xte, pos)

            th = weekly_thresholds(s_tr, ytr["evt"], itr["week"])
            pred = apply_thresholds(s_te, ite["week"], th)

            tag = "%s/%s" % (kind, fset)
            daily_all = evaluate(s_te, pred, yte["evt"], ite)
            daily_cap = evaluate(s_te, pred, yte["evt"], ite, ite["capable"])
            ws, wp, wo, wi = weekly_event(s_te, pred, yte["evt"], ite)
            week_all = evaluate(ws, wp, wo, wi)

            results[tag] = {"daily_all": daily_all, "daily_capable": daily_cap,
                            "weekly_all": week_all,
                            "thresholds": {str(k): float(v) for k, v in th.items()},
                            "seconds": time.time() - t1}
            if model is not None and hasattr(model, "feature_importances_"):
                imp = np.asarray(model.feature_importances_, dtype="f8")
                imp = imp / max(imp.sum(), 1e-9)
                order = np.argsort(imp)[::-1][:15]
                results[tag]["top_features"] = [(names[i], float(imp[i])) for i in order]

            print("\n[%s]  (%.0fs)" % (tag, time.time() - t1), flush=True)
            print(fmt_table("daily, all rows", daily_all), flush=True)
            print(fmt_table("daily, heat-capable cells only", daily_cap), flush=True)
            print(fmt_table("weekly event (>=1 EHD day in the week)", week_all), flush=True)

        del cache[fset]

    with open(args.out, "w") as fh:
        json.dump(results, fh, indent=1, default=float)
    print("\n[done] -> %s" % args.out, flush=True)


if __name__ == "__main__":
    main()