"""Robust statistics for the controlled predictive test (source of the README numbers).

Implements the second external review's four requested analyses (2026-08-11), from committed
artifacts only:

1. condition-clustered bootstrap CIs for each predictor's Spearman rho, and PAIRED bootstrap
   differences between predictors (a significant rho for one and not another is not a significant
   difference — the paired contrast is the honest comparison);
2. sample-structure disclosure (13 conditions x 3 seeds = 39 rows; parents are retrained per
   condition x seed but share task-data seeds across conditions within a seed, so rows are not
   independent — hence clustering by condition);
3. between- vs within-axis decomposition (pooled correlations are partly axis discrimination);
4. the outcome under three references: oracle parent potential (pre-registered primary),
   best parent, and mean parent — reported because the predictor ordering is sensitive to it.

Plus leave-one-condition-out (LOCO) held-out prediction per predictor.

Usage:  python figures/stats_llm_epistasis.py
"""

from __future__ import annotations

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

PREDICTORS = ["epi_conf", "dis_raw", "grad_cos", "delta_cos", "delta_l2", "cross_perf"]


def load() -> pd.DataFrame:
    a = pd.read_parquet("results/llm_epistasis/results.parquet")
    b = pd.read_parquet("results/llm_epistasis_compat/results.parquet")
    df = pd.concat([a, b], ignore_index=True)
    df["cond"] = df["mode"] + "_" + df["x"].astype(str)
    df["parent_a_overall"] = df[["pa_fam_a", "pa_fam_b", "pa_coh"]].mean(axis=1)
    df["parent_b_overall"] = df[["pb_fam_a", "pb_fam_b", "pb_coh"]].mean(axis=1)
    df["pen_oracle"] = df["merge_penalty"]                       # pre-registered primary
    df["pen_best"] = df[["parent_a_overall", "parent_b_overall"]].max(axis=1) - df["merged_overall"]
    df["pen_mean"] = df[["parent_a_overall", "parent_b_overall"]].mean(axis=1) - df["merged_overall"]
    return df


def clustered_bootstrap(df: pd.DataFrame, outcome: str = "pen_oracle", B: int = 4000, seed: int = 0):
    """Percentile CIs for each predictor's rho, resampling CONDITIONS (13 clusters) with replacement."""
    rng = np.random.default_rng(seed)
    conds = df["cond"].unique()
    groups = {c: df[df["cond"] == c] for c in conds}
    boot = {p: np.empty(B) for p in PREDICTORS}
    for i in range(B):
        bs = pd.concat([groups[c] for c in rng.choice(conds, size=len(conds), replace=True)],
                       ignore_index=True)
        for p in PREDICTORS:
            boot[p][i] = spearmanr(bs[p], bs[outcome])[0]
    return boot


def loco(df: pd.DataFrame, outcome: str = "pen_oracle"):
    """Leave-one-condition-out held-out prediction (linear fit per predictor)."""
    out = {}
    for p in PREDICTORS:
        pr, ac = [], []
        for c in df["cond"].unique():
            tr, te = df[df["cond"] != c], df[df["cond"] == c]
            coef = np.polyfit(tr[p], tr[outcome], 1)
            pr += list(np.polyval(coef, te[p])); ac += list(te[outcome])
        rho, pv = spearmanr(pr, ac)
        out[p] = (rho, pv, float(np.sqrt(np.mean((np.array(pr) - np.array(ac)) ** 2))))
    return out


def main() -> None:
    df = load()
    print(f"sample: {df['cond'].nunique()} conditions x {df['seed'].nunique()} seeds = {len(df)} rows")

    print("\n== league table under three outcome references (Spearman rho) ==")
    print(f"{'predictor':>11} {'oracle*':>8} {'best':>8} {'mean':>8}   (*pre-registered primary)")
    for p in PREDICTORS:
        r = [spearmanr(df[p], df[o])[0] for o in ["pen_oracle", "pen_best", "pen_mean"]]
        print(f"{p:>11} {r[0]:+8.3f} {r[1]:+8.3f} {r[2]:+8.3f}")

    boot = clustered_bootstrap(df)
    print("\n== condition-clustered bootstrap 95% CIs (primary outcome) ==")
    for p in PREDICTORS:
        v = boot[p][~np.isnan(boot[p])]
        print(f"{p:>11}: {spearmanr(df[p], df['pen_oracle'])[0]:+.3f}"
              f"  [{np.percentile(v, 2.5):+.3f}, {np.percentile(v, 97.5):+.3f}]")

    print("\n== paired bootstrap |rho| differences ==")
    for a_, b_ in [("epi_conf", "dis_raw"), ("dis_raw", "delta_cos"),
                   ("dis_raw", "grad_cos"), ("epi_conf", "delta_cos")]:
        d = np.abs(boot[a_]) - np.abs(boot[b_]); d = d[~np.isnan(d)]
        print(f"|rho({a_})| - |rho({b_})|: {np.mean(d):+.3f}"
              f"  [{np.percentile(d, 2.5):+.3f}, {np.percentile(d, 97.5):+.3f}]")

    print("\n== leave-one-condition-out held-out prediction ==")
    for p, (rho, pv, rmse) in loco(df).items():
        print(f"{p:>11}: LOCO rho={rho:+.3f} (p={pv:.3g})  rmse={rmse:.3f}")

    print("\n== between- vs within-axis ==")
    print("mean penalty by axis:", df.groupby("mode")["pen_oracle"].mean().round(3).to_dict())
    c_df = df[df["mode"] == "conflict"]
    for p in PREDICTORS:
        r, pv = spearmanr(c_df[p], c_df["pen_oracle"])
        print(f"{p:>11} (conflict axis only, n={len(c_df)}): {r:+.3f} (p={pv:.2g})")


if __name__ == "__main__":
    main()