"""q_common.py: the analysis functions pre-registered in Phase 3 and used unchanged in Phase 4.

Phase 3 calls only the parts that describe the population or restate results already seen in Phase 2
(run-held-out predictions of the Phase 2 screen model, marginal counts). The tool-family-held-out
predictions, the agreement statistic and the by-tool contrast are computed only in Phase 4, after the
plan is frozen. This file is frozen with the plan: its sha256 is recorded at the Phase 3 stop, and any
change in Phase 4 is a deviation that goes to the deviations log.

Conventions
  family      tool family: the tool name without the asterisk (A and A* form family A)
  run_id      machine and run together; the unit of resampling
  scores      out-of-fold probabilities; the run scheme keeps a row of scores per repeat, and its statistic
              is the mean over repeats of each repeat's AUC (auc_over_repeats)
  weights     a run bootstrap is a matrix of path weights (the multiplicity of each path's run in a resample)
"""
import json
import numpy as np
import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import make_pipeline

SEED = 20260922                      # registered as seed_analysis (q02_judgments.py); the check asserts equality
FEATS = {
    "highrate": ["acc0_rms", "acc1_rms", "s0_rms", "s1_rms", "s2_rms"],
    "highrate_no_s0": ["acc0_rms", "acc1_rms", "s1_rms", "s2_rms"],
    "mtc": ["mtc_smarms", "mtc_smapeak", "mtc_smcrest", "mtc_smvrms", "mtc_S_Axis_Load", "mtc_watt"],
    "clip": ["s0_fullscale"],
    "s2only": ["s2_rms"],
}
META_NUM = ["N", "ap", "ae", "fz"]
SEG = "/home/claude/mill/raw/Phase2/derived/seg_features"


def load_paths(machine="imi_vm20i", labels=(0, 1)):
    T = pd.read_parquet("/home/claude/mill/outputs/segments.parquet")
    T = T[(T.machine == machine) & T.label.isin(labels)].copy()
    T["family"] = T.Tool.str.replace("*", "", regex=False)
    T["y"] = (T.label == 1).astype(int)
    return T.reset_index(drop=True)


def design(df, fset):
    """Feature matrix: logs of amplitudes and MTConnect levels; a one-hot design for the metadata set (E3)."""
    if fset in ("meta", "meta_no_family"):
        X = pd.DataFrame({c: np.log(df[c].astype(float)) for c in META_NUM})
        X["up"] = (df.direction.str.upper() == "UP").astype(float)
        X = X.join(pd.get_dummies(df.Workpiece, prefix="wp", dtype=float))
        if fset == "meta":
            X = X.join(pd.get_dummies(df.family, prefix="fam", dtype=float))
        return X
    cols = FEATS[fset]
    X = df[cols].astype(float).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], 1e-9, None))
    return X


def model(kind="lr", C=1.0, balanced=True):
    cw = "balanced" if balanced else None
    if kind == "lr":
        return make_pipeline(StandardScaler(), LogisticRegression(max_iter=5000, C=C, class_weight=cw))
    if kind == "rf":
        return RandomForestClassifier(n_estimators=500, min_samples_leaf=5, class_weight=cw, random_state=SEED, n_jobs=-1)
    raise ValueError(kind)


def winsorize_fit(Xtr, Xte, q=0.01):
    lo, hi = Xtr.quantile(q), Xtr.quantile(1 - q)
    return Xtr.clip(lo, hi, axis=1), Xte.clip(lo, hi, axis=1)


def fit_predict(Xtr, ytr, Xte, kind="lr", C=1.0, winsor=False, balanced=True):
    Xtr, Xte = Xtr.align(Xte, join="left", axis=1, fill_value=0.0)
    if winsor:
        Xtr, Xte = winsorize_fit(Xtr, Xte)
    m = model(kind, C, balanced)
    m.fit(Xtr.values, ytr)
    return m.predict_proba(Xte.values)[:, 1]


def run_folds(run_ids, k, rng):
    runs = np.array(sorted(set(run_ids)))
    rng.shuffle(runs)
    f = {r: i % k for i, r in enumerate(runs)}
    return np.array([f[r] for r in run_ids])


def oof_run_matrix(df, fset, repeats=20, k=5, kind="lr", C=1.0, winsor=False, balanced=True, seed=SEED):
    """Run-held-out out-of-fold scores, one row per repeat (every path is out of fold once per repeat), so that
    every score in a row comes from a single fitted model, as in the family scheme."""
    X = design(df, fset); y = df.y.values
    rng = np.random.default_rng(seed)
    S = np.full((repeats, len(df)), np.nan)
    for rep in range(repeats):
        folds = run_folds(df.run_id.values, k, rng)
        for f in range(k):
            tr, te = folds != f, folds == f
            if len(np.unique(y[tr])) < 2:
                raise RuntimeError("a run-held-out training fold lacks a label")
            S[rep, te] = fit_predict(X[tr], y[tr], X[te], kind, C, winsor, balanced)
    return S


def oof_run(df, fset, repeats=20, k=5, kind="lr", C=1.0, winsor=False, balanced=True, seed=SEED):
    """Run-held-out scores averaged over repeats (used only to restate the Phase 2 screen in Phase 3)."""
    return oof_run_matrix(df, fset, repeats, k, kind, C, winsor, balanced, seed).mean(axis=0)


def auc_over_repeats(y, S, w=None):
    """The confirmatory run-held-out statistic: mean over repeats of the pooled AUC of each repeat's scores."""
    return float(np.mean([auc_w(y, s, w) for s in np.atleast_2d(S)]))


def oof_family(df, fset, kind="lr", C=1.0, winsor=False, balanced=True):
    """Tool-family-held-out out-of-fold scores: one fold per family."""
    X = design(df, fset); y = df.y.values
    s = np.full(len(df), np.nan)
    for fam in sorted(df.family.unique()):
        te = (df.family == fam).values; tr = ~te
        if len(np.unique(y[tr])) < 2:
            raise RuntimeError(f"training set without family {fam} lacks a label")
        s[te] = fit_predict(X[tr], y[tr], X[te], kind, C, winsor, balanced)
    return s


# ------------------------------------------------------------------ weighted statistics for the run bootstrap
def auc_w(y, s, w=None):
    """Weighted Mann-Whitney AUC with midrank ties."""
    y = np.asarray(y); s = np.asarray(s, float)
    w = np.ones(len(y)) if w is None else np.asarray(w, float)
    u, inv = np.unique(s, return_inverse=True)
    pw = np.bincount(inv, weights=w * (y == 1), minlength=len(u))
    nw = np.bincount(inv, weights=w * (y == 0), minlength=len(u))
    below = np.concatenate([[0.0], np.cumsum(nw)[:-1]])
    Wp, Wn = pw.sum(), nw.sum()
    if Wp == 0 or Wn == 0:
        return np.nan
    return float((pw * (below + 0.5 * nw)).sum() / (Wp * Wn))


def within_family_auc(y, s, fam, families, w=None):
    """Pair-weighted mean of the AUCs computed inside each listed family (no comparison across families)."""
    y = np.asarray(y); fam = np.asarray(fam)
    w = np.ones(len(y)) if w is None else np.asarray(w, float)
    num = den = 0.0
    for f in families:
        m = fam == f
        pairs = (w[m] * (y[m] == 1)).sum() * (w[m] * (y[m] == 0)).sum()
        if pairs > 0:
            num += pairs * auc_w(y[m], np.asarray(s)[m], w[m]); den += pairs
    return num / den if den > 0 else np.nan


def kappa_w(a, b, w=None):
    """Cohen's kappa of two binary indicators, path weights w."""
    a = np.asarray(a).astype(int); b = np.asarray(b).astype(int)
    w = np.ones(len(a)) if w is None else np.asarray(w, float)
    W = w.sum()
    po = (w * (a == b)).sum() / W
    pa, pb = (w * a).sum() / W, (w * b).sum() / W
    pe = pa * pb + (1 - pa) * (1 - pb)
    return float((po - pe) / (1 - pe)) if pe < 1 else np.nan


def _strata_shares(a, b, strata, w):
    a = np.asarray(a).astype(int); b = np.asarray(b).astype(int); strata = np.asarray(strata)
    w = np.ones(len(a)) if w is None else np.asarray(w, float)
    W = w.sum(); out = []
    for s in np.unique(strata):
        m = strata == s; ws = w[m].sum()
        if ws > 0:
            out.append((ws / W, (w[m] * a[m]).sum() / ws, (w[m] * b[m]).sum() / ws))
    return np.array(out)


def _kappa_from(po, ps):
    ma, mb = (ps[:, 0] * ps[:, 1]).sum(), (ps[:, 0] * ps[:, 2]).sum()
    pe = ma * mb + (1 - ma) * (1 - mb)
    return float((po - pe) / (1 - pe)) if pe < 1 else np.nan


def kappa_null_w(a, b, strata, w=None):
    """Kappa expected if a and b were independent within each stratum (the strata of H3: the label)."""
    ps = _strata_shares(a, b, strata, w)
    po = (ps[:, 0] * (ps[:, 1] * ps[:, 2] + (1 - ps[:, 1]) * (1 - ps[:, 2]))).sum()
    return _kappa_from(po, ps)


def kappa_max_w(a, b, strata, w=None):
    """Largest kappa the per-stratum shares of a and b allow."""
    ps = _strata_shares(a, b, strata, w)
    return _kappa_from((ps[:, 0] * (1 - np.abs(ps[:, 1] - ps[:, 2]))).sum(), ps)


def kappa_threshold(k_null, k_max, frac):
    """The H3 rule: the kappa that lies frac of the way from the label-only expectation to the attainable maximum."""
    return k_null + frac * (k_max - k_null)


def share_diff_w(ind, in_group, w=None, by_run=None):
    """Share of ind in the group minus the share outside it; path-weighted, or run-weighted when by_run is given."""
    ind = np.asarray(ind, float); g = np.asarray(in_group, bool)
    w = np.ones(len(ind)) if w is None else np.asarray(w, float)
    if by_run is not None:          # every run counts once: mean of the within-run shares, run weights from w
        r = np.asarray(by_run)
        vals = []
        for grp in (True, False):
            num = den = 0.0
            for run in np.unique(r[g == grp]):
                m = (r == run) & (g == grp)
                wr = w[m].max() if w[m].size else 0.0
                if wr > 0:
                    num += wr * ind[m].mean(); den += wr
            vals.append(num / den if den > 0 else np.nan)
        return vals[0] - vals[1]
    sa = (w[g] * ind[g]).sum() / w[g].sum() if w[g].sum() > 0 else np.nan
    so = (w[~g] * ind[~g]).sum() / w[~g].sum() if w[~g].sum() > 0 else np.nan
    return sa - so


def run_weights(run_ids, B, seed=SEED):
    """B run-bootstrap resamples as path weights (the multiplicity of each path's run)."""
    runs = np.array(sorted(set(run_ids)))
    idx = {r: i for i, r in enumerate(runs)}
    col = np.array([idx[r] for r in run_ids])
    rng = np.random.default_rng(seed)
    draws = rng.integers(0, len(runs), size=(B, len(runs)))
    counts = np.stack([np.bincount(d, minlength=len(runs)) for d in draws])
    return counts[:, col].astype(float)


def run_weights_strat(run_ids, run_group, B, seed=SEED):
    """Run bootstrap within groups of runs (H4: family A runs and the other runs are resampled separately)."""
    run_ids = np.asarray(run_ids); run_group = np.asarray(run_group)
    rng = np.random.default_rng(seed)
    out = np.zeros((B, len(run_ids)))
    for g in sorted(set(run_group)):
        m = run_group == g
        runs = np.array(sorted(set(run_ids[m])))
        idx = {r: i for i, r in enumerate(runs)}
        col = np.array([idx[r] for r in run_ids[m]])
        draws = rng.integers(0, len(runs), size=(B, len(runs)))
        counts = np.stack([np.bincount(d, minlength=len(runs)) for d in draws])
        out[:, m] = counts[:, col]
    return out


def nonfinite_share(draws):
    d = np.asarray(draws, float)
    return float(1 - np.isfinite(d).mean()) if d.size else 1.0


def percentile_ci(draws, level=0.95):
    """Percentile interval of the finite draws; NaN when none is finite (the caller reports nonfinite_share)."""
    d = np.asarray(draws, float); d = d[np.isfinite(d)]
    if not d.size:
        return float("nan"), float("nan")
    return float(np.quantile(d, (1 - level) / 2)), float(np.quantile(d, 1 - (1 - level) / 2))


def one_sided_p(draws, null):
    """Bootstrap p-value for 'statistic <= null' (share of resamples at or below the null, with the +1 rule);
    NaN when no draw is finite, so that the classifier records NOT_TESTED rather than a failure."""
    d = np.asarray(draws, float); d = d[np.isfinite(d)]
    if not d.size:
        return float("nan")
    return float((np.sum(d <= null) + 1) / (len(d) + 1))


def holm(pvals):
    """Holm step-down: adjusted p-values in the input order (untested predictions are left out by the caller)."""
    p = np.asarray(pvals, float); m = len(p); order = np.argsort(p)
    adj = np.empty(m); running = 0.0
    for rank, i in enumerate(order):
        running = max(running, min(1.0, (m - rank) * p[i])); adj[i] = running
    return adj


# ------------------------------------------------------------------ spectral indicator from the stored spectra
def nh_indicator(df, ch, comb_half_bins=2, band=(200.0, 10000.0), min_spindle_hz=50.0, rel_speed_tol=0.0):
    """1 where the largest line in the band lies off every spindle harmonic (outside +/- comb bins), 0 otherwise,
    NaN where the spindle is too slow for clean comb gaps. Recomputed from the Welch spectra of the device pass;
    with the default arguments it reproduces the Phase 2 indicator (D5), which the Phase 3 check asserts."""
    out = np.full(len(df), np.nan)
    cache = {}
    for i, r in enumerate(df.itertuples()):
        key = (r.machine, int(r.run))
        if key not in cache:
            z = np.load(f"{SEG}/{r.machine}_dataset{int(r.run)}.npz", allow_pickle=True)
            meta = json.loads(str(z["meta"]))
            cache[key] = (z, meta)
        z, meta = cache[key]
        if f"{ch}_psd" not in z.files:
            continue
        j = int(np.argmin(np.abs(z["start"] - r.start) + np.abs(z["end"] - r.end)))
        P = z[f"{ch}_psd"][j].astype(float)
        df_hz = meta["channels"][ch]["df"]
        f = np.arange(len(P)) * df_hz
        fsp = r.N / 60.0
        chk = (f >= 50.0) & (f < 12000.0)
        if fsp < min_spindle_hz or not np.isfinite(P[chk]).all():
            continue
        inband = (f >= band[0]) & (f < band[1])
        harm = np.zeros(len(f), bool)
        kk = np.arange(1, int(10000.0 / fsp) + 2)
        idx = np.round(kk * fsp / df_hz).astype(int)
        if rel_speed_tol > 0:         # variant V6c: each zone widens with the harmonic's frequency
            half = np.maximum(comb_half_bins, np.ceil(rel_speed_tol * kk * fsp / df_hz)).astype(int)
            for i0, h in zip(idx, half):
                harm[max(i0 - h, 0):min(i0 + h + 1, len(f))] = True
        else:
            for d in range(-comb_half_bins, comb_half_bins + 1):
                ii = idx + d; harm[ii[(ii >= 0) & (ii < len(f))]] = True
        nh, hh = inband & ~harm, inband & harm
        if not nh.any() or not hh.any():
            continue
        out[i] = float(P[nh].max() > P[hh].max())
    return out