"""Cheap ablation: does adding DEM elevation help the stage-2 EHDE classifier?

Method (matches the project's own '3 points x 7 regions' test protocol):
  * pick 21 gridpoints (3 per climate region),
  * build the operational F3 weekly-aggregated design matrix,
  * train nationally-pooled stage-2 models on TRAIN_YEARS, score on VALID_YEARS,
  * compare WITHOUT vs WITH two static elevation columns (elev_mean, elev_std),
  * report AUC and max-SEDI (base-rate robust) per lead week, plus the elevation
    coefficient (LogReg) and gain rank (LGBM) so we can see if the model even uses it.

Elevation is a per-point static field; its first-order signal is already carried by
thr_obs/thr_mod (local Tx90p) and the per-gridpoint z-score, so the honest question
is whether the RESIDUAL (35C-cap interaction, terrain ruggedness) adds skill.
"""
import os
import sys

import numpy as np
from netCDF4 import Dataset
from sklearn.linear_model import LogisticRegression
import lightgbm as lgb

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import ehd_common as C
from ai_data import Store
from phase1_arch import build_window
from ehd_metrics import auc_score, best_threshold, score_binary

DEM = os.path.join(C.OUT, "DEM", "dem_features.nc")


def pick_points(store, per_region=3):
    reg = np.asarray(store.st["region"]).astype(int)
    pts = []
    for r in range(1, 8):
        idx = np.nonzero(reg == r)[0]
        if idx.size == 0:
            continue
        sel = idx[np.linspace(0, idx.size - 1, per_region).round().astype(int)]
        pts.extend(sel.tolist())
    return np.array(sorted(set(pts)), dtype=int)


def load_elev():
    with Dataset(DEM) as ds:
        return (np.asarray(ds["p_elev_mean"][:], "f8"),
                np.asarray(ds["p_elev_std"][:], "f8"))


def prep(Xtr, Xte, etr, ete, use_elev):
    """Impute (train median) + standardize (train stats); optionally append elev."""
    if use_elev:
        Xtr = np.column_stack([Xtr, etr])
        Xte = np.column_stack([Xte, ete])
    med = np.nanmedian(Xtr, axis=0)
    Xtr = np.where(np.isnan(Xtr), med, Xtr)
    Xte = np.where(np.isnan(Xte), med, Xte)
    mu = Xtr.mean(0); sd = Xtr.std(0); sd[sd == 0] = 1.0
    return (Xtr - mu) / sd, (Xte - mu) / sd


def fit_eval(Xtr, ytr, Xte, yte, kind):
    if kind == "logreg":
        m = LogisticRegression(C=1.0, penalty="l2", max_iter=2000, class_weight=None)
        m.fit(Xtr, ytr)
        p = m.predict_proba(Xte)[:, 1]
        return p, m
    else:
        m = lgb.LGBMClassifier(n_estimators=300, learning_rate=0.05, num_leaves=31,
                               min_child_samples=50, subsample=0.8,
                               colsample_bytree=0.8, verbose=-1, n_jobs=8)
        m.fit(Xtr, ytr)
        p = m.predict_proba(Xte)[:, 1]
        return p, m


def main():
    store = Store()
    pts = pick_points(store)
    reg = np.asarray(store.st["region"]).astype(int)
    print("test points (%d):" % pts.size,
          {r: int((reg[pts] == r).sum()) for r in range(1, 8)})
    elev_mean, elev_std = load_elev()

    print("\n%-3s %6s | %-28s | %-28s" % (
        "W", "base", "LogReg  AUC(base/elev) dAUC  SEDI(b/e)",
        "LGBM AUC(base/elev) dAUC  elevRank"))
    for wname, lts in C.WEEKS:
        Atr, _, Ntr, _, Ptr, names = build_window(store, pts, C.TRAIN_YEARS, "F3", lts)
        Ate, _, Nte, _, Pte, _ = build_window(store, pts, C.VALID_YEARS, "F3", lts)
        ytr = (Ntr >= 1).astype(int)
        yte = (Nte >= 1).astype(int)
        etr = np.column_stack([elev_mean[Ptr], elev_std[Ptr]])
        ete = np.column_stack([elev_mean[Pte], elev_std[Pte]])
        base = yte.mean()

        out = {}
        for kind in ("logreg", "lgbm"):
            res = {}
            for tag, ue in (("base", False), ("elev", True)):
                Xtr, Xte = prep(Atr, Ate, etr, ete, ue)
                p, m = fit_eval(Xtr, ytr, Xte, yte, kind)
                auc = auc_score(p, yte)
                t, sedi = best_threshold(p, yte, metric="SEDI")
                res[tag] = (auc, sedi, m, Xtr.shape[1])
            out[kind] = res

        la, lsb = out["logreg"]["base"][0], out["logreg"]["base"][1]
        le, lse = out["logreg"]["elev"][0], out["logreg"]["elev"][1]
        ga, _ = out["lgbm"]["base"][0], out["lgbm"]["base"][1]
        ge, _ = out["lgbm"]["elev"][0], out["lgbm"]["elev"][1]
        # elevation coefficient (LogReg, standardized) and LGBM gain rank
        mL = out["logreg"]["elev"][2]
        coef = mL.coef_.ravel()
        coef_elev_mean, coef_elev_std = coef[-2], coef[-1]
        mG = out["lgbm"]["elev"][2]
        imp = mG.booster_.feature_importance(importance_type="gain")
        order = np.argsort(imp)[::-1]
        nf = imp.size
        rank_mean = int(np.nonzero(order == nf - 2)[0][0]) + 1
        rank_std = int(np.nonzero(order == nf - 1)[0][0]) + 1

        print("%-3s %6.3f | AUC %.3f/%.3f d%+.3f  SEDI %.3f/%.3f d%+.3f | "
              "AUC %.3f/%.3f d%+.3f | coef(m/s)=%+.2f/%+.2f rank(m/s)=%d/%d of %d"
              % (wname, base, la, le, le - la, lsb, lse, lse - lsb,
                 ga, ge, ge - ga, coef_elev_mean, coef_elev_std,
                 rank_mean, rank_std, nf))


if __name__ == "__main__":
    main()