"""q05_synthetic.py: exercise the frozen estimators of q_confirm.py on synthetic data with planted effects.

Called by q04_check.py. Builds small datasets in memory with the structure of the real one (tool families,
runs nested in families, labels clustered by run, positive amplitude features, spindle speeds) and runs
auc_drop, cooccurrence, boot_weights and holm_adjust on them. It never reads the dataset: the check's static
scan asserts that this file loads no data file. Returns a list of (condition, description) pairs.
"""
import numpy as np
import pandas as pd

import q_common as QC
import q_confirm as QF

FAMS = {"A": 8, "E": 6, "B": 1, "C1": 2, "D": 3}


def make_paths(seed, family_shift, paths_per_run=24):
    """Paths with a latent condition that raises label 1 odds and the amplitudes; family_shift adds a
    family-specific level to the amplitudes and to the label odds, which a family-held-out model cannot learn."""
    rng = np.random.default_rng(seed)
    rows = []
    run = 0
    for fam, nruns in FAMS.items():
        off = rng.normal(0, 1.0) * family_shift
        base_rate = {"A": 0.5, "E": 0.25, "B": 0.0, "C1": 0.0, "D": 0.1}[fam]
        for _ in range(nruns):
            run += 1
            run_rate = np.clip(base_rate + rng.normal(0, 0.1), 0, 0.9) if base_rate > 0 else 0.0
            speed = rng.choice([4000, 8000, 10000, 12000])
            for _ in range(paths_per_run):
                y = int(rng.random() < run_rate)
                lvl = 1.2 * y + off + rng.normal(0, 0.6)
                rows.append({"machine": "synthetic", "run": run, "run_id": f"s_{run}", "family": fam, "Tool": fam, "N": float(speed), "y": y, "label": y,
                             **{c: float(np.exp(lvl + rng.normal(0, 0.3))) for c in QC.FEATS["highrate"]}})
    return pd.DataFrame(rows)


def run(B=300, repeats=3, k=3):
    checks = []
    add = lambda cond, what: checks.append((bool(cond), what))
    # ---------------------------------------------------------------- weights: stratified by family
    df = make_paths(1, family_shift=1.5)
    W = QF.boot_weights(df, B, 7)
    fam_present = all(all(W[b, (df.family == f).values].sum() > 0 for f in FAMS) for b in range(B))
    add(fam_present, "S boot_weights keeps every family in every resample")
    same_run = all(len(set(W[b, (df.run_id == r).values])) == 1 for b in range(0, B, 50) for r in df.run_id.unique())
    add(same_run, "S boot_weights gives every path of a run the same weight")
    tot = W[:, (df.family == "A").values].sum(axis=1)
    add(np.allclose(tot, (df.family == "A").sum()), "S boot_weights keeps each family's size (runs of equal size)")
    # ---------------------------------------------------------------- auc_drop: planted family shift against none
    res, sc = QF.auc_drop(df, "highrate", B, 7, 0.95, repeats, k)
    add(all(np.isfinite([res["auc_run"], res["auc_family"], res["drop"], res["drop_lo"], res["drop_hi"], res["drop_p"]])), "S auc_drop returns finite results")
    add(res["nonfinite"] == 0, "S auc_drop has no undefined draw on complete data")
    add(res["drop_lo"] <= res["drop"] <= res["drop_hi"], "S auc_drop interval contains its estimate")
    add(abs(res["drop"] - (res["auc_run"] - res["auc_family"])) < 1e-12, "S auc_drop drop equals the difference of its AUCs")
    add(sc["scores_run"].shape == (repeats, len(df)) and np.isfinite(sc["scores_run"]).all(), "S every path is out of fold once per repeat")
    add(res["drop"] > 0 and res["drop_p"] < 0.05, f"S auc_drop detects a planted family shift (drop {res['drop']:.3f}, p {res['drop_p']:.4f})")
    df0 = make_paths(2, family_shift=0.0)
    res0, _ = QF.auc_drop(df0, "highrate", B, 7, 0.95, repeats, k)
    add(abs(res0["drop"]) < 0.05, f"S auc_drop stays near zero without a family shift (drop {res0['drop']:.3f})")
    res_w, _ = QF.auc_drop(df, "highrate", B, 7, 0.95, repeats, k, W=W)
    add(abs(res_w["drop"] - res["drop"]) < 1e-12 and abs(res_w["drop_p"] - res["drop_p"]) < 1e-12, "S auc_drop with a shared weight matrix reproduces its own resamples")
    # ---------------------------------------------------------------- cooccurrence: planted shared line against independence
    rng = np.random.default_rng(3)
    lab1 = df[df.y == 1].reset_index(drop=True)
    cls = QF.speed_class(lab1.N.values, 9000)
    shared = rng.random(len(lab1)) < np.where(cls == 1, 0.7, 0.4)
    ia = np.where(rng.random(len(lab1)) < 0.9, shared, rng.random(len(lab1)) < 0.3).astype(float)
    ib = np.where(rng.random(len(lab1)) < 0.85, shared, rng.random(len(lab1)) < 0.3).astype(float)
    ia[:3] = np.nan
    rc = QF.cooccurrence(ia, ib, lab1, cls, B, 7, 0.95, 0.5)
    add(rc["n_dropped"] == 3 and rc["n_paths"] == len(lab1) - 3, "S cooccurrence drops undefined indicators and counts them")
    add(rc["kappa_null"] < rc["thr"] < rc["kappa_max"], "S cooccurrence threshold lies strictly between the null and the maximum")
    add(abs(rc["thr"] - (rc["kappa_null"] + 0.5 * (rc["kappa_max"] - rc["kappa_null"]))) < 1e-12, "S cooccurrence threshold follows the registered rule")
    add(rc["kappa"] > rc["thr"] and rc["excess_p"] < 0.05, f"S cooccurrence detects planted agreement (kappa {rc['kappa']:.3f}, threshold {rc['thr']:.3f})")
    add(rc["kappa_lo"] <= rc["kappa"] <= rc["kappa_hi"], "S cooccurrence interval contains its estimate")
    ja = (rng.random(len(lab1)) < np.where(cls == 1, 0.7, 0.4)).astype(float)
    jb = (rng.random(len(lab1)) < np.where(cls == 1, 0.6, 0.3)).astype(float)
    r0 = QF.cooccurrence(ja, jb, lab1, cls, B, 7, 0.95, 0.5)
    add(abs(r0["excess"]) < 0.15 and r0["kappa"] < r0["thr"], f"S cooccurrence finds no excess for indicators independent within classes (excess {r0['excess']:.3f})")
    add(r0["kappa_null"] > 0, "S the stratified null is positive when both indicators depend on the speed class")
    # ---------------------------------------------------------------- undefined draws become NaN, never a verdict
    add(np.isnan(QC.one_sided_p([np.nan, np.nan], 0.0)) and all(np.isnan(QC.percentile_ci([np.nan])) ), "S all-undefined draws give NaN p-value and interval")
    add(abs(QC.nonfinite_share([1.0, np.nan, 2.0, np.nan]) - 0.5) < 1e-12, "S undefined draws are counted")
    # ---------------------------------------------------------------- Holm over tested predictions only
    h = QF.holm_adjust({"H1": 0.01, "H2": 0.04, "H3": 0.03}, {"H1", "H2", "H3"})
    add(abs(h["H1"] - 0.03) < 1e-12 and abs(h["H3"] - 0.06) < 1e-12 and abs(h["H2"] - 0.06) < 1e-12, "S holm_adjust matches the Holm step-down by hand")
    h2 = QF.holm_adjust({"H1": 0.01, "H2": 0.04, "H3": 0.03}, {"H1", "H3"})
    add(np.isnan(h2["H2"]) and abs(h2["H1"] - 0.02) < 1e-12 and abs(h2["H3"] - 0.03) < 1e-12, "S holm_adjust leaves untested predictions out of the family")
    return checks