"""Publication figures for the PNAS draft — unified, lettered, codename-free.

Renders fig1 (the experimental-programme schematic) and re-plots every data panel directly from the
committed results artifacts (figs/fig2.pdf .. fig7.pdf): no experiment codenames, no suptitles, no per-panel headline titles
(interpretation lives in the captions), bold panel letters, one consistent style. The per-experiment
figures under results/ remain the exploratory versions; these are the manuscript's.

Usage:  python paper/pnas/make_figs.py
"""

from __future__ import annotations

import sys
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "figures"))
sys.path.insert(0, str(ROOT / "src"))
import os
os.chdir(ROOT)                                            # load_bundle uses repo-relative paths
from _figlib import load_bundle, mean_ci  # noqa: E402

OUT = ROOT / "paper" / "pnas" / "figs"

plt.rcParams.update({
    "font.size": 8, "axes.labelsize": 8.5, "legend.fontsize": 7, "legend.frameon": False,
    "lines.markersize": 3.6, "axes.spines.top": False, "axes.spines.right": False,
})


def letter(ax, s, x=-0.14):
    ax.text(x, 1.02, s, transform=ax.transAxes, fontsize=12, fontweight="bold", va="bottom")


def save(fig, name):
    OUT.mkdir(exist_ok=True)
    fig.savefig(OUT / f"{name}.pdf", bbox_inches="tight")
    plt.close(fig)
    print("wrote", OUT / f"{name}.pdf")


def _icon(svg_name: str):
    """Rasterise a committed icon SVG at 2048 px (print-lossless at the ~0.3 in placed size).

    Requires ``rsvg-convert`` (librsvg). The SVGs are the committed source of truth; no derived
    PNGs are kept in the repo.
    """
    import subprocess
    import tempfile

    svg = OUT / "icons" / svg_name
    with tempfile.NamedTemporaryFile(suffix=".png") as f:
        try:
            subprocess.run(["rsvg-convert", "-w", "1024", "-h", "1024", "-o", f.name, str(svg)],
                           check=True, capture_output=True)
        except FileNotFoundError as e:
            raise RuntimeError("rsvg-convert (librsvg) is required to rasterise the icon SVGs "
                               "for fig1a") from e
        return plt.imread(f.name)


# ---------------------------------------------------------------- fig 1: experimental programme
def fig1a():
    from matplotlib.patches import FancyBboxPatch

    # (name, architecture, guarantee, edge, cell face, header fill, header text colour, icon)
    # Icons: Flaticon #2347052 (green pea, for Mendel) and #10479785 (robot) as committed SVGs,
    # used under GG's paid Flaticon licence; rasterised at build time by _icon().
    TIERS = [
        ("Biological model", "Wright\u2013Fisher simulator (NumPy)", "closed forms \u00b7 bitwise-reproducible",
         "#4e8d4e", "#eef6ec", "#c5e0bd", "#2d5b2d", "pea.svg"),
        ("Trained networks", "RNN \u00b7 MLP \u00b7 VAE on a synthetic oracle;\nconvolutional VAE on MNIST",
         "sign-level tests \u00b7 exact oracles", "#5b9bc9", "#eff6fb", "#c9e2f2", "#1f4e79",
         "robot.svg"),
        ("Language models", "LoRA specialists on Qwen 0.5B & 7B;\nexact-match verifier",
         "seed-replicated signs", "#3c6ea5", "#e7eef8", "#adc8e8", "#1d3f66", "robot.svg"),
    ]
    ROWS = [
        ("Grounding = immigration",
         "fresh verified samples from a\nfixed external source enter the\ntraining mix every generation",
         ["immigration\u2013drift equilibrium:\n$g \\approx 0.05$ retains $\\geq$95% diversity;\nobservation floor $1-e^{-mp}$",
          "collapse & rescue in every\narchitecture; MNIST: dry 30$\\to$1 modes,\n10% grounding holds 30/30;\nestimator-bias learning kernel",
          "LIT:established at LLM scale in\nprior work (refs. 21, 30);\nnot re-run here"]),
        ("Recombination = sex",
         "a child inherits from several\nparents, reassembling variants\nthat arose in different lineages",
         ["blending conservation law\n(first-order cancellation);\nunion-operator gain; Fisher\u2013Muller",
          "merge rescues two forgetting\nspecialists ($\\approx$0.50 $\\to$ 0.955)",
          "merged specialists beat every parent\n(5 seeds at 0.5B; 7B); routing vs\naveraging: the headroom rule"]),
        ("Epistasis (entangled skills)",
         "a variant's fitness contribution\ndepends on the variants present\nat the other loci",
         ["NK landscapes: outbreeding\ndepression; directed sex restores\nthe gain; mate-pool breadth optimum",
          None,
          "bred-and-screened offspring beat\nthe blind blend in every seed\n(hard, unsaturated tasks)"]),
        ("The composed society",
         "selection, recombination,\ndiversity preservation and\ngrounding on one population",
         ["four-arm ablation: grounding, sex,\ndiversity each removed\n$\\to$ three distinct failures",
          None,
          "OPEN"]),
        ("Speciation",
         "reproductive isolation: diverged\nlineages no longer produce\nviable (mergeable) offspring",
         ["BDM incompatibility model:\nisolation cliff; quadratic snowball",
          "barrier decomposition under\npermutation+rescaling; conflict\nsweep 0.97$\\to$0.03; emergent null",
          "convention conflict $\\to$ hybrid\nbreakdown; duration null; pre-merge\npredictive test (13 cond. $\\times$ 3 seeds)"]),
    ]

    TAGS = [("Fig. 2A", "Fig. 2B", None),
            ("Fig. 4", "Table S2", "Fig. 3A\u2013B"),
            ("Fig. 5", None, "Table S2"),
            ("Fig. 6", None, None),
            ("Fig. 7A\u2013B", "Fig. 7C\u2013D", "Figs. 7E\u2013F, 3C\u2013D")]

    fig, ax = plt.subplots(figsize=(11.4, 5.3))
    ax.set_axis_off()
    ax.set_xlim(0, 1)
    ax.set_ylim(0, 1)
    ax.text(0.0, 0.995, "A", fontsize=13, fontweight="bold", va="top")
    x0, gap, sep = 0.205, 0.008, 0.02              # sep: extra gutter between theory and the AI pair
    cw = (1.0 - x0 - sep) / 3
    xs = [x0, x0 + cw + sep, x0 + 2 * cw + sep]
    row_h, row_top = 0.157, 0.805

    for j2, (name, arch, guarantee, edge, face, headfill, textcol, icon) in enumerate(TIERS):
        x = xs[j2]
        xc = x + cw / 2 - 0.02                              # text centred left of the icon slot
        ax.add_patch(FancyBboxPatch((x + gap, 0.825), cw - 2 * gap, 0.170,
                                    boxstyle="round,pad=0.004", fc=headfill, ec=edge, lw=1.6))
        ax.text(xc, 0.988, name, ha="center", va="top", fontsize=9.5,
                fontweight="bold", color=textcol)
        ax.text(x + cw / 2 - 0.030, 0.938, arch, ha="center", va="top", fontsize=6.6,
                linespacing=1.3, color=textcol)
        ax.text(x + cw / 2 - 0.030, 0.831, guarantee, ha="center", va="bottom", fontsize=6.4,
                style="italic", color=textcol, alpha=0.85)
        # Reason: imshow + interpolation="none" embeds the icon unsampled in the PDF (an
        # OffsetImage is always composited at figure dpi, i.e. ~39 px, whatever the source).
        img = _icon(icon)
        iw = 56.0 / 1140.0                                  # 56 display px on an 11.4in/100dpi fig
        ih = iw * 11.4 / 5.3
        icx, icy = x + cw - gap - 0.030, 0.906
        ax.imshow(img, extent=(icx - iw / 2, icx + iw / 2, icy - ih / 2, icy + ih / 2),
                  interpolation="none", aspect="auto", zorder=5)
        ax.set_xlim(0, 1)
        ax.set_ylim(0, 1)

    for i2, (label, definition, cells) in enumerate(ROWS):
        tags = TAGS[i2]
        y1 = row_top - i2 * row_h
        y0 = y1 - row_h + 2 * gap
        yc = (y0 + y1) / 2
        ax.text(0.0, y1 - 0.014, label, ha="left", va="top", fontsize=8, fontweight="bold")
        ax.text(0.0, y1 - 0.054, definition, ha="left", va="top", fontsize=6.2,
                style="italic", color="#555", linespacing=1.35)
        for j2, cell in enumerate(cells):
            x = xs[j2]
            edge = TIERS[j2][3]
            face = TIERS[j2][4]
            if cell is None:
                ax.add_patch(FancyBboxPatch((x + gap, y0), cw - 2 * gap, y1 - y0,
                                            boxstyle="round,pad=0.004", fc="white", ec="#bbbbbb",
                                            lw=0.8, ls=(0, (3, 2))))
                ax.text(x + cw / 2, yc, "adds no discriminating test\nat this tier", ha="center",
                        va="center", fontsize=6.4, style="italic", color="#999", linespacing=1.35)
            elif cell.startswith("LIT:"):
                ax.add_patch(FancyBboxPatch((x + gap, y0), cw - 2 * gap, y1 - y0,
                                            boxstyle="round,pad=0.004", fc="white", ec="#bbbbbb",
                                            lw=0.8, ls=(0, (3, 2))))
                ax.text(x + cw / 2, yc, cell[4:], ha="center", va="center",
                        fontsize=6.4, style="italic", color="#777", linespacing=1.35)
            elif cell == "OPEN":
                ax.add_patch(FancyBboxPatch((x + gap, y0), cw - 2 * gap, y1 - y0,
                                            boxstyle="round,pad=0.004", fc="white", ec=edge,
                                            lw=0.8, ls=(0, (3, 2))))
                ax.text(x + cw / 2, yc, "open \u2014 the stated gap", ha="center", va="center",
                        fontsize=6.6, style="italic", color=edge)
            else:
                ax.add_patch(FancyBboxPatch((x + gap, y0), cw - 2 * gap, y1 - y0,
                                            boxstyle="round,pad=0.004", fc=face, ec=edge, lw=0.9))
                ax.text(x + cw / 2, yc + 0.008, cell, ha="center", va="center",
                        fontsize=6.4, linespacing=1.35)
            if tags[j2]:                                    # where the result lives (the ToC role)
                ax.text(x + cw - gap - 0.005, y0 + 0.006, tags[j2], ha="right", va="bottom",
                        fontsize=5.6, style="italic", color=edge)
    save(fig, "fig1a")

# ------------------------------------------------------- fig 1B: society in space -> in time
BLUE, GREEN, GOLD = "#4292c6", "#41ab5d", "#d4a017"


def _robot(ax, x, y, fc="#e8eef5", dots=(), lost=(), w=0.52, h=0.52):
    """A friendly robot head with capability dots beneath (gold = the rare skill)."""
    from matplotlib.patches import Arc, Circle, FancyBboxPatch

    ax.add_patch(FancyBboxPatch((x - w / 2, y - h / 2), w, h, boxstyle="round,pad=0.03",
                                fc=fc, ec="#445", lw=1.1))
    for dx in (-w / 4.5, w / 4.5):
        ax.add_patch(Circle((x + dx, y + h / 7), w / 9, fc="white", ec="#445", lw=0.7))
        ax.add_patch(Circle((x + dx, y + h / 7), w / 24, fc="#334", ec="none"))
    ax.add_patch(Arc((x, y - h / 7), w / 2.4, h / 3.2, theta1=195, theta2=345, ec="#445", lw=1.0))
    ax.plot([x, x], [y + h / 2 + 0.03, y + h / 2 + 0.10], color="#445", lw=1.1)
    ax.add_patch(Circle((x, y + h / 2 + 0.135), 0.04, fc="#445", ec="none"))
    marks = [(c, False) for c in dots] + [(c, True) for c in lost]
    n = len(marks)
    for i, (c, is_lost) in enumerate(marks):
        cx = x + (i - (n - 1) / 2) * 0.19
        cy = y - h / 2 - 0.17
        if is_lost:
            ax.add_patch(Circle((cx, cy), 0.07, fc="white", ec=c, lw=0.9, ls=(0, (2, 2))))
            ax.text(cx, cy - 0.005, "\u00d7", ha="center", va="center", fontsize=6, color=c)
        else:
            ax.add_patch(Circle((cx, cy), 0.07, fc=c, ec="none"))


def fig1b():
    from matplotlib.patches import Circle, FancyArrowPatch

    W, H = 11.4, 4.75
    fig, ax = plt.subplots(figsize=(W, H))
    ax.set_xlim(0, W)
    ax.set_ylim(0, H)
    ax.set_aspect("equal")
    ax.set_axis_off()
    ax.text(0.05, H - 0.05, "B", fontsize=13, fontweight="bold", va="top")

    def arrow(p, q, color="#666", lw=1.2, style="-|>", shrink=2.0, ls="-"):
        ax.add_patch(FancyArrowPatch(p, q, arrowstyle=style, mutation_scale=10, color=color,
                                     lw=lw, linestyle=ls, shrinkA=shrink, shrinkB=shrink))

    # ---------------- left: a society in space (contemporaries exchanging messages)
    cx, cy, r = 2.45, 2.95, 0.95
    dotsets = [(BLUE, GOLD), (GREEN,), (BLUE, GREEN), (GOLD, GREEN), (BLUE,)]
    pos = []
    for i, ds in enumerate(dotsets):
        a = np.pi / 2 + i * 2 * np.pi / 5
        x, y = cx + r * np.cos(a) * 1.4, cy + r * np.sin(a) * 0.95
        pos.append((x, y))
        _robot(ax, x, y, dots=ds)
    for i, j2 in [(0, 2), (1, 3), (2, 4), (0, 3), (1, 4)]:
        arrow(pos[i], pos[j2], color="#99a", lw=0.9, style="<|-|>", shrink=26, ls=(0, (4, 2)))
    ax.text(pos[0][0] + 0.72, pos[0][1] + 0.38, "hi!", fontsize=8, ha="center",
            bbox=dict(boxstyle="round,pad=0.25", fc="white", ec="#99a", lw=0.8))
    clk = (0.55, 4.33)
    ax.add_patch(Circle(clk, 0.21, fc="white", ec="#445", lw=1.1))
    ax.plot([clk[0], clk[0]], [clk[1], clk[1] + 0.13], color="#445", lw=1.0)
    ax.plot([clk[0], clk[0] + 0.10], [clk[1], clk[1]], color="#445", lw=1.0)
    ax.text(clk[0], clk[1] - 0.34, "one moment", ha="center", fontsize=7, style="italic",
            color="#555")
    ax.text(2.45, 0.80, "a society in space", ha="center", fontsize=11, fontweight="bold")
    ax.text(2.45, 0.50, "contemporaries exchanging messages \u2014 multi-agent systems, agent economies",
            ha="center", fontsize=7.2, style="italic", color="#555")
    ax.text(2.45, 0.24, "information is passed on, but not easily stored: it fades with the conversation",
            ha="center", fontsize=7.2, style="italic", color="#555")

    # ---------------- middle: the shift of perspective
    arrow((4.60, 2.75), (5.95, 2.75), color="#445", lw=2.0, style="-|>")
    ax.text(5.27, 2.91, "the same ecosystem,\nseen along its time axis", ha="center",
            va="bottom", fontsize=8, style="italic", color="#334", linespacing=1.3)

    # ---------------- right: a society in time (a pedigree)
    axx = 6.8
    arrow((axx, 4.55), (axx, 1.25), color="#445", lw=1.3)
    for gy, lab in [(4.05, "gen 0"), (2.80, "gen 1"), (1.60, "gen 2")]:
        ax.text(axx - 0.12, gy, lab, ha="right", va="center", fontsize=7.5, color="#445")

    p1, p2 = (8.0, 4.05), (9.5, 4.05)
    c1, c2 = (7.5, 2.80), (9.0, 2.80)
    g2 = (9.0, 1.60)
    _robot(ax, *p1, dots=(BLUE, GOLD))
    _robot(ax, *p2, dots=(GREEN, BLUE))
    _robot(ax, *c1, dots=(BLUE,), lost=(GOLD,))
    _robot(ax, *c2, dots=(BLUE, GREEN, GOLD))
    _robot(ax, *g2, dots=(BLUE, GREEN, GOLD))
    arrow((7.87, 3.51), (7.56, 3.29), color="#666")
    ax.text(7.42, 3.41, "inherit", ha="right", fontsize=7, style="italic", color="#555")
    arrow((8.15, 3.49), (8.85, 3.29), color="#666")
    arrow((9.40, 3.49), (9.13, 3.29), color="#666")
    ax.text(9.0, 3.39, "merge (sex)", ha="center", fontsize=7, style="italic", color="#555",
            bbox=dict(boxstyle="round,pad=0.12", fc="white", ec="none"))
    ax.text(7.5, 2.08, "rare skill lost", ha="center", fontsize=6.8, style="italic",
            color="#a33")
    arrow((9.0, 2.24), (9.0, 2.06), color="#666")
    globe = (10.55, 2.10)
    ax.add_patch(Circle(globe, 0.32, fc="#eaf4fb", ec="#2c7fb8", lw=1.2))
    from matplotlib.patches import Arc as _Arc
    ax.add_patch(_Arc(globe, 0.32, 0.64, theta1=90, theta2=270, ec="#2c7fb8", lw=0.8))
    ax.add_patch(_Arc(globe, 0.32, 0.64, theta1=270, theta2=90, ec="#2c7fb8", lw=0.8))
    ax.plot([globe[0] - 0.32, globe[0] + 0.32], [globe[1], globe[1]], color="#2c7fb8", lw=0.8)
    ax.text(globe[0], globe[1] - 0.44, "reality\n(verifier)", ha="center", va="top", fontsize=7,
            color="#2c7fb8", linespacing=1.2)
    arrow((10.22, 1.95), (9.38, 1.70), color="#2c7fb8", lw=1.2)
    ax.text(9.82, 2.03, "ground\n(immigrate)", ha="center", va="bottom", fontsize=7,
            style="italic", color="#2c7fb8", linespacing=1.2)
    ax.text(8.9, 0.80, "a society in time", ha="center", fontsize=11, fontweight="bold")
    ax.text(8.9, 0.50, "information is inherited, evolutionarily selected, and passed on again \u2014",
            ha="center", fontsize=7.2, style="italic", color="#555")
    ax.text(8.9, 0.24, "from parent model to child model \u2014 where population genetics applies",
            ha="center", fontsize=7.2, style="italic", color="#555")
    save(fig, "fig1b")

# ---------------------------------------------------------------- fig 2: grounding + MNIST
def fig2():
    from knowledge.analysis import critical_grounding, reduce_to_stationary
    from knowledge.metrics import heterozygosity
    from knowledge.truth import make_true_distribution

    df, cfg = load_bundle("results/E2")
    n = cfg["dynamics"]["n"]
    td = make_true_distribution(cfg["truth"]["K"], 1, "zipf", cfg["truth"]["tail_frac"],
                                cfg["truth"]["zipf_s"], 0, tail_threshold=cfg["truth"]["tail_threshold"])
    H_star = heterozygosity(td.p_star)
    last = int(cfg["generations"] * 0.8)
    stat = df[df["generation"] >= last]

    fig, axes = plt.subplots(1, 2, figsize=(10.6, 3.5), gridspec_kw={"width_ratios": [1, 1.35]})

    ax = axes[0]
    st = reduce_to_stationary(stat, value_col="heterozygosity", replicate_col="replicate", last_frac=1.0)
    gg, Hm, Hci = mean_ci(stat, "g", "heterozygosity")
    nz = gg > 0
    ax.errorbar(gg[nz], Hm[nz], yerr=Hci[nz], fmt="o", color="#1f77b4", capsize=2, label="simulation")
    ax.plot(gg[~nz], Hm[~nz], "o", mfc="white", mec="#1f77b4")
    m_of_g = stat.groupby("g")["m"].first().to_numpy()
    m_grid = np.linspace(0, m_of_g.max(), 400)

    def H_eq(m):
        m = np.asarray(m, float)
        return np.where(m <= 0, 0.0, H_star * m * (2 * n + m - 1) / (n + 2 * n * m + m * m))

    ax.plot(m_grid / (n + m_grid), H_eq(m_grid), "k--", lw=1, label="exact equilibrium")
    ax.axhline(H_star, ls=":", color="gray", lw=1, label="source diversity $H^*$")
    r = critical_grounding(st, H_star=H_star, frac=0.95, seed=7)
    ax.axvspan(r["ci_low"], r["ci_high"], color="#d62728", alpha=0.15)
    ax.axvline(r["g_star"], color="#d62728", lw=1.1,
               label=f"95%-retention threshold $g\\approx{r['g_star']:.3f}$")
    ax.set(xlabel="grounding fraction $g = m/(n+m)$", ylabel="stationary diversity $H$")
    ax.legend()
    letter(ax, "A")

    ax = axes[1]
    from PIL import Image
    im = np.asarray(Image.open("results/mnist_collapse/mnist_montage.png"))
    # Strip the baked-in title band and left label margin (raster text is unreadable at panel
    # size); measured on the committed montage: boxes span y >= 69, x >= 75, row centres below.
    top, left = 60, 68
    ax.imshow(im[top:, left:], interpolation="bilinear")
    for yc, g in zip((101.5, 191.5, 282.0, 372.5, 462.5), (0, 4, 8, 12, 15)):
        ax.text(-10, yc - top, str(g), ha="right", va="center", fontsize=8.5)
    ax.text(-0.055, 0.5, "generation", transform=ax.transAxes, rotation=90,
            ha="center", va="center", fontsize=8.5)
    ax.set_axis_off()
    letter(ax, "B", x=-0.02)
    save(fig, "fig2")


# ---------------------------------------------------------------- fig 4: blending vs union + Fisher–Muller
def fig4():
    fig, axes = plt.subplots(1, 2, figsize=(10.6, 3.5))

    df, _ = load_bundle("results/E4")
    r0 = df[(df["g"] == 0.0) & (df["rho"] == 0.0)]
    mx = r0.groupby("K_T")["surviving_max"].agg(["mean", "sem"])
    mn = r0.groupby("K_T")["surviving_mean"].agg(["mean", "sem"])
    ax = axes[0]
    ax.errorbar(mx.index, mx["mean"], yerr=1.96 * mx["sem"], fmt="-o", color="#1f77b4",
                capsize=2, label="union operator (strongest source)")
    ax.errorbar(mn.index, mn["mean"], yerr=1.96 * mn["sem"], fmt="--s", color="#d62728",
                capsize=2, label="output-mean (blending)")
    ax.set(xlabel="number of parents", ylabel="rare capabilities surviving in the child",
           xticks=sorted(r0["K_T"].unique()))
    ax.legend()
    letter(ax, "A")

    df8, cfg8 = load_bundle("results/E8")
    L = cfg8["society"]["L"]
    d0 = df8[df8["rho"] == 0.0]
    ax = axes[1]
    for col, c, lab in [("best_parent", "#7f7f7f", "best single parent"),
                        ("average", "#1f77b4", "blended average"),
                        ("sexual", "#d62728", "recombined offspring")]:
        k, m, ci = mean_ci(d0, "K_T", col)
        ax.errorbar(k, m, yerr=ci, fmt="-o", color=c, capsize=2, label=lab)
    ax.axhline(L, ls=":", color="green", lw=1, label="optimum")
    ax.set(xlabel="number of parents", ylabel="offspring capability")
    ax.legend()
    letter(ax, "B")
    save(fig, "fig4")


# ---------------------------------------------------------------- fig 5: rugged landscapes
def fig5():
    fig, axes = plt.subplots(2, 2, figsize=(10.6, 6.8))

    df9, _ = load_bundle("results/E9")
    Ks = sorted(df9["K"].unique())
    colors = plt.cm.viridis(np.linspace(0, 0.85, len(Ks)))
    bp = df9.groupby("K")["best_parent"].mean()
    ax = axes[0, 0]
    for K, c in zip(Ks, colors):
        s = df9[df9["K"] == K].groupby("rate")["mean_offspring"].mean() - bp[K]
        ax.plot(s.index, s.values, "-o", color=c, label=f"$K$={K}")
    ax.axhline(0, ls=":", color="gray", lw=1)
    ax.set(xlabel="recombination rate", ylabel="mean offspring − best parent")
    ax.legend(title="ruggedness", ncol=2)
    letter(ax, "A")

    df10, _ = load_bundle("results/E10")
    ax = axes[0, 1]
    for col, c, lab in [("global_opt", "green", "global optimum"),
                        ("directed_sex", "#d62728", "screened recombination (directed)"),
                        ("best_parent", "#7f7f7f", "best single parent"),
                        ("random_sex", "#1f77b4", "blind recombination")]:
        k, m, ci = mean_ci(df10, "K", col)
        if col == "global_opt":
            ax.plot(k, m, ":", color=c, label=lab)
        else:
            ax.errorbar(k, m, yerr=ci, fmt="-o", color=c, capsize=2, label=lab)
    ax.set(xlabel="landscape ruggedness $K$", ylabel="offspring capability")
    ax.legend()
    letter(ax, "B")

    df14, _ = load_bundle("results/E14")
    last = df14[df14["generation"] == df14["generation"].max()].copy()
    last["best_n"] = last["best_fitness"] / last["global_opt"]
    K14 = sorted(last["K"].unique())
    cmap = plt.get_cmap("viridis")
    c14 = {K: cmap(i / max(1, len(K14) - 1)) for i, K in enumerate(K14)}
    for ax, col, ylab, L in [(axes[1, 0], "best_n", "best fitness / optimum", "C"),
                             (axes[1, 1], "diversity", "population diversity", "D")]:
        for K in K14:
            g = last[last["K"] == K].groupby("breadth")[col].agg(["mean", "sem"]).reset_index()
            ax.errorbar(g["breadth"], g["mean"], yerr=1.96 * g["sem"].fillna(0), fmt="-o",
                        color=c14[K], capsize=2, label=f"$K$={K}")
        ax.set_xscale("log")
        ax.set(xlabel="mate-pool breadth (monogamous → panmictic)", ylabel=ylab)
        ax.legend(title="ruggedness")
        letter(ax, L)
    save(fig, "fig5")


# ---------------------------------------------------------------- fig 6: the society
def fig6():
    df, _ = load_bundle("results/E11")
    arms = [("full", "#2ca02c", "full system"),
            ("no_sex", "#ff7f0e", "no recombination"),
            ("no_diversity", "#9467bd", "no diversity preservation"),
            ("no_grounding", "#d62728", "no grounded evaluation")]
    arms = [a for a in arms if a[0] in set(df["arm"].unique())]
    g_opt = df["global_opt"].mean()

    fig, axes = plt.subplots(1, 3, figsize=(11.4, 3.2))
    panels = [("best_fitness", "best real fitness", "A"),
              ("diversity", "population diversity", "B"),
              ("conformity_true_gap", "conformity − true fitness", "C")]
    for ax, (col, ylab, L) in zip(axes, panels):
        for name, c, lab in arms:
            sub = df[df["arm"] == name]
            g, m, ci = mean_ci(sub, "generation", col)
            ax.plot(g, m, "-", color=c, lw=1.6, label=lab)
            ax.fill_between(g, m - ci, m + ci, color=c, alpha=0.15)
        if col == "best_fitness":
            ax.axhline(g_opt, ls=":", color="gray", lw=1, label="global optimum")
            ax.legend()
        ax.set(xlabel="generation", ylabel=ylab)
        letter(ax, L)
    save(fig, "fig6")


# ---------------------------------------------------------------- fig 7: speciation, three tiers
def fig7():
    fig, axes = plt.subplots(2, 3, figsize=(11.4, 6.6))

    bdm, _ = load_bundle("results/E12")
    rhos = sorted(bdm["rho"].unique())
    colors = plt.cm.viridis(np.linspace(0.15, 0.85, len(rhos)))

    def agg(df, keys, value):
        g = df.groupby(keys)[value].agg(["mean", "std", "count"]).reset_index()
        g["se"] = g["std"] / np.sqrt(g["count"].clip(lower=1))
        return g

    ax = axes[0, 0]
    par = agg(bdm, "divergence", "parent_fitness")
    ax.plot(par["divergence"], par["mean"], "k--", lw=1.3, label="parents")
    for rho, c in zip(rhos, colors):
        g = agg(bdm[bdm["rho"] == rho], "divergence", "offspring_fitness")
        ax.plot(g["divergence"], g["mean"], "-o", color=c, label=f"hybrid, density {rho:g}")
        ax.fill_between(g["divergence"], g["mean"] - g["se"], g["mean"] + g["se"], color=c, alpha=0.15)
    ax.axhline(0, color="#999", lw=0.7, ls=":")
    ax.set(xlabel="parental divergence (substitutions)", ylabel="fitness")
    ax.legend()
    letter(ax, "A")

    ax = axes[0, 1]
    for rho, c in zip(rhos, colors):
        g = agg(bdm[bdm["rho"] == rho], "divergence", "isolation")
        ax.plot(g["divergence"], g["mean"], "-o", color=c, label=f"{rho:g}")
    ax.set(xlabel="parental divergence (substitutions)", ylabel="P(hybrid inviable)", ylim=(-0.02, 1.02))
    ax.legend(title="incompatibility density")
    letter(ax, "B")

    dec, _ = load_bundle("results/speciation_real")
    order = [c for c in ["shared", "independent", "conflict"] if c in set(dec["condition"])]
    g = dec.groupby("condition").agg(naive=("barrier_naive", "mean"),
                                     res=("residual_scale", "mean")).reindex(order)
    ax = axes[0, 2]
    x = np.arange(len(order)); w = 0.38
    ax.bar(x - w / 2, g["naive"], w, color="#9ecae1", label="before alignment")
    ax.bar(x + w / 2, g["res"], w, color="#d62728", label="after alignment (residual)")
    ax.set_xticks(x)
    ax.set_xticklabels(["same task,\nshared start", "same task,\ndifferent start", "conflicting\ntasks"])
    ax.set(ylabel="merge error barrier")
    ax.legend()
    letter(ax, "C")

    cliff, _ = load_bundle("results/speciation_real_cliff")
    cg = cliff.groupby("conflict_frac").agg(res=("residual_scale", "mean"),
                                            hyb=("acc_merge_scale", "mean")).reset_index()
    ax = axes[1, 0]
    ax.plot(cg["conflict_frac"], cg["res"], "-o", color="#d62728", label="residual barrier")
    ax2 = ax.twinx()
    ax2.plot(cg["conflict_frac"], cg["hyb"], "-s", color="#2c7fb8", label="merged-model accuracy")
    ax2.set_ylabel("merged accuracy", color="#2c7fb8")
    ax2.tick_params(axis="y", labelcolor="#2c7fb8")
    ax2.spines["right"].set_visible(True)
    ax.set(xlabel="fraction of classes in conflict", ylabel="residual barrier")
    l1, la1 = ax.get_legend_handles_labels(); l2, la2 = ax2.get_legend_handles_labels()
    ax.legend(l1 + l2, la1 + la2, loc="center left")
    letter(ax, "D")

    rep, _ = load_bundle("results/llm_speciation")

    def series(df, mode, model, metric):
        sub = df[(df["mode"] == mode) & (df["model"] == model) & (df["metric"] == metric)]
        g = sub.groupby("x")["accuracy"].mean().reset_index()
        return g["x"], g["accuracy"]

    ax = axes[1, 1]
    x_, y_ = series(rep, "conflict", "parent_a", "ambig_asc")
    ax.plot(x_, y_, "--o", color="#9ecae1", label="parent A, own convention")
    x_, y_ = series(rep, "conflict", "parent_b", "ambig_desc")
    ax.plot(x_, y_, "--o", color="#a1d99b", label="parent B, own convention")
    x_, y_ = series(rep, "conflict", "merge_soup", "coherence")
    ax.plot(x_, y_, "-s", color="#d62728", label="merge, best convention")
    ax.set(xlabel="fraction of training in conflict", ylabel="accuracy, shared prompts")
    ax.legend()
    letter(ax, "E")

    ax = axes[1, 2]
    x_, y_ = series(rep, "duration", "merge_soup", "mean_private")
    ax.plot(x_, y_, "-o", color="#d62728", label="merged model")
    x_, y_ = series(rep, "duration", "parent_a", "strings")
    ax.plot(x_, y_, "--o", color="#9ecae1", label="parent A, own task")
    x_, y_ = series(rep, "duration", "parent_b", "arith")
    ax.plot(x_, y_, "--o", color="#a1d99b", label="parent B, own task")
    ax.set(xlabel="specialist training (epochs)", ylabel="accuracy", ylim=(0, 1.02))
    ax.legend()
    letter(ax, "F")
    save(fig, "fig7")


# ---------------------------------------------------------------- fig 3: the language-model tier
def fig3():
    import pandas as pd
    from scipy.stats import spearmanr

    fig, axes = plt.subplots(2, 2, figsize=(10.6, 6.6))

    dfm, _ = load_bundle("results/llm_merge_seeds")
    specs = sorted(m for m in dfm["model"].unique() if m.startswith("spec_"))
    rows = []
    for s, sub in dfm.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)
    dfm = pd.concat([dfm] + rows, ignore_index=True)
    models = ["base", "best_specialist", "merge_soup", "merge_ties"]
    labels = ["base", "best\nspecialist", "merged\n(average)", "merged\n(interference-aware)"]
    ax = axes[0, 0]
    x = np.arange(len(models))
    for off, metric, c, lab in ((-0.19, "overall", "#2c7fb8", "overall"),
                                (0.19, "worst_family", "#d62728", "worst task family")):
        vals, errs = [], []
        for m in models:
            v = dfm[(dfm["model"] == m) & (dfm["metric"] == metric)].groupby("seed")["accuracy"].mean()
            vals.append(v.mean()); errs.append(1.96 * v.std(ddof=1) / max(1, np.sqrt(len(v))))
        ax.bar(x + off, vals, 0.36, yerr=errs, capsize=2, color=c, label=lab)
    ax.set_xticks(x); ax.set_xticklabels(labels, fontsize=7)
    ax.set(ylabel="verifier accuracy")
    ax.legend()
    letter(ax, "A")

    df7, _ = load_bundle("results/llm_moe_hard_hpc")

    def acc7(model, metric):
        r = df7[(df7["model"] == model) & (df7["metric"] == metric)]["accuracy"]
        return float(r.iloc[0]) if len(r) else np.nan
    specs7 = sorted(m for m in df7["model"].unique() if m.startswith("spec_"))
    best7 = max(specs7, key=lambda m: acc7(m, "overall"))
    models7 = [best7, "merge_soup", "merge_ties", "moe_oracle"]
    labels7 = ["best\nspecialist", "merged\n(average)", "merged\n(interference-aware)", "routed\n(kept separate)"]
    ax = axes[0, 1]
    x = np.arange(len(models7))
    for off, metric, c, lab in ((-0.19, "overall", "#2c7fb8", "overall"),
                                (0.19, "worst_family", "#d62728", "worst task family")):
        ax.bar(x + off, [acc7(m, metric) for m in models7], 0.36, color=c, label=lab)
    ax.set_xticks(x); ax.set_xticklabels(labels7, fontsize=7)
    ax.set(ylabel="verifier accuracy")
    ax.legend()
    letter(ax, "B")

    a = pd.read_parquet("results/llm_epistasis/results.parquet")
    b = pd.read_parquet("results/llm_epistasis_compat/results.parquet")
    dfe = pd.concat([a, b], ignore_index=True)
    ax = axes[1, 0]
    for mode, c, mk, lab in (("conflict", "#d62728", "o", "conflicting conventions"),
                             ("duration", "#2c7fb8", "s", "divergence only"),
                             ("compat", "#41ab5d", "^", "overlap, no conflict")):
        sub = dfe[dfe["mode"] == mode]
        ax.scatter(sub["epi_conf"], sub["merge_penalty"], c=c, marker=mk, s=26, alpha=0.75, label=lab)
    ax.axhline(0, color="#999", lw=0.6)
    ax.set(xlabel="pre-merge functional conflict (confidence-weighted)",
           ylabel="merge penalty")
    ax.legend()
    letter(ax, "C")

    preds = [("dis_raw", "raw\ndisagreement"), ("epi_conf", "conf-weighted\nconflict"),
             ("grad_cos", "gradient\nalignment"), ("delta_cos", "weight\ncosine"),
             ("delta_l2", "weight\ndistance"), ("cross_perf", "cross-task\naccuracy")]
    ax = axes[1, 1]
    rhos_ = [abs(spearmanr(dfe[c], dfe["merge_penalty"])[0]) for c, _ in preds]
    cols = ["#fc9272", "#d62728", "#9ecae1", "#9ecae1", "#9ecae1", "#9ecae1"]
    ax.bar(np.arange(len(preds)), rhos_, 0.6, color=cols)
    ax.set_xticks(np.arange(len(preds)))
    ax.set_xticklabels([l for _, l in preds], fontsize=6.5)
    ax.set(ylabel="|Spearman ρ| vs merge penalty", ylim=(0, 0.8))
    letter(ax, "D")
    save(fig, "fig3")


if __name__ == "__main__":
    for f in (fig1a, fig1b, fig2, fig3, fig4, fig5, fig6, fig7):
        f()