"""E13 figure — real-weight model speciation: what alignment can and cannot merge, and what emerges.

(A) The barrier decomposition per condition, at two alignment strengths: the linear-mode-connectivity
error barrier between two merged MLPs, naive vs after Git Re-Basin permutation alignment vs after
alignment modulo the FULL ReLU unit symmetry group (scale-canonicalisation + permutation, E13c).
`independent` (same task, different init) is a coordinate artefact — either alignment removes ~all of
it; `conflict` (contradictory label maps) survives both — real reproductive isolation, not a missed
symmetry (cf. arXiv:2606.23607).

(B) The isolation cliff as hybrid fitness: sweeping the fraction of conflicting classes, the residual
(full-symmetry) barrier rises while the merged (midpoint) model's accuracy falls 0.97 -> 0.03 — the
real-weight image of E12's compatible -> depression -> inviability trajectory.

(C) Emergent divergence (E13b): children specialising on disjoint classes (or divergent input
conventions) from a shared fork develop NO residual barrier at any divergence — instead the merge
RESCUES the two forgetting specialists (Fisher-Muller), holding ~0.95 while the parents decay.
Speciation requires functional conflict; it does not emerge from compatible specialisation here.

Usage:  python figures/plot_speciation_real.py
"""

from __future__ import annotations

import sys
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np

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


def main() -> None:
    dec, _ = load_bundle("results/speciation_real")
    cliff, _ = load_bundle("results/speciation_real_cliff")
    emer, _ = load_bundle("results/speciation_real_emergent")

    fig, axes = plt.subplots(1, 3, figsize=(17.5, 5))

    # Panel A: naive / residual(perm) / residual(perm+scale) per condition.
    ax = axes[0]
    order = [c for c in ["shared", "independent", "conflict"] if c in set(dec["condition"])]
    g = dec.groupby("condition").agg(naive=("barrier_naive", "mean"),
                                     res_p=("residual", "mean"),
                                     res_s=("residual_scale", "mean")).reindex(order)
    x = np.arange(len(order)); w = 0.27
    ax.bar(x - w, g["naive"], w, label="naive (no alignment)", color="#9ecae1")
    ax.bar(x, g["res_p"], w, label="residual after permutation\n(Git Re-Basin)", color="#fc9272")
    ax.bar(x + w, g["res_s"], w, label="residual after FULL symmetry group\n(scale + permutation)", color="#d62728")
    ax.set_xticks(x); ax.set_xticklabels(order)
    ax.set(ylabel="linear-mode-connectivity error barrier",
           title="(A) coordinate artefact vs functional isolation\n"
                 "(conflict survives the full ReLU symmetry group)")
    ax.legend(frameon=False, fontsize=7.5)

    # Panel B: the cliff — residual barrier and hybrid fitness vs conflict fraction.
    ax = axes[1]
    cg = cliff.groupby("conflict_frac").agg(res_s=("residual_scale", "mean"), sd=("residual_scale", "std"),
                                            nai=("barrier_naive", "mean"),
                                            hyb=("acc_merge_scale", "mean")).reset_index()
    ax.plot(cg["conflict_frac"], cg["nai"], "--o", color="#999", lw=1.4, label="naive barrier")
    ax.plot(cg["conflict_frac"], cg["res_s"], "-o", color="#d62728", lw=2,
            label="residual (full-symmetry alignment)")
    ax.fill_between(cg["conflict_frac"], cg["res_s"] - cg["sd"], cg["res_s"] + cg["sd"],
                    color="#d62728", alpha=0.15)
    ax.set(xlabel="fraction of classes with conflicting labels", ylabel="error barrier",
           ylim=(-0.02, None),
           title="(B) the isolation cliff, in real weights\n(hybrid fitness falls as conflict rises)")
    ax2 = ax.twinx()
    ax2.plot(cg["conflict_frac"], cg["hyb"], "-s", color="#2c7fb8", lw=1.8, label="merged-model accuracy")
    ax2.set_ylabel("merged (hybrid) accuracy", color="#2c7fb8")
    ax2.tick_params(axis="y", labelcolor="#2c7fb8"); ax2.set_ylim(-0.02, 1.02)
    lines, labels = ax.get_legend_handles_labels()
    l2, la2 = ax2.get_legend_handles_labels()
    ax.legend(lines + l2, labels + la2, frameon=False, fontsize=7.5, loc="center left")

    # Panel C: emergent divergence — no isolation; the merge rescues the forgetting specialists.
    ax = axes[2]
    colors = {"shared": "#999999", "disjoint": "#2c7fb8", "augment": "#41ab5d"}
    for cond in ["shared", "disjoint", "augment"]:
        sub = emer[emer["condition"] == cond]
        if sub.empty:
            continue
        m = sub.groupby("t_div").agg(merge=("acc_merge_scale", "mean"),
                                     pa=("acc_parent_a", "mean"), pb=("acc_parent_b", "mean"),
                                     res=("residual_scale", "mean")).reset_index()
        ax.plot(m["t_div"], m["merge"], "-o", color=colors[cond], lw=2, label=f"{cond}: merged")
        if cond == "disjoint":
            ax.plot(m["t_div"], (m["pa"] + m["pb"]) / 2, "--", color=colors[cond], lw=1.2,
                    label="disjoint: parents (forgetting)")
        max_res = float(emer[emer["condition"] != "shared"]["residual_scale"].max())
    ax.set_xscale("log")
    ax.set(xlabel="divergence (post-fork training steps)", ylabel="accuracy on the full task",
           ylim=(0, 1.02),
           title="(C) emergent divergence does NOT speciate —\n"
                 f"the merge rescues the specialists (max residual = {max_res:.3f})")
    ax.legend(frameon=False, fontsize=7.5, loc="center left")

    fig.suptitle("E13 — real-weight model speciation: isolation requires functional conflict; "
                 "alignment (even modulo the full symmetry group) cannot remove it, and compatible "
                 "specialists merge into a rescuing generalist", y=1.03, fontsize=12)
    fig.tight_layout()
    savefig(fig, "results/speciation_real", "speciation_real")


if __name__ == "__main__":
    main()