#!/usr/bin/env python3
"""Weekly multivariate features at the 541 observation points, and what a plain
pooled logistic regression can already do with them.

This is the honest floor-to-ceiling measurement. The single-predictor baseline
(weekly max Tx alone) says what the raw model carries; this says how much of the
gap to the required AUC multivariate information closes, BEFORE any downscaling
or model-family search is attempted. If a linear model on 16 z-scored variables
lands far short at W3-W6, no amount of threshold tuning will rescue those windows.

Features per (init, week, point): each of the 16 flux variables aggregated over
the week's lead days (max for tmax, min for tmin, mean otherwise), z-scored
against the MODEL's own (week, day of year, point) climatology from training
years, plus two annual harmonics. Observations never enter the features.
"""

from __future__ import annotations

import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import cfs_common as C
import build_obs as B
import build_daily as D
from weekly_skill import auc, climz

FEAT = os.path.join(C.OUT, "Cube", "weekly_feats.npz")


def build():
    n, nlt, nv, ny, nx = D.shapes()
    inits = C.both_inits()
    daily = np.memmap(D.DAILY, dtype="float32", mode="r", shape=(n, nlt, nv, ny, nx))
    z = np.load(B.TX_NPZ)
    olat, olon = z["lat"][z["jj"]], z["lon"][z["ii"]]
    g = C.flux_grid()
    pj = np.abs(olat[:, None] - g["wlat"][None, :]).argmin(1)
    pi = np.abs(olon[:, None] - g["wlon"][None, :]).argmin(1)
    npt, nw = len(pj), len(C.WEEKS)

    X = np.full((n, nw, nv, npt), np.nan, dtype="float32")
    for i in range(n):
        cube = daily[i][:, :, pj, pi]                    # (46, 16, npt)
        for w, (_, idx) in enumerate(C.WEEKS):
            sub = cube[idx]
            with np.errstate(invalid="ignore"):
                X[i, w] = np.nanmean(sub, axis=0)
                X[i, w, C.FLUX_NAMES.index("tmax")] = np.nanmax(
                    sub[:, C.FLUX_NAMES.index("tmax")], axis=0)
                X[i, w, C.FLUX_NAMES.index("tmin")] = np.nanmin(
                    sub[:, C.FLUX_NAMES.index("tmin")], axis=0)
        if (i + 1) % 500 == 0:
            print("  %d/%d" % (i + 1, n), flush=True)
    np.savez(FEAT, X=X, inits=np.array(inits))
    return X


def main():
    if not os.path.exists(FEAT):
        build()
    X = np.load(FEAT)["X"]
    wk = np.load(os.path.join(C.OUT, "Cube", "weekly.npz"))
    yrs, E, doy = wk["yrs"], wk["E"], wk["doy"]
    train = np.isin(yrs, C.TRAIN_YEARS)
    valid = np.isin(yrs, C.VALID_YEARS)

    from sklearn.linear_model import LogisticRegression
    from signal_ceiling import far_to_pofd, auc_needed

    print("\nPOOLED LOGISTIC REGRESSION ON 16 z-SCORED VARIABLES + 2 HARMONICS")
    print("trained on %d inits of %s, scored on %d inits of %s"
          % (int(train.sum()), C.TRAIN_YEARS, int(valid.sum()), C.VALID_YEARS))
    print("%-4s %8s %8s %8s %9s %8s"
          % ("week", "AUC_1var", "AUC_16var", "AUC_req", "gap", "base"))
    for w, name in enumerate(C.WEEK_NAMES):
        cols = []
        for v in range(X.shape[2]):
            cols.append(climz(X[:, [w], v], doy[:, [w]], yrs, train)[:, 0])
        Z = np.stack(cols, -1)                            # (n, npt, nvar)
        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, :], Z.shape[1], axis=1)
        F = np.concatenate([Z, H], -1)
        y = E[:, w]

        def flat(mask):
            f = F[mask].reshape(-1, F.shape[-1])
            t = y[mask].ravel()
            ok = np.isfinite(f).all(1)
            return f[ok], t[ok]

        ftr, ytr = flat(train)
        fva, yva = flat(valid)
        mu, sd = ftr.mean(0), ftr.std(0) + 1e-9
        clf = LogisticRegression(max_iter=2000, C=1.0)
        clf.fit((ftr - mu) / sd, ytr)
        p = clf.predict_proba((fva - mu) / sd)[:, 1]
        a16 = auc(p, yva)
        a1 = auc(fva[:, C.FLUX_NAMES.index("tmax")], yva)
        br = float(yva.mean())
        req = auc_needed(0.60, min(far_to_pofd(0.60, 0.20, br), 0.999))
        print("%-4s %8.4f %8.4f %8.4f %9.4f %8.4f"
              % (name, a1, a16, req, a16 - req, br))


if __name__ == "__main__":
    main()