"""r01_confirm.py: the pre-registered confirmatory analysis (H1, H2, H3 and their robustness variants).

Runs only after r_freeze.assert_freeze() has checked the four files frozen at the Phase 3 stop. It calls the
frozen estimators of q_confirm.py with the settings registered in values3.json and values.json, loops over the
variants listed in q_confirm.VARIANTS_FOR, applies the Holm step over the predictions that pass both
testability conditions (q_rules.testable), and registers every result in values4.json under the names the plan
lists (inputs/plan.json, "quantities" and "variant_naming"). It classifies nothing: the Phase 4 check does that
with q_rules.classify and writes the tally.

Also saved for later scripts (never used to change a confirmatory result): the out-of-fold scores of the
primary H1 and H2 models and the primary H3 indicators (outputs/r01_arrays.npz).
"""
import sys, json, time
sys.path.insert(0, "/home/claude/mill/phase4/code")
sys.path.insert(0, "/home/claude/mill/phase3/code")
sys.path.insert(0, "/home/claude/mill/code")
import r_freeze
FREEZE = r_freeze.assert_freeze()

import numpy as np
import pandas as pd
from vals import Values
import q_common as QC
import q_confirm as QF
import q_rules as QR

P3, P4 = "/home/claude/mill/phase3", "/home/claude/mill/phase4"
V2 = json.load(open("/home/claude/mill/outputs/values.json"))
V3 = json.load(open(f"{P3}/outputs/values3.json"))
PLAN = json.load(open(f"{P3}/inputs/plan.json"))
V = Values(path=f"{P4}/outputs/values4.json", script="r01_confirm.py")
g = lambda n: (V3[n] if n in V3 else V2[n])["value"]
POP = "labelled cutting paths (label.csv rows)"
MISSING = []          # (name, reason) for any result that could not be registered

# ------------------------------------------------------------------ registered settings (asserted against the frozen code)
B, SEED, LEVEL = int(g("boot_b")), int(g("seed_analysis")), float(g("ci_level"))
REP, K, C = int(g("s1_repeats")), int(g("s1_folds")), float(g("lr_c"))
EDGE, FRAC = float(g("speed_class_edge_rpm")), float(g("thr_kappa_frac"))
assert SEED == QC.SEED
assert QF.VARIANTS["v6a"]["comb_half_bins"] == g("comb_bins_v6a") and QF.VARIANTS["v6b"]["comb_half_bins"] == g("comb_bins_v6b")
assert QF.VARIANTS["v6c"]["rel_speed_tol"] == g("rel_speed_tol") and QF.VARIANTS["v7"]["band"] == (g("d5_band_lo_hz"), g("band_hi_v7_hz"))
BASE_IND = dict(comb_half_bins=int(g("d5_comb_bins")), band=(float(g("d5_band_lo_hz")), float(g("d5_band_hi_hz"))),
                min_spindle_hz=float(g("d5_min_rpm")) / 60.0, rel_speed_tol=0.0)

# ------------------------------------------------------------------ population
paths = QC.load_paths()                                   # imi_vm20i, labels 0 and 1
inv = pd.read_csv("/home/claude/mill/outputs/inventory_runs.csv")
remark = inv.summary_remark.fillna("").str.lower()
# An early stop is a remark that starts with "early stopped"; the Phase 3 substring rule also matched "no early
# stopped" (imi_vmx30ui run 17), which is outside this population, so the imi_vm20i set is the same under both.
EARLY = set(inv[(inv.machine == "imi_vm20i") & remark.str.startswith("early stop")].run)
EARLY_P3 = set(inv[(inv.machine == "imi_vm20i") & remark.str.contains("early stop")].run)
assert EARLY == EARLY_P3, "early-stop rule changes the imi_vm20i population"
VAR_NAME = {v["id"].lower(): v["name"] for v in PLAN["variants"]}


def est(filt, stat):
    return {"population": POP, "filter": filt, "statistic": stat}


def setv(name, value, unit, e, kind="value"):
    if value is None or (isinstance(value, float) and not np.isfinite(value)):
        MISSING.append((name, "not finite"))
        return None
    return V.set(name, float(value) if isinstance(value, (float, np.floating)) else value, unit, e, kind=kind)


def seti(name, point, lo, hi, e, method):
    if not (np.isfinite(lo) and np.isfinite(hi)):
        MISSING.append((name, "interval not finite"))
        return
    V.interval(name, point, float(lo), float(hi), LEVEL, method, e)


Q = PLAN["quantities"]
BOOT_METHOD = "percentile interval at ci_level over boot_b run resamples drawn within tool families (q_confirm.boot_weights), predictions held fixed"

# ------------------------------------------------------------------ H1 and H2
t0 = time.time()
arrays = {}
for H, fset in (("h1", "highrate"), ("h2", "mtc")):
    base_filter = Q[f"{H}_auc_run"]["filter"]
    for vid in ["primary"] + QF.VARIANTS_FOR[H.upper()]:
        setting = {} if vid == "primary" else QF.VARIANTS[vid]
        pop = setting.get("population", "all")
        d = QF.population(paths, pop, EARLY).dropna(subset=QC.FEATS[fset]).reset_index(drop=True)
        kw = dict(kind=setting.get("kind", "lr"), C=C, winsor=setting.get("winsor", False))
        pre = H if vid == "primary" else f"{H}_{vid}"
        filt = base_filter if vid == "primary" else f"{base_filter}, variant {vid.upper()} ({VAR_NAME[vid]})"
        try:
            res, sc = QF.auc_drop(d, fset, B, SEED, LEVEL, REP, K, **kw)
        except Exception as exc:                          # a variant that cannot be computed is reported, never dropped silently
            MISSING.append((pre, f"{type(exc).__name__}: {exc}"))
            continue
        for part, unit in (("auc_run", "AUC"), ("auc_family", "AUC"), ("drop", "AUC")):
            stat = Q[f"{H}_{part}"]["statistic"] if vid == "primary" else f"as {H}_{part}, with the single change of variant {vid.upper()}"
            e = est(filt, stat)
            setv(f"{pre}_{part}", res[part], unit, e)
            seti(f"{pre}_{part}_ci", f"{pre}_{part}", res[f"{part}_lo"], res[f"{part}_hi"], e, BOOT_METHOD)
        setv(f"{pre}_p", res["drop_p"], "probability", est(filt, Q[f"{H}_p"]["statistic"] if vid == "primary" else f"as {H}_p, variant {vid.upper()}; unadjusted, outside the verdict"))
        setv(f"{pre}_nonfinite", res["nonfinite"], "fraction", est(filt, Q[f"{H}_nonfinite"]["statistic"] if vid == "primary" else f"as {H}_nonfinite, variant {vid.upper()}"))
        for n in ("n_paths", "n_runs", "n_families"):
            setv(f"{pre}_{n}", int(res[n]), n.split("_")[1], est(filt, f"{n.split('_')[1]} in the analysis population"))
        if vid == "primary":
            arrays[f"{H}_scores_run"] = sc["scores_run"]
            arrays[f"{H}_scores_family"] = sc["scores_family"]
            arrays[f"{H}_run_id"] = d.run_id.values.astype(str)
        print(f"{pre}: drop {res['drop']:.4f} [{res['drop_lo']:.4f}, {res['drop_hi']:.4f}] p {res['drop_p']:.4f}  ({time.time() - t0:.0f} s)", flush=True)

# ------------------------------------------------------------------ H3
lab1 = paths[paths.y == 1].reset_index(drop=True)
h3_filter = Q["h3_kappa"]["filter"]
ind_cache = {}


def indicators(settings):
    key = json.dumps(settings, sort_keys=True)
    if key not in ind_cache:
        ind_cache[key] = (QC.nh_indicator(lab1, "acc0", **settings), QC.nh_indicator(lab1, "s2", **settings))
    return ind_cache[key]


for vid in ["primary"] + QF.VARIANTS_FOR["H3"]:
    setting = {} if vid == "primary" else QF.VARIANTS[vid]
    ind_set = dict(BASE_IND)
    for k in ("comb_half_bins", "band", "rel_speed_tol"):
        if k in setting:
            ind_set[k] = setting[k]
    ia_all, ib_all = indicators(ind_set)
    keep = QF.population(lab1.assign(_i=np.arange(len(lab1))), setting.get("population", "all"), EARLY)._i.values
    d = lab1.iloc[keep].reset_index(drop=True)
    ia, ib = ia_all[keep], ib_all[keep]
    strata = QF.h3_strata(d, EDGE, setting.get("strata", "speed"))
    pre = "h3" if vid == "primary" else f"h3_{vid}"
    filt = h3_filter if vid == "primary" else f"{h3_filter}, variant {vid.upper()} ({VAR_NAME[vid]})"
    try:
        r = QF.cooccurrence(ia, ib, d, strata, B, SEED, LEVEL, FRAC)
    except Exception as exc:
        MISSING.append((pre, f"{type(exc).__name__}: {exc}"))
        continue
    S = lambda part: Q[f"h3_{part}"]["statistic"] if vid == "primary" else f"as h3_{part}, with the single change of variant {vid.upper()}"
    e_k = est(filt, S("kappa"))
    setv(f"{pre}_kappa", r["kappa"], "kappa", e_k)
    seti(f"{pre}_kappa_ci", f"{pre}_kappa", r["kappa_lo"], r["kappa_hi"], e_k, BOOT_METHOD)
    e_x = est(filt, S("excess"))
    setv(f"{pre}_excess", r["excess"], "kappa", e_x)
    seti(f"{pre}_excess_ci", f"{pre}_excess", r["excess_lo"], r["excess_hi"], e_x, BOOT_METHOD + "; the null recomputed in every resample")
    setv(f"{pre}_kappa_null", r["kappa_null"], "kappa", est(filt, S("kappa_null") if vid != "v9" else "kappa expected if the indicators were independent within strata of tool family crossed with spindle-speed class"))
    setv(f"{pre}_kappa_max", r["kappa_max"], "kappa", est(filt, S("kappa_max") if vid != "v9" else "largest kappa the per-stratum shares allow, strata of tool family crossed with spindle-speed class"))
    setv(f"{pre}_thr", r["thr"], "kappa", est(filt, S("thr")))
    setv(f"{pre}_p", r["excess_p"], "probability", est(filt, S("p") if vid == "primary" else f"as h3_p, variant {vid.upper()}; unadjusted, outside the verdict"))
    setv(f"{pre}_nonfinite", r["nonfinite"], "fraction", est(filt, S("nonfinite")))
    for n in ("n_paths", "n_dropped", "n_runs"):
        setv(f"{pre}_{n}", int(r[n]), "runs" if n == "n_runs" else "paths", est(filt, {"n_paths": "paths with both indicators defined", "n_dropped": "paths dropped because an indicator is undefined", "n_runs": "runs holding those paths"}[n]))
    if vid == "primary":
        ok = np.isfinite(ia) & np.isfinite(ib)
        cls = QF.speed_class(d.N.values, EDGE)
        for c, cname in ((0, "lo"), (1, "hi")):
            for combo, (va, vb) in {"both": (1, 1), "acc0_only": (1, 0), "s2_only": (0, 1), "neither": (0, 0)}.items():
                m = ok & (cls == c) & (ia == va) & (ib == vb)
                setv(f"h3_tab_{cname}_{combo}", int(m.sum()), "paths",
                     est(f"{h3_filter}, spindle-speed class {'below' if c == 0 else 'at or above'} speed_class_edge_rpm",
                         f"paths with the Acc0 indicator {'on' if va else 'off'} and the Sound2 indicator {'on' if vb else 'off'} (D5 settings)"))
        arrays["h3_acc0"], arrays["h3_s2"], arrays["h3_run_id"] = ia, ib, d.run_id.values.astype(str)
        arrays["h3_start"], arrays["h3_end"] = d.start.values, d.end.values
    print(f"{pre}: kappa {r['kappa']:.3f} [{r['kappa_lo']:.3f}, {r['kappa_hi']:.3f}] null {r['kappa_null']:.3f} max {r['kappa_max']:.3f} thr {r['thr']:.3f} p {r['excess_p']:.4f} n {r['n_paths']} ({time.time() - t0:.0f} s)", flush=True)

# ------------------------------------------------------------------ testability and the Holm step (q_rules, q_confirm)
POPVALS = {**V2, **V3}
now = {**V2, **V3, **V.v}
tested, reasons = set(), {}
for pred in PLAN["predictions"]:
    ok, why = QR.testable(pred, now, POPVALS)
    reasons[pred["id"]] = why
    if ok:
        tested.add(pred["id"])
pv = {p["id"]: (V.v[f"{p['id'].lower()}_p"]["value"] if f"{p['id'].lower()}_p" in V.v else float("nan")) for p in PLAN["predictions"]}
adj = QF.holm_adjust(pv, tested)
for pid, a in adj.items():
    h = pid.lower()
    setv(f"{h}_p_holm", a, "probability", est(Q[f"{h}_p_holm"]["filter"], Q[f"{h}_p_holm"]["statistic"] + " (Holm step-down over the tested predictions, q_confirm.holm_adjust)"))
setv("holm_n_tested", len(tested), "predictions", est("pre-registered predictions H1 to H3", "predictions that passed both testability conditions and entered the Holm step"))
for pid in reasons:
    print(pid, "testable" if pid in tested else f"NOT testable: {reasons[pid]}", "p", pv[pid], "p_holm", adj[pid])

np.savez_compressed(f"{P4}/outputs/r01_arrays.npz", **arrays)
json.dump({"missing": MISSING, "testability": {k: (k in tested, reasons[k]) for k in reasons}, "early_runs": sorted(int(x) for x in EARLY),
           "freeze": FREEZE, "seconds": round(time.time() - t0, 1)}, open(f"{P4}/outputs/r01_log.json", "w"), indent=1)
V.save()
print("r01_confirm: registered", sum(1 for r in V.v.values() if r["source"] == "r01_confirm.py"), "values; missing", MISSING)