"""v2 society figure — E11's three panels at the language-model tier, plus the competence genotype.

Pre-registered layout (tasks/prereg-llm-society-v2.md §8), written before unblinding and run on the
smoke bundle first. Reads only committed bundles: a single bundle directory, or a campaign directory
whose sub-directories ``s{seed}_{arm}/`` each hold a bundle (the PBS array writes one per element).

(A) Best-agent overall test accuracy per arm over generations (solid) with the best *newborn* of each
    generation (dotted) — a climb carried by a surviving founder is visible as such; B₀ (best founder
    at gen 0) dashed. Mean ± 95% CI over seeds.
(B) Behavioural diversity of the population (mean pairwise disagreement).
(C) The self-consumption signature: mean conformity − mean true accuracy.
(D) Competence genotype of the ``full`` arm's best agent: per-family test accuracy × generation, mean
    over seeds — E8's "a genotype no parent had", if it happens.

Usage:  python figures/plot_llm_society.py [results/llm_society_v2 | results/llm_society_v2_smoke]
"""

from __future__ import annotations

import sys
from pathlib import Path

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

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

_ARMS = [("full", "#2ca02c", "full society"),
         ("no_sex", "#ff7f0e", "no sex (no recombination)"),
         ("no_diversity", "#9467bd", "no diversity (greedy)"),
         ("no_grounding", "#d62728", "no grounding (self-consumption)"),
         ("sex_linear", "#1f77b4", "sex by linear blend (H2 control)")]


def load_any(results_dir: Path) -> tuple[pd.DataFrame, list[str]]:
    """One bundle, or every ``*/results.parquet`` below the directory (the campaign layout)."""
    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}")
    df = pd.concat([pd.read_parquet(p / "results.parquet") for p in paths], ignore_index=True)
    fams = yaml.safe_load((paths[0] / "resolved_config.yaml").read_text())["source_config"]["families"]
    return df, list(fams)


def main(results_dir: str = "results/llm_society_v2") -> None:
    rd = Path(results_dir)
    df, fams = load_any(rd)
    pop = df[df.role == "population"]
    summ = df[df.role == "summary"]
    arms = [a for a in _ARMS if a[0] in set(df.arm.unique())]

    best = (pop[pop.metric == "test_overall"].groupby(["arm", "seed", "generation"]).value.max()
            .rename("best").reset_index())
    newborn = summ[summ.metric == "best_newborn_overall"][["arm", "seed", "generation", "value"]]
    b0 = best[best.generation == 0].groupby("seed").best.mean().mean()

    fig, axes = plt.subplots(1, 4, figsize=(21, 4.8))

    def traj(ax, frame, col, title, ylabel, style="-", label_suffix=""):
        for arm, color, label in arms:
            sub = frame[frame.arm == arm]
            if sub.empty:
                continue
            x, m, h = mean_ci(sub, "generation", col)
            ax.plot(x, m, style, color=color, label=(label + label_suffix) if style == "-" else None, lw=2)
            if style == "-":
                ax.fill_between(x, m - h, m + h, color=color, alpha=0.15, lw=0)
        if title:                                   # overlay calls pass "" and must not wipe labels
            ax.set_title(title, fontsize=10); ax.set_xlabel("generation"); ax.set_ylabel(ylabel)

    traj(axes[0], best, "best", "(A) best agent (solid) and best newborn (dotted)\nB₀ = best founder, dashed",
         "overall test accuracy")
    traj(axes[0], newborn.rename(columns={"value": "best"}), "best", "", "", style=":")
    axes[0].axhline(b0, color="k", ls="--", lw=1, label=f"B₀ = {b0:.2f}")
    axes[0].legend(fontsize=8, loc="best")

    div = summ[summ.metric == "diversity_behav"]
    traj(axes[1], div, "value", "(B) population diversity\n(mean pairwise disagreement)", "diversity")
    gap = summ[summ.metric == "gap_conformity_minus_truth"]
    traj(axes[2], gap, "value", "(C) self-consumption signature\nconformity − true accuracy", "gap")
    axes[2].axhline(0, color="k", lw=0.8)

    # (D) competence genotype of the full arm's best agent, families × generations, mean over seeds
    full = pop[pop.arm == ("full" if "full" in set(pop.arm) else arms[0][0])]
    fam_cols = [f"test_{f}" for f in fams]
    idx = full[full.metric == "test_overall"].sort_values("value").groupby(["seed", "generation"]).tail(1)
    keyed = full.set_index(["seed", "generation", "agent", "metric"]).value
    gens = sorted(full.generation.unique())
    heat = np.full((len(fams), len(gens)), np.nan)
    for gi, g in enumerate(gens):
        rows = idx[idx.generation == g]
        vals = np.array([[keyed.get((r.seed, g, r.agent, c), np.nan) for c in fam_cols] for r in rows.itertuples()])
        if len(vals):
            heat[:, gi] = np.nanmean(vals, axis=0)
    im = axes[3].imshow(heat, aspect="auto", cmap="viridis", vmin=0, vmax=1)
    axes[3].set_yticks(range(len(fams))); axes[3].set_yticklabels(fams, fontsize=8)
    axes[3].set_xticks(range(len(gens))); axes[3].set_xticklabels(gens, fontsize=8)
    axes[3].set_xlabel("generation"); axes[3].set_title("(D) competence genotype of the best agent\n(full arm; per-family accuracy)", fontsize=10)
    fig.colorbar(im, ax=axes[3], fraction=0.046, pad=0.02)

    n_seeds = df.seed.nunique()
    fig.suptitle(f"llm_society_v2 — the composed society at LLM scale ({n_seeds} seed{'s' if n_seeds != 1 else ''}, "
                 f"L={len(fams)} families, mean ± 95% CI)", y=1.02)
    fig.tight_layout()
    savefig(fig, rd, "llm_society_v2")


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