"""p04_screen.py: screening computations used only to score ideas (not findings).

S1. How much do random path-level splits overstate out-of-fold AUC compared with splits that keep
    each run whole? One simple model (standardised logistic regression) on four feature sets,
    per laboratory machine, label 1 against label 0. Spread over repeated random splits is reported;
    it describes split-to-split variation, not sampling uncertainty.
S2. How far does one clipping statistic (share of Sound0 samples at full scale) go on its own?
S3. How far does tool identity go on its own (a one-bit baseline: tool A or its defective twin A*, against every other tool)?
None of these analyses was in the pre-specified list; they are screening evidence and labelled as such.
"""
import json, sys
import numpy as np, pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline
from sklearn.metrics import roc_auc_score
sys.path.insert(0, "/home/claude/mill/code")
from vals import Values

OUT = "/home/claude/mill/outputs"
V = Values(script="p04_screen.py")
T = pd.read_parquet(f"{OUT}/segments.parquet")
REPEATS = 20
SEED = 20260922
V.set("s1_folds", 5, "folds", {"population": "screening setting", "filter": "p04_screen.py", "statistic": "folds per split"})
V.set("s1_repeats", REPEATS, "repeats", {"population": "screening setting", "filter": "p04_screen.py", "statistic": "repeated random splits per scheme"})

SETS = {
    "highrate": ["acc0_rms", "acc1_rms", "s0_rms", "s1_rms", "s2_rms"],
    "mtc": ["mtc_smarms", "mtc_smapeak", "mtc_smcrest", "mtc_smvrms", "mtc_S_Axis_Load", "mtc_watt"],
    "s2only": ["s2_rms"],
    "clip": ["s0_fullscale"],
}


def oof_auc(d, cols, groups, k, rng):
    X = d[cols].copy()
    for c in cols:
        if c.endswith("_rms") or c.startswith("mtc_sm") or c == "mtc_watt":
            X[c] = np.log(np.clip(X[c].astype(float), 1e-9, None))
    y = d.label.values
    pred = np.full(len(d), np.nan)
    if groups is None:
        # random path-level folds, stratified by label
        idx = np.arange(len(d)); folds = np.zeros(len(d), int)
        for lab in (0, 1):
            ii = idx[y == lab]; rng.shuffle(ii)
            folds[ii] = np.arange(len(ii)) % k
    else:
        g = pd.Series(groups).astype("category").cat.codes.values
        ug = np.unique(g); rng.shuffle(ug)
        gf = {gg: i % k for i, gg in enumerate(ug)}
        folds = np.array([gf[x] for x in g])
    for f in range(k):
        tr, te = folds != f, folds == f
        if len(np.unique(y[tr])) < 2 or te.sum() == 0:
            continue
        m = make_pipeline(StandardScaler(), LogisticRegression(max_iter=2000))
        m.fit(X.values[tr], y[tr])
        pred[te] = m.predict_proba(X.values[te])[:, 1]
    ok = np.isfinite(pred)
    return roc_auc_score(y[ok], pred[ok]) if len(np.unique(y[ok])) == 2 else np.nan


out = []
for mach in ("imi_vm20i", "imi_vmx30ui"):
    d0 = T[(T.machine == mach) & T.label.isin([0, 1])]
    for name, cols in SETS.items():
        cols = [c for c in cols if c in d0 and d0[c].notna().any()]
        if not cols:
            continue
        d = d0.dropna(subset=cols)
        V.set(f"s1_nfeat_{mach}_{name}", len(cols), "features", {"population": "screening setting", "filter": f"{mach}, feature set {name}", "statistic": "features available and used"})
        res = {}
        for scheme in ("random", "run"):
            rng = np.random.default_rng(SEED)
            vals = [oof_auc(d, cols, None if scheme == "random" else d.run_id.values, 5, rng) for _ in range(REPEATS)]
            vals = np.array([v for v in vals if np.isfinite(v)])
            res[scheme] = vals
            est = {"population": "labelled cutting paths", "filter": f"{mach}, labels 0 and 1, feature set {name}",
                   "statistic": f"median over repeats of out-of-fold AUC, five folds, {scheme} split"}
            V.set(f"s1_auc_{mach}_{name}_{scheme}", float(np.median(vals)), "AUC", est)
            V.set(f"s1_auc_{mach}_{name}_{scheme}_min", float(vals.min()), "AUC", {**est, "statistic": f"min over repeats, {scheme} split"})
            V.set(f"s1_auc_{mach}_{name}_{scheme}_max", float(vals.max()), "AUC", {**est, "statistic": f"max over repeats, {scheme} split"})
        V.set(f"s1_gap_{mach}_{name}", float(np.median(res["random"]) - np.median(res["run"])), "AUC",
              {"population": "labelled cutting paths", "filter": f"{mach}, labels 0 and 1, feature set {name}",
               "statistic": "median random-split AUC minus median run-split AUC"})
        out.append({"machine": mach, "set": name, "n": int(len(d)), "random": float(np.median(res["random"])), "run": float(np.median(res["run"]))})
# S3: tool identity alone (no fitting, so no split is involved): AUC of the indicator "tool A or A*" for label 1 against label 0
for mach in ("imi_vm20i", "imi_vmx30ui"):
    d = T[(T.machine == mach) & T.label.isin([0, 1])]
    isA = (d.Tool.str.replace("*", "", regex=False) == "A").astype(float).values
    V.set(f"s3_auc_toola_{mach}", float(roc_auc_score(d.label.values, isA)), "AUC",
          {"population": "labelled cutting paths", "filter": f"{mach}, labels 0 and 1", "statistic": "AUC of the indicator tool A or A* (ties count half)"})
json.dump(out, open(f"{OUT}/screen_s1.json", "w"), indent=1)
V.save()
print(pd.DataFrame(out).to_string())