"""Pre-registered analysis for the composition experiment (prereg v3 §3) — written before unblinding.

Prints each hypothesis, its per-seed quantities, the paired mean ± 95% CI, and PASS / FAIL against
the threshold fixed in the pre-registration. Nothing here is chosen after seeing the data.

    H1 gate  S_0 >= +0.05, union-exceedance >= 0.03, cat > linear by >= 0.03  (>=2/3 seeds)
    H2       S_t declines (Spearman <= -0.7) and composition's fractional loss exceeds each parent's
    H3       Ĉ_t (one parameter, fixed at t=0) predicts observed with MAE <= 0.05 and beats a
             two-parameter exponential on AIC
    H4       S_G(grounded) - S_G(dry) >= +0.08, 3/3 seeds positive
    H5       rho_t rises in dry (Spearman >= +0.7); partial corr of S_t with rho_t given q_t < 0
    H6       dry_linear: S_0 <= +0.02 and union-exceedance <= 0.01 at every generation

Usage:  python figures/stats_llm_compose.py [results/llm_compose]
"""

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))
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
from plot_llm_compose import load_any, series  # noqa: E402
from llm.compose import predicted_composition  # noqa: E402


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


def spearman(y) -> float:
    from scipy import stats
    y = np.asarray(y, dtype=float)
    return float(stats.spearmanr(np.arange(len(y)), y).statistic) if len(y) > 2 else float("nan")


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


def main(results_dir: str = "results/llm_compose") -> None:
    df = load_any(Path(results_dir))
    seeds = sorted(df.seed.unique())
    G = int(df.generation.max())
    arms = set(df.arm.unique())
    print(f"bundle {results_dir}   seeds {seeds}   G = {G}   arms {sorted(arms)}\n")

    def at(arm, metric, gen):
        s = series(df, arm, metric)
        return {int(r.seed): float(r.value) for r in s[s.generation == gen].itertuples()}

    def traj(arm, metric, seed):
        s = series(df, arm, metric)
        s = s[s.seed == seed].sort_values("generation")
        return s.value.to_numpy()

    # ---------------- H1
    print("H1 — generation-0 gate (does the published effect reproduce here?)")
    s0, u0 = at("dry", "surplus", 0), at("dry", "union_exceedance", 0)
    print(f"  surplus at t=0            {np.round(list(s0.values()), 3).tolist()}   "
          f"{verdict(sum(v >= 0.05 for v in s0.values()) >= max(1, len(s0) - 1))}  (>= +0.05)")
    print(f"  union-exceedance at t=0   {np.round(list(u0.values()), 3).tolist()}   "
          f"{verdict(sum(v >= 0.03 for v in u0.values()) >= max(1, len(u0) - 1))}  (>= 0.03)")
    # The operator arms are named by their operator (`dry` = linear, `dry_cat` = cat) after the
    # gen-0 sweep; compare whichever two are present rather than assuming a name.
    if {"dry", "dry_cat"} <= arms:
        lin0, cat0 = at("dry", "composed_acc", 0), at("dry_cat", "composed_acc", 0)
        d = [lin0[s] - cat0[s] for s in lin0 if s in cat0]
        print(f"  linear − cat at t=0       {np.round(d, 3).tolist()}   (sweep found the ordering is "
              f"weight-dependent; reported, not gated)")

    # ---------------- H2
    print("\nH2 — composition decays, and faster than its parents")
    for seed in seeds:
        s = traj("dry", "surplus", seed)
        c = traj("dry", "composed_acc", seed)
        qm, qc = traj("dry", "q_math", seed), traj("dry", "q_code", seed)
        if len(c) < 3:
            continue
        frac_c = c[-1] / c[0] if c[0] > 0 else np.nan
        frac_m = qm[-1] / qm[0] if qm[0] > 0 else np.nan
        frac_q = qc[-1] / qc[0] if qc[0] > 0 else np.nan
        print(f"  seed {seed}: spearman(S_t) {spearman(s):+.2f}   retained: composed {frac_c:.2f} "
              f"vs math {frac_m:.2f}, code {frac_q:.2f}   "
              f"{'faster' if frac_c < min(frac_m, frac_q) else 'NOT faster'}")

    # ---------------- H3
    print("\nH3 — the closed form predicts the trajectory (the paper's predictive claim)")
    for arm in ("dry", "grounded"):
        if arm not in arms:
            continue
        maes, aics = [], []
        for seed in seeds:
            obs = traj(arm, "composed_acc", seed)
            qm, qc = traj(arm, "q_math", seed), traj(arm, "q_code", seed)
            rho = traj(arm, "rho_behav", seed)
            if len(obs) < 3 or not (len(obs) == len(qm) == len(qc) == len(rho)):
                continue
            pred = predicted_composition(qm, qc, rho, float(obs[0]))
            mae = float(np.mean(np.abs(pred - obs)))
            # two-parameter baseline: a*exp(-b t), least squares on the same points
            t = np.arange(len(obs), dtype=float)
            pos = obs > 1e-6
            if pos.sum() >= 2:
                b, loga = np.polyfit(t[pos], np.log(obs[pos]), 1)
                base = np.exp(loga) * np.exp(b * t)
            else:
                base = np.full_like(obs, obs.mean())
            n = len(obs)
            aic = lambda resid, k: n * np.log(max(1e-12, np.mean(resid ** 2))) + 2 * k
            maes.append(mae); aics.append(aic(pred - obs, 1) - aic(base - obs, 2))
        if maes:
            m, h = ci95(maes)
            print(f"  {arm:9s} MAE {m:.3f} ± {h:.3f}  {verdict(m <= 0.05)} (<= 0.05)   "
                  f"ΔAIC vs exponential {np.mean(aics):+.1f} "
                  f"({'closed form wins' if np.mean(aics) < 0 else 'baseline wins'})")

    # ---------------- H4
    print("\nH4 — grounding arrests the decay")
    if {"dry", "grounded"} <= arms:
        d, g = at("dry", "surplus", G), at("grounded", "surplus", G)
        diff = [g[s] - d[s] for s in g if s in d]
        m, h = ci95(diff)
        print(f"  S_G(grounded) − S_G(dry)  {np.round(diff, 3).tolist()}  mean {m:+.3f} ± {h:.3f}   "
              f"{verdict(m >= 0.08 and all(v > 0 for v in diff))}  (>= +0.08, all seeds positive)")

    # ---------------- H5
    print("\nH5 — rising ρ is the mechanism")
    for arm in ("dry", "grounded"):
        if arm not in arms:
            continue
        sp = [spearman(traj(arm, "rho_behav", s)) for s in seeds if len(traj(arm, "rho_behav", s)) > 2]
        if sp:
            print(f"  {arm:9s} spearman(ρ_t) {np.round(sp, 2).tolist()}  mean {np.mean(sp):+.2f}"
                  + (f"   {verdict(np.mean(sp) >= 0.7)}  (>= +0.7)" if arm == "dry" else ""))
    try:
        from scipy import stats
        rows = []
        for seed in seeds:
            s, r = traj("dry", "surplus", seed), traj("dry", "rho_behav", seed)
            qm, qc = traj("dry", "q_math", seed), traj("dry", "q_code", seed)
            if len(s) > 3 and len(s) == len(r) == len(qm) == len(qc):
                rows.append(np.column_stack([s, r, qm * qc]))
        if rows:
            a = np.vstack(rows)
            # partial correlation of S with rho, controlling for q_math*q_code
            res_s = a[:, 0] - np.poly1d(np.polyfit(a[:, 2], a[:, 0], 1))(a[:, 2])
            res_r = a[:, 1] - np.poly1d(np.polyfit(a[:, 2], a[:, 1], 1))(a[:, 2])
            pr = float(stats.pearsonr(res_s, res_r).statistic)
            print(f"  partial corr(S, ρ | q_math·q_code) = {pr:+.2f}   "
                  f"{verdict(pr < 0)}  (negative = lost complementarity, not just retention loss)")
    except Exception as e:                                     # descriptive only, never fatal
        print(f"  partial correlation unavailable ({type(e).__name__})")

    # ---------------- H6
    print("\nH6 (revised) — does the operator ordering hold across generations, or only at gen 0?")
    if {"dry", "dry_cat"} <= arms:
        for gen in range(G + 1):
            lin, cat = at("dry", "surplus", gen), at("dry_cat", "surplus", gen)
            d = [lin[s] - cat[s] for s in lin if s in cat]
            if d:
                print(f"  gen {gen}: surplus linear − cat = {np.mean(d):+.3f}   "
                      f"(linear {np.mean(list(lin.values())):+.3f}, cat {np.mean(list(cat.values())):+.3f})")


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