"""r02_checks.py: the assumption checks the plan attaches to the confirmatory predictions (A1, A4, A6).

A1  every confirmatory statistic recomputed with each run left out in turn (the frozen estimators, the
    registered settings, the Holm step over the three predictions of each leave-one-run-out set); reports the
    largest change of each statistic and the runs whose removal would flip a decision rule. The class of a
    prediction never changes here; a flip is flagged in the text.
A4  H3 again on paths of at least d7_min_dur_s seconds, where entry and exit weigh less.
A6  H3 again without the paths whose Sound2 channel has more than d9_threshold of its samples at full scale.
Checks are reported next to their prediction and never enter the verdict.
"""
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
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

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"))
V4 = json.load(open(f"{P4}/outputs/values4.json"))
V = Values(path=f"{P4}/outputs/values4.json", script="r02_checks.py")
g = lambda n: (V4[n] if n in V4 else V3[n] if n in V3 else V2[n])["value"]
POP = "labelled cutting paths (label.csv rows)"
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, ALPHA, THR = float(g("speed_class_edge_rpm")), float(g("thr_kappa_frac")), float(g("alpha_family")), float(g("thr_auc_drop"))
est = lambda f, s: {"population": POP, "filter": f, "statistic": s}
BOOT_METHOD = "percentile interval at ci_level over boot_b run resamples drawn within tool families (q_confirm.boot_weights), predictions held fixed"

paths = QC.load_paths()
lab1 = paths[paths.y == 1].reset_index(drop=True)
arr = np.load(f"{P4}/outputs/r01_arrays.npz", allow_pickle=True)
ia_all, ib_all = arr["h3_acc0"], arr["h3_s2"]
assert (arr["h3_run_id"] == lab1.run_id.values.astype(str)).all(), "H3 indicators out of order"


def h3_on(mask):
    d = lab1[mask].reset_index(drop=True)
    return QF.cooccurrence(ia_all[mask], ib_all[mask], d, QF.h3_strata(d, EDGE, "speed"), B, SEED, LEVEL, FRAC), d


# ------------------------------------------------------------------ A4 and A6: H3 on subsets
t0 = time.time()
for tag, mask, filt in (
        ("a4", (lab1.dur >= g("d7_min_dur_s")).values, "imi_vm20i, label 1, Acc0 and Sound2 spectra, paths of at least d7_min_dur_s seconds (check A4)"),
        ("a6", (lab1.s2_fullscale <= g("d9_threshold")).values, "imi_vm20i, label 1, Acc0 and Sound2 spectra, without paths whose Sound2 share of samples at full scale exceeds d9_threshold (check A6)")):
    r, d = h3_on(mask)
    ek = est(filt, "as h3_kappa, on this subset")
    V.set(f"{tag}_kappa", r["kappa"], "kappa", ek)
    V.interval(f"{tag}_kappa_ci", f"{tag}_kappa", r["kappa_lo"], r["kappa_hi"], LEVEL, BOOT_METHOD, ek)
    V.set(f"{tag}_kappa_null", r["kappa_null"], "kappa", est(filt, "as h3_kappa_null, on this subset"))
    V.set(f"{tag}_kappa_max", r["kappa_max"], "kappa", est(filt, "as h3_kappa_max, on this subset"))
    V.set(f"{tag}_thr", r["thr"], "kappa", est(filt, "as h3_thr, on this subset"))
    V.set(f"{tag}_p", r["excess_p"], "probability", est(filt, "as h3_p, on this subset; unadjusted, a check outside the verdict"))
    V.set(f"{tag}_n_paths", r["n_paths"], "paths", est(filt, "paths with both indicators defined"))
    V.set(f"{tag}_n_runs", r["n_runs"], "runs", est(filt, "runs holding those paths"))
    V.set(f"{tag}_n_lo", int((d.N < EDGE).sum()), "paths", est(filt, "paths below speed_class_edge_rpm"))
    V.set(f"{tag}_n_hi", int((d.N >= EDGE).sum()), "paths", est(filt, "paths at or above speed_class_edge_rpm"))
    print(tag, f"kappa {r['kappa']:.3f} [{r['kappa_lo']:.3f}, {r['kappa_hi']:.3f}] thr {r['thr']:.3f} p {r['excess_p']:.4f} n {r['n_paths']}", flush=True)

# ------------------------------------------------------------------ A1: leave one run out
base = {"H1": V4["h1_drop"]["value"], "H2": V4["h2_drop"]["value"], "H3": V4["h3_kappa"]["value"]}
base_rule = {"H1": V4["h1_drop"]["value"] >= THR and V4["h1_p_holm"]["value"] < ALPHA,
             "H2": V4["h2_drop"]["value"] >= THR and V4["h2_p_holm"]["value"] < ALPHA,
             "H3": V4["h3_kappa"]["value"] >= V4["h3_thr"]["value"] and V4["h3_p_holm"]["value"] < ALPHA}
rows = []
for run in sorted(paths.run_id.unique()):
    d = paths[paths.run_id != run].reset_index(drop=True)
    r1, _ = QF.auc_drop(d, "highrate", B, SEED, LEVEL, REP, K, "lr", C, False)
    r2, _ = QF.auc_drop(d, "mtc", B, SEED, LEVEL, REP, K, "lr", C, False)
    m = (lab1.run_id != run).values
    r3, _ = h3_on(m)
    adj = QF.holm_adjust({"H1": r1["drop_p"], "H2": r2["drop_p"], "H3": r3["excess_p"]}, {"H1", "H2", "H3"})
    rule = {"H1": r1["drop"] >= THR and adj["H1"] < ALPHA, "H2": r2["drop"] >= THR and adj["H2"] < ALPHA,
            "H3": r3["kappa"] >= r3["thr"] and adj["H3"] < ALPHA}
    rows.append({"run": run, "h1": r1["drop"], "h2": r2["drop"], "h3": r3["kappa"], "h3_thr": r3["thr"],
                 "p1": adj["H1"], "p2": adj["H2"], "p3": adj["H3"], **{f"rule_{k}": v for k, v in rule.items()}})
    print(run, f"h1 {r1['drop']:.4f} h2 {r2['drop']:.4f} h3 {r3['kappa']:.3f}/{r3['thr']:.3f} rules {rule} ({time.time() - t0:.0f} s)", flush=True)
L = pd.DataFrame(rows)
L.to_csv(f"{P4}/outputs/a1_leave_one_run_out.csv", index=False)
for H, col in (("H1", "h1"), ("H2", "h2"), ("H3", "h3")):
    h = H.lower()
    stat = {"h1": "h1_drop", "h2": "h2_drop", "h3": "h3_kappa"}[h]
    filt = f"imi_vm20i, labels 0 and 1, each of the {len(L)} runs left out in turn (check A1)"
    ch = (L[col] - base[H]).abs()
    V.set(f"a1_{h}_max_change", float(ch.max()), "AUC" if h != "h3" else "kappa", est(filt, f"largest absolute change of {stat} when a single run is left out"))
    V.set(f"a1_{h}_max_change_run", str(L.run[ch.idxmax()]), "text", est(filt, f"run whose removal changes {stat} most"), kind="text")
    V.set(f"a1_{h}_min", float(L[col].min()), "AUC" if h != "h3" else "kappa", est(filt, f"smallest {stat} over the leave-one-run-out sets"))
    V.set(f"a1_{h}_max", float(L[col].max()), "AUC" if h != "h3" else "kappa", est(filt, f"largest {stat} over the leave-one-run-out sets"))
    flips = int((L[f"rule_{H}"] != base_rule[H]).sum())
    V.set(f"a1_{h}_flips", flips, "runs", est(filt, f"runs whose removal would flip the decision rule of {H} (Holm over the three predictions of each set)"))
    if h == "h3":
        V.set("a1_h3_margin_min", float((L.h3 - L.h3_thr).min()), "kappa", est(filt, "smallest h3_kappa minus h3_thr over the leave-one-run-out sets"))
V.set("a1_n_runs", int(len(L)), "runs", est("imi_vm20i, labels 0 and 1", "runs left out one at a time in check A1"))
V.save()
print("r02_checks done", round(time.time() - t0), "s")