"""Per-seed paired contrasts for the 7B language-model runs (source of the SI Table S2 numbers).

The three 7B experiments were single-seed until 2026-09-11; seeds 2–3 run via hpc/llm_7b_seeds.pbs
into ``results/llm_<name>_hpc/s{seed}/``. This script reports, per seed and as mean ± 95% CI:

- Fisher–Muller (``llm_merge_hpc``): merged (soup, ties) − best single specialist, overall and
  worst-family;
- union vs fusion on hard tasks (``llm_moe_hard_hpc``): routing (oracle, learned) − soup;
- directed selection on hard tasks (``llm_directed_hard_hpc``): directed (overall, balanced) − soup.

Reads committed artifacts only; with one seed the CI is reported as n/a rather than invented.

Usage:  python figures/stats_llm_7b_seeds.py
"""

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import pandas as pd
from scipy.stats import ttest_rel

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

METRICS = ("overall", "worst_family")


def with_best_specialist(df: pd.DataFrame) -> pd.DataFrame:
    """Add a ``best_specialist`` model per seed (the spec_* with the highest overall accuracy)."""
    specs = sorted(m for m in df["model"].unique() if m.startswith("spec_"))
    rows = []
    for _, sub in df.groupby("seed"):
        ov = {m: sub[(sub["model"] == m) & (sub["metric"] == "overall")]["accuracy"].mean() for m in specs}
        b = sub[sub["model"] == max(ov, key=ov.get)].copy()
        b["model"] = "best_specialist"
        rows.append(b)
    return pd.concat([df] + rows, ignore_index=True)


def table(df: pd.DataFrame, models: list[str]) -> pd.DataFrame:
    """Per-seed accuracy of each model on each metric, mean ± CI over seeds."""
    rows = []
    for m in models:
        for met in METRICS:
            v = df[(df["model"] == m) & (df["metric"] == met)].groupby("seed")["accuracy"].mean()
            rows.append({"model": m, "metric": met, "n_seeds": len(v),
                         **{f"s{s}": round(a, 3) for s, a in v.items()},
                         "mean": round(v.mean(), 3),
                         "ci95": round(1.96 * v.std(ddof=1) / np.sqrt(len(v)), 3) if len(v) > 1 else np.nan})
    return pd.DataFrame(rows)


def contrasts(df: pd.DataFrame, pairs: list[tuple[str, str]]) -> pd.DataFrame:
    """Per-seed paired differences a − b on each metric."""
    rows = []
    for a, b in pairs:
        for met in METRICS:
            piv = (df[(df["metric"] == met) & df["model"].isin([a, b])]
                   .pivot_table(index="seed", columns="model", values="accuracy"))
            if a not in piv or b not in piv:
                continue
            d = (piv[a] - piv[b]).dropna()
            # paired per-seed t-test (the figures' significance brackets and the captions' p-values)
            p_paired = float(ttest_rel(piv.loc[d.index, a], piv.loc[d.index, b]).pvalue) if len(d) > 1 else np.nan
            rows.append({"contrast": f"{a} − {b}", "metric": met, "n_seeds": len(d), "p_paired": round(p_paired, 4),
                         **{f"s{s}": round(x, 3) for s, x in d.items()},
                         "mean": round(d.mean(), 3),
                         "ci95": round(1.96 * d.std(ddof=1) / np.sqrt(len(d)), 3) if len(d) > 1 else np.nan,
                         "sign_agrees": f"{int((np.sign(d) == np.sign(d.mean())).sum())}/{len(d)}"})
    return pd.DataFrame(rows)


def main() -> None:
    runs = {
        "llm_merge_hpc": (["best_specialist", "merge_soup", "merge_ties"],
                          [("merge_soup", "best_specialist"), ("merge_ties", "best_specialist")]),
        "llm_moe_hard_hpc": (["best_specialist", "merge_soup", "merge_ties", "moe_oracle", "moe_learned",
                              "max_merge"],
                             [("moe_oracle", "merge_soup"), ("moe_learned", "merge_soup"),
                              ("merge_soup", "best_specialist")]),
        "llm_directed_hard_hpc": (["merge_soup", "directed_overall", "directed_balanced"],
                                  [("directed_overall", "merge_soup"), ("directed_balanced", "merge_soup")]),
    }
    for name, (models, pairs) in runs.items():
        d = Path("results") / name
        if not d.exists():
            print(f"## {name}: missing\n")
            continue
        df, cfg = load_seed_bundles(d)
        df = with_best_specialist(df)
        present = [m for m in models if m in set(df["model"])]
        print(f"## {name} — {cfg.get('base_model')}, seeds {sorted(df['seed'].unique())}")
        print(table(df, present).to_string(index=False))
        print()
        print(contrasts(df, pairs).to_string(index=False))
        print()


if __name__ == "__main__":
    main()