#!/usr/bin/env python3
"""Does upper-air information, or a nonlinear model, close the gap to AUC_req?

Four configurations, same inits, same split, same scoring -- so each line differs
from the one above it by exactly one thing:

    L_flux    logistic regression, 16 flux variables            (the floor)
    L_both    logistic regression, flux + 27 upper-air features (what pgb buys)
    G_flux    LightGBM, flux only                               (what nonlinearity buys)
    G_both    LightGBM, flux + upper-air                        (both)

Everything is z-scored against the model's own (week, day of year, point)
climatology from training years. Inits without a pgb file are dropped from ALL
four so the comparison is not confounded by sample size.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import warnings

import numpy as np

warnings.filterwarnings("ignore")
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import cfs_common as C
import build_obs as B
from weekly_skill import auc, climz
from signal_ceiling import far_to_pofd, auc_needed

RES = os.environ.get("EXP_RES",
                     os.path.join(C.OUT, "AI_Results", "exp_features.json"))
PGBF = os.environ.get("PGB_WEEKLY",
                      os.path.join(C.OUT, "Cube", "pgb_weekly.npz"))


def zblock(X, doy, yrs, train, w):
    """z-score every variable of week w -> (ninit, npt, nvar)."""
    return np.stack([climz(X[:, [w], v], doy[:, [w]], yrs, train)[:, 0]
                     for v in range(X.shape[2])], -1)


def part_path(w):
    return "%s.w%d.json" % (RES, w)


def merge():
    """Combine the per-week part files written by the array tasks."""
    res = {}
    for w, wname in enumerate(C.WEEK_NAMES):
        f = part_path(w)
        if not os.path.exists(f):
            print("THIEU %s" % f)
            continue
        res[wname] = json.load(open(f))
    json.dump(res, open(RES, "w"), indent=1)
    print("%-4s %8s %8s %8s %8s %9s" %
          ("week", "L_flux", "L_both", "G_flux", "G_both", "AUC_req"))
    for wname, r in res.items():
        print("%-4s %8.4f %8.4f %8.4f %8.4f %9.4f"
              % (wname, r["L_flux"], r["L_both"], r["G_flux"], r["G_both"],
                 r["AUC_req"]))
    print("\nwhat each step bought (AUC)")
    print("%-4s %10s %10s %10s" % ("week", "upper-air", "nonlinear", "both"))
    for wname, r in res.items():
        print("%-4s %+10.4f %+10.4f %+10.4f"
              % (wname, r["L_both"] - r["L_flux"], r["G_flux"] - r["L_flux"],
                 r["G_both"] - r["L_flux"]))
    print("\nsaved", RES)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--week", type=int, default=-1,
                    help="compute only this week (0-5) into its own part file")
    ap.add_argument("--merge", action="store_true")
    a = ap.parse_args()
    if a.merge:
        return merge()
    ncore = int(os.environ.get("SLURM_CPUS_PER_TASK", 8))

    fx = np.load(os.path.join(C.OUT, "Cube", "weekly_feats.npz"))["X"]
    wk = np.load(os.path.join(C.OUT, "Cube", "weekly.npz"))
    yrs, E, doy = wk["yrs"], wk["E"], wk["doy"]
    pg = np.load(PGBF)
    px, have = pg["X"], pg["have"]

    train = np.isin(yrs, C.TRAIN_YEARS) & have
    valid = np.isin(yrs, C.VALID_YEARS) & have
    print("inits with pgb: train %d, valid %d (of %d)"
          % (int(train.sum()), int(valid.sum()), len(yrs)), flush=True)

    from sklearn.linear_model import LogisticRegression
    import lightgbm as lgb

    res = {}
    print("\n%-4s %8s %8s %8s %8s %9s" %
          ("week", "L_flux", "L_both", "G_flux", "G_both", "AUC_req"))
    for w, name in enumerate(C.WEEK_NAMES):
        if a.week >= 0 and w != a.week:
            continue
        Zf = zblock(fx, doy, yrs, train, w)
        Zp = zblock(px, doy, yrs, train, w)
        d = doy[:, w].astype(float)
        h = np.stack([np.cos(2 * np.pi * d / 365), np.sin(2 * np.pi * d / 365),
                      np.cos(4 * np.pi * d / 365), np.sin(4 * np.pi * d / 365)], -1)
        H = np.repeat(h[:, None, :], Zf.shape[1], axis=1)
        y = E[:, w]

        def mats(Z, mask):
            f = np.concatenate([Z, H], -1)[mask].reshape(-1, Z.shape[-1] + 4)
            t = y[mask].ravel()
            ok = np.isfinite(f).all(1)
            return f[ok], t[ok]

        row = {}
        for tag, Z in (("flux", Zf), ("both", np.concatenate([Zf, Zp], -1))):
            ftr, ytr = mats(Z, train)
            fva, yva = mats(Z, valid)
            mu, sd = ftr.mean(0), ftr.std(0) + 1e-9
            lr = LogisticRegression(max_iter=3000, C=1.0)
            lr.fit((ftr - mu) / sd, ytr)
            row["L_" + tag] = auc(lr.predict_proba((fva - mu) / sd)[:, 1], yva)
            gb = lgb.LGBMClassifier(n_estimators=400, learning_rate=0.05,
                                    num_leaves=63, min_child_samples=200,
                                    subsample=0.8, colsample_bytree=0.8,
                                    verbose=-1, n_jobs=ncore)
            gb.fit(ftr, ytr)
            row["G_" + tag] = auc(gb.predict_proba(fva)[:, 1], yva)
            row["base"] = float(yva.mean())
        row["AUC_req"] = auc_needed(0.60, min(far_to_pofd(0.60, 0.20, row["base"]), .999))
        res[name] = row
        print("%-4s %8.4f %8.4f %8.4f %8.4f %9.4f"
              % (name, row["L_flux"], row["L_both"], row["G_flux"], row["G_both"],
                 row["AUC_req"]), flush=True)

    os.makedirs(os.path.dirname(RES), exist_ok=True)
    if a.week >= 0:
        out = part_path(a.week)
        json.dump(res[C.WEEK_NAMES[a.week]], open(out, "w"), indent=1)
    else:
        out = RES
        json.dump(res, open(out, "w"), indent=1)
    print("\nsaved", out)

    if a.week < 0:
        print("\nwhat each step bought (AUC)")
        print("%-4s %10s %10s %10s" % ("week", "upper-air", "nonlinear", "both"))
        for name in C.WEEK_NAMES:
            r = res[name]
            print("%-4s %+10.4f %+10.4f %+10.4f"
                  % (name, r["L_both"] - r["L_flux"], r["G_flux"] - r["L_flux"],
                     r["G_both"] - r["L_flux"]))


if __name__ == "__main__":
    main()