"""Composition-decay figure (prereg v3 §8) — written before unblinding.

(A) Composition **surplus** S_t = composed − best single parent, per arm over generations, with the
    zero line: the vertical claim, and whether it survives inheritance.
(B) Own-skill retention q_t per lineage (math on GSM8K, code on MBPP), dry vs grounded — the
    denominators of the prediction.
(C) rho_t, the behavioural correlation between the two lineages: the mechanism, if it rises.
(D) Observed composed accuracy against the framework's forecast Ĉ_t (one free scale, fixed at
    generation 0) — H3, the paper's predictive claim, drawn as a line the data can miss.

Reads only committed bundles: one bundle directory, or a campaign directory of ``s*/`` bundles.
Usage:  python figures/plot_llm_compose.py [results/llm_compose]
"""

from __future__ import annotations

import sys
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

sys.path.insert(0, str(Path(__file__).parent))
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
from _figlib import mean_ci, savefig  # noqa: E402
from llm.compose import predicted_composition  # noqa: E402

ARMS = [("dry", "#d62728", "dry · blending operator"),
        ("grounded", "#2ca02c", "grounded (g = 0.10) · blending"),
        ("dry_cat", "#1f77b4", "dry · union operator (cat)")]


def load_any(results_dir: Path) -> pd.DataFrame:
    if (results_dir / "results.parquet").exists():
        paths = [results_dir]
    else:
        paths = sorted(p.parent for p in results_dir.glob("*/results.parquet"))
    if not paths:
        raise SystemExit(f"no results.parquet under {results_dir}")
    return pd.concat([pd.read_parquet(p / "results.parquet") for p in paths], ignore_index=True)


def series(df: pd.DataFrame, arm: str, metric: str) -> pd.DataFrame:
    return df[(df.arm == arm) & (df.metric == metric)][["seed", "generation", "value"]]


def main(results_dir: str = "results/llm_compose") -> None:
    rd = Path(results_dir)
    df = load_any(rd)
    arms = [a for a in ARMS if a[0] in set(df.arm.unique())]
    n_seeds = df.seed.nunique()
    fig, ax = plt.subplots(1, 4, figsize=(21, 4.6))

    # (A) surplus
    for arm, color, label in arms:
        s = series(df, arm, "surplus")
        if s.empty:
            continue
        x, m, h = mean_ci(s, "generation", "value")
        ax[0].plot(x, m, "-o", color=color, label=label, lw=2, ms=4)
        ax[0].fill_between(x, m - h, m + h, color=color, alpha=0.15, lw=0)
    ax[0].axhline(0, color="k", lw=1, ls="--")
    ax[0].set_title("(A) composition surplus\ncomposed − best single parent", fontsize=10)
    ax[0].set_xlabel("generation"); ax[0].set_ylabel("surplus"); ax[0].legend(fontsize=8)

    # (B) own-skill retention
    for arm, color, _ in arms:
        for metric, ls in (("q_math", "-"), ("q_code", ":")):
            s = series(df, arm, metric)
            if s.empty:
                continue
            x, m, _h = mean_ci(s, "generation", "value")
            ax[1].plot(x, m, ls, color=color, lw=2,
                       label=f"{arm} · {metric.split('_')[1]}" if arm != "dry_linear" else None)
    ax[1].set_title("(B) own-skill retention q_t\nsolid math (GSM8K), dotted code (MBPP)", fontsize=10)
    ax[1].set_xlabel("generation"); ax[1].set_ylabel("accuracy"); ax[1].legend(fontsize=8)

    # (C) rho
    for arm, color, label in arms:
        s = series(df, arm, "rho_behav")
        if s.empty:
            continue
        x, m, h = mean_ci(s, "generation", "value")
        ax[2].plot(x, m, "-o", color=color, label=label, lw=2, ms=4)
        ax[2].fill_between(x, m - h, m + h, color=color, alpha=0.15, lw=0)
    ax[2].set_title("(C) lineage correlation ρ_t\n(agreement on a shared probe)", fontsize=10)
    ax[2].set_xlabel("generation"); ax[2].set_ylabel("ρ"); ax[2].legend(fontsize=8)

    # (D) observed vs predicted, dry arm
    for arm, color, label in arms:
        obs = series(df, arm, "composed_acc").groupby("generation").value.mean()
        qm = series(df, arm, "q_math").groupby("generation").value.mean()
        qc = series(df, arm, "q_code").groupby("generation").value.mean()
        rho = series(df, arm, "rho_behav").groupby("generation").value.mean()
        if obs.empty or len(obs) < 2:
            continue
        pred = predicted_composition(qm.to_numpy(), qc.to_numpy(), rho.to_numpy(), float(obs.iloc[0]))
        ax[3].plot(obs.index, obs.to_numpy(), "-o", color=color, lw=2, ms=4, label=f"{label} observed")
        ax[3].plot(obs.index, pred, "--", color=color, lw=1.5, alpha=0.8,
                   label=f"{label} predicted Ĉ")
    ax[3].set_title("(D) H3: observed vs the closed form\nĈ = c₀·q_math·q_code·(1−ρ)/(1−ρ₀)", fontsize=10)
    ax[3].set_xlabel("generation"); ax[3].set_ylabel("composed accuracy"); ax[3].legend(fontsize=7)

    fig.suptitle(f"llm_compose — does a composed capability survive inheritance? "
                 f"({n_seeds} seed{'s' if n_seeds != 1 else ''}, mean ± 95% CI)", y=1.03)
    fig.tight_layout()
    savefig(fig, rd, "llm_compose")


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