"""E13 figure — real-weight model speciation with Git Re-Basin.

(A) The barrier decomposition per condition: the linear-mode-connectivity error barrier between two
merged MLPs, split into the part permutation alignment REMOVES (coordinate artefact) and the RESIDUAL it
cannot (reproductive isolation). `shared` ≈ 0; `independent` (same task, different init) is almost all
removable (residual ≈ 0 — same species, different basis); `conflict` (conflicting tasks) is almost all
residual (real isolation). (B) The isolation cliff: residual barrier vs the fraction of conflicting
classes — the real-weight image of E12's cliff, after alignment (so it is not a coordinate artefact).

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")

    fig, axes = plt.subplots(1, 2, figsize=(13, 5))

    # Panel A: removable (coordinate artefact) vs residual (isolation), stacked, per condition.
    ax = axes[0]
    order = [c for c in ["shared", "independent", "conflict"] if c in set(dec["condition"])]
    g = dec.groupby("condition").agg(removable=("removable", "mean"),
                                     residual=("residual", "mean")).reindex(order)
    x = np.arange(len(order))
    ax.bar(x, g["removable"], 0.6, label="removable by alignment\n(coordinate artefact)", color="#9ecae1")
    ax.bar(x, g["residual"], 0.6, bottom=g["removable"], label="residual after alignment\n(reproductive isolation)",
           color="#d62728")
    ax.set_xticks(x); ax.set_xticklabels(order)
    ax.set(ylabel="linear-mode-connectivity error barrier",
           title="Merge barrier = coordinate artefact + residual isolation\n"
           "(same task even across inits is coordinate; conflict is real)")
    ax.legend(frameon=False, fontsize=8)

    # Panel B: the isolation cliff — residual barrier vs conflict fraction.
    ax = axes[1]
    cg = cliff.groupby("conflict_frac").agg(res_m=("residual", "mean"), res_s=("residual", "std"),
                                            nai_m=("barrier_naive", "mean")).reset_index()
    ax.plot(cg["conflict_frac"], cg["nai_m"], "--o", color="#999", lw=1.4, label="naive barrier")
    ax.plot(cg["conflict_frac"], cg["res_m"], "-o", color="#d62728", lw=2, label="residual (after alignment)")
    ax.fill_between(cg["conflict_frac"], cg["res_m"] - cg["res_s"], cg["res_m"] + cg["res_s"],
                    color="#d62728", alpha=0.15)
    ax.set(xlabel="fraction of classes with conflicting labels", ylabel="error barrier",
           ylim=(-0.02, None),
           title="The reproductive-isolation cliff, in real weights\n"
           "(residual rises with task conflict — not removable by alignment)")
    ax.legend(frameon=False, fontsize=9)

    fig.suptitle("E13 — real-weight model speciation: what permutation alignment can and cannot merge",
                 y=1.02, fontsize=13)
    fig.tight_layout()
    savefig(fig, "results/speciation_real", "speciation_real")


if __name__ == "__main__":
    main()