"""Pre-registered analysis for the v2 society (tasks/prereg-llm-society-v2.md §5, §8).

Prints, for each hypothesis, the per-seed quantities, the paired mean ± 95% CI over seeds, the sign
count, and PASS / FAIL against the pre-set threshold. Written before unblinding and exercised on the
smoke bundle; nothing here is chosen after seeing the campaign. Reads only committed bundles (one
bundle directory, or a campaign directory of ``s{seed}_{arm}/`` bundles).

    H1 vertical climb   full best(G) − B₀ ≥ 0.20 ; best newborn(G) − B₀ ≥ 0.15 ; ≥ 6 families ≥ 0.6
    H3 self-consumption no_grounding best(G) ≤ B₀ + 0.05 ; gap(no_grounding) − gap(full) ≥ 0.30
    H4 sex necessity    no_sex best(G) ≤ B₀ + 0.05 in every seed
    H5 diversity        AUC(diversity) full > no_diversity ; no_diversity diversity < 0.1 by gen 6
    H6 where skills die ≤ 20% of family losses in `full` were supplied at ≥ 0.6 by the child's source
    (H2 is deferred: the sex_linear arm is not in the first campaign.)

Usage:  python figures/stats_llm_society.py [results/llm_society_v2]
"""

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import pandas as pd

sys.path.insert(0, str(Path(__file__).parent))
from plot_llm_society import load_any  # noqa: E402

COMPETENT = 0.6


def ci95(x: np.ndarray) -> tuple[float, float]:
    x = np.asarray(x, dtype=float)
    if len(x) < 2:
        return float(x.mean()), float("nan")
    from scipy import stats
    h = stats.t.ppf(0.975, len(x) - 1) * x.std(ddof=1) / np.sqrt(len(x))
    return float(x.mean()), float(h)


def verdict(ok: bool | None) -> str:
    return "n/a " if ok is None else ("PASS" if ok else "FAIL")


def main(results_dir: str = "results/llm_society_v2") -> None:
    df, fams = load_any(Path(results_dir))
    pop, summ, child = (df[df.role == r] for r in ("population", "summary", "child"))
    src = df[df.role == "child_source"]
    seeds = sorted(df.seed.unique())
    G = int(pop.generation.max())
    arms = set(df.arm.unique())
    print(f"bundle: {results_dir}   seeds {seeds}   G = {G}   L = {len(fams)}   arms {sorted(arms)}\n")

    best = pop[pop.metric == "test_overall"].groupby(["arm", "seed", "generation"]).value.max()
    B0 = {s: float(best.xs(s, level="seed").xs(0, level="generation").mean()) for s in seeds}
    print("B₀ (best founder, gen 0) per seed:", {s: round(v, 3) for s, v in B0.items()})

    def at_G(arm, metric_frame, metric=None, gen=G, agg="max"):
        out = {}
        for s in seeds:
            f = metric_frame[(metric_frame.arm == arm) & (metric_frame.seed == s) & (metric_frame.generation == gen)]
            if metric is not None:
                f = f[f.metric == metric]
            if not f.empty:
                out[s] = float(f.value.max() if agg == "max" else f.value.mean())
        return out

    def report(name, per_seed, thr, direction, note=""):
        vals = np.array(list(per_seed.values()))
        if len(vals) == 0:
            print(f"  {name:38s} {verdict(None)}"); return None
        m, h = ci95(vals)
        ok_each = (vals >= thr) if direction == ">=" else (vals <= thr)
        ok = bool(ok_each.sum() >= max(3, len(vals)) if len(vals) >= 3 else ok_each.all())
        print(f"  {name:38s} {verdict(ok)}  mean {m:+.3f} ± {h:.3f}   per seed "
              f"{np.round(vals, 3).tolist()}   {int(ok_each.sum())}/{len(vals)} meet {direction} {thr}  {note}")
        return ok

    # ---------------- H1
    print("\nH1 — vertical climb (full arm)")
    if "full" in arms:
        gain = {s: at_G("full", pop, "test_overall")[s] - B0[s] for s in seeds if s in at_G("full", pop, "test_overall")}
        nb = at_G("full", summ, "best_newborn_overall", gen=G - 1)
        gain_nb = {s: nb[s] - B0[s] for s in nb}
        # families the best agent is competent on, at G
        comp = {}
        for s in seeds:
            f = pop[(pop.arm == "full") & (pop.seed == s) & (pop.generation == G)]
            if f.empty:
                continue
            ov = f[f.metric == "test_overall"].set_index("agent").value
            a = int(ov.idxmax())
            per = f[(f.agent == a) & f.metric.isin([f"test_{x}" for x in fams])].value
            comp[s] = float((per >= COMPETENT).sum())
        report("best agent − B₀ (≥ 0.20)", gain, 0.20, ">=")
        report("best newborn − B₀ (≥ 0.15)", gain_nb, 0.15, ">=")
        report("families competent in best agent (≥ 6)", comp, 6, ">=")
    else:
        print("  full arm absent")

    # ---------------- H3
    print("\nH3 — self-consumption (no_grounding)")
    if {"no_grounding", "full"} <= arms:
        ng = at_G("no_grounding", pop, "test_overall")
        report("no_grounding best − B₀ (≤ 0.05)", {s: ng[s] - B0[s] for s in ng}, 0.05, "<=")
        gap_ng = at_G("no_grounding", summ, "gap_conformity_minus_truth", agg="mean")
        gap_f = at_G("full", summ, "gap_conformity_minus_truth", agg="mean")
        report("gap(no_grounding) − gap(full) (≥ 0.30)", {s: gap_ng[s] - gap_f[s] for s in gap_ng if s in gap_f}, 0.30, ">=")
        ca = summ[(summ.arm == "no_grounding") & (summ.metric == "consensus_acc")]
        slope = {s: float(np.polyfit(g.generation, g.value, 1)[0]) for s, g in ca.groupby("seed") if len(g) > 1}
        report("consensus-accuracy slope, no_grounding (≤ 0)", slope, 0.0, "<=", note="(non-increasing)")
    else:
        print("  arms absent")

    # ---------------- H4
    print("\nH4 — sex necessity (no_sex ceiling)")
    if "no_sex" in arms:
        ns = at_G("no_sex", pop, "test_overall")
        vals = {s: ns[s] - B0[s] for s in ns}
        ok = all(v <= 0.05 for v in vals.values()) if vals else None
        print(f"  {'no_sex best − B₀ (≤ 0.05 in EVERY seed)':38s} {verdict(ok)}  per seed {np.round(list(vals.values()), 3).tolist()}")
    else:
        print("  no_sex arm absent")

    # ---------------- H5
    print("\nH5 — diversity (full vs no_diversity)")
    if {"full", "no_diversity"} <= arms:
        div = summ[summ.metric == "diversity_behav"]
        auc = lambda arm, s: float(np.trapezoid(div[(div.arm == arm) & (div.seed == s)].sort_values("generation").value))
        d_auc = {s: auc("full", s) - auc("no_diversity", s) for s in seeds
                 if not div[(div.arm == "full") & (div.seed == s)].empty and not div[(div.arm == "no_diversity") & (div.seed == s)].empty}
        report("AUC(diversity) full − no_diversity (> 0)", d_auc, 1e-9, ">=")
        g6 = min(6, G)
        nd6 = at_G("no_diversity", summ, "diversity_behav", gen=g6, agg="mean")
        report(f"no_diversity diversity at gen {g6} (< 0.1)", nd6, 0.1, "<=")
    else:
        print("  arms absent")

    # ---------------- H6
    print("\nH6 — where skills die (full arm)")
    if "full" in arms and not src.empty:
        losses, supplied_ok = 0, 0
        for s in seeds:
            fpop = pop[(pop.arm == "full") & (pop.seed == s)]
            fsrc = src[(src.arm == "full") & (src.seed == s)]
            for t in range(G):
                alive_t = {f for f in fams if (fpop[(fpop.generation == t) & (fpop.metric == f"test_{f}")].value >= COMPETENT).any()}
                alive_t1 = {f for f in fams if (fpop[(fpop.generation == t + 1) & (fpop.metric == f"test_{f}")].value >= COMPETENT).any()}
                for f in alive_t - alive_t1:
                    losses += 1
                    sup = fsrc[(fsrc.generation == t) & (fsrc.metric == f"source_{f}")].value
                    supplied_ok += int((sup >= COMPETENT).any())
        frac = supplied_ok / losses if losses else float("nan")
        ok = None if not losses else frac <= 0.20
        print(f"  {'family losses supplied at ≥0.6 (≤ 20%)':38s} {verdict(ok)}  {supplied_ok}/{losses} losses "
              f"({frac:.0%} if any) — skills should die because they arrived diluted, not despite competent supply")
    else:
        print("  no source diagnostics")

    print("\nH2 (union vs linear blend) — deferred: sex_linear not in the first campaign (prereg §12).")


if __name__ == "__main__":
    main(*sys.argv[1:])