"""LLM-tier model speciation figure — conflict coherence cliff, de-confounded interference, duration null.

(A) The conflict cliff, read where it is clean: on the shared ambiguous prompts, each parent performs
under its own convention while the 50/50 merge scores below BOTH under either grading — the hybrid
loses precisely the conflicted function (the mu(S) floor made visible). From the "replace" design
(results/llm_speciation).

(B) The de-confounded private-family readout ("add" design, results/llm_speciation_add: private
training held fixed, conflict data added on top): whether the merge's private-family competence
tracks its parents (conflict damage localised to the conflicted function) or falls below them
(interference spreading to shared circuitry).

(C) The duration (emergent) null: over-trained disjoint specialists keep merging well — the merged
model's private-family accuracy stays above the best parent at every duration. The MLP tier's
"no emergent isolation" null generalises to LLM weights in this regime.

The shared frozen base controls a major source of coordinate mismatch (LoRA deltas share its
coordinates), allowing a cleaner test of conflict-associated merging failure — though averaging can
still fail for non-conflict reasons (nonlinear interaction, scaling, capacity).

Usage:  python figures/plot_llm_speciation.py
"""

from __future__ import annotations

import sys
from pathlib import Path

import matplotlib.pyplot as plt

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


def _series(df, mode, model, metric):
    sub = df[(df["mode"] == mode) & (df["model"] == model) & (df["metric"] == metric)]
    g = sub.groupby("x")["accuracy"].agg(["mean", "std"]).reset_index().fillna(0.0)
    return g["x"], g["mean"], g["std"]


def main() -> None:
    rep, _ = load_bundle("results/llm_speciation")
    add, _ = load_bundle("results/llm_speciation_add")
    fam_a = "strings" if (rep["metric"] == "strings").any() else "lists"
    fam_b = "arith"

    fig, axes = plt.subplots(1, 3, figsize=(16.5, 4.9))

    # (A) coherence on the conflicted function (replace design).
    ax = axes[0]
    x, y, _ = _series(rep, "conflict", "parent_a", "ambig_asc")
    ax.plot(x, y, "--o", color="#9ecae1", lw=1.5, label="parent A under its convention (asc)")
    x, y, _ = _series(rep, "conflict", "parent_b", "ambig_desc")
    ax.plot(x, y, "--o", color="#a1d99b", lw=1.5, label="parent B under its convention (desc)")
    x, y, _ = _series(rep, "conflict", "merge_soup", "coherence")
    ax.plot(x, y, "-s", color="#d62728", lw=2.2, label="merge under its BEST convention")
    ax.set(xlabel="fraction of training carrying the conflicting convention",
           ylabel="accuracy on the shared ambiguous prompts", ylim=(-0.02, None),
           title="(A) the hybrid loses the conflicted function\n(below BOTH parents under either grading)")
    ax.legend(frameon=False, fontsize=8)

    # (B) de-confounded private families (add design).
    ax = axes[1]
    mode = "conflict_add"
    for model, color, style, lw in (("merge_soup", "#d62728", "-s", 2.2),
                                    ("parent_a", "#9ecae1", "--o", 1.5),
                                    ("parent_b", "#a1d99b", "--o", 1.5)):
        x, y, s = _series(add, mode, model, "mean_private")
        ax.plot(x, y, style, color=color, lw=lw, label=f"{model}: private families (mean)")
        ax.fill_between(x, y - s, y + s, color=color, alpha=0.15)   # +-1 sd over seeds
    ax.set(xlabel="conflict data added on top of fixed private training",
           ylabel="verifier accuracy", ylim=(-0.02, 1.02),
           title="(B) conflict damage does NOT spread: private families\ntrack the parents at every conflict level (3 seeds, ±1 sd)")
    ax.legend(frameon=False, fontsize=8)

    # (C) duration null.
    ax = axes[2]
    x, y, _ = _series(rep, "duration", "merge_soup", "mean_private")
    ax.plot(x, y, "-o", color="#d62728", lw=2.2, label="merge: private families (mean)")
    x, y, _ = _series(rep, "duration", "parent_a", fam_a)
    ax.plot(x, y, "--o", color="#9ecae1", lw=1.5, label=f"parent A on its own family ({fam_a})")
    x, y, _ = _series(rep, "duration", "parent_b", fam_b)
    ax.plot(x, y, "--o", color="#a1d99b", lw=1.5, label=f"parent B on its own family ({fam_b})")
    ax.set(xlabel="specialist training duration (epochs)", ylabel="verifier accuracy",
           ylim=(-0.02, 1.02),
           title="(C) the emergent test: over-specialisation\ndoes not erode mergeability here")
    ax.legend(frameon=False, fontsize=8)

    fig.suptitle("LLM-tier model speciation: conflict provokes function-specific hybrid breakdown; "
                 "no isolation emerges from duration alone (shared base controls coordinate mismatch — "
                 "a cleaner test of conflict-associated failure)", y=1.03, fontsize=11.5)
    fig.tight_layout()
    savefig(fig, "results/llm_speciation", "llm_speciation")


if __name__ == "__main__":
    main()