"""`recombination` figure — the E4 "merge, don't average" finding in REAL weights (↔ E4/C4).

K_T specialist GRUs are trained on assignments from the exact shared-switch retention
construction, so union coverage matches the closed form U(K_T,rho,q) exactly. The pupil then
recombines the *measured* teacher distributions two ways: mean (naive pooling) vs oracle-guided
max-merge (M2N2-style union). The Layer-1 conservation law — mean-mixture is flat in K_T while
max-merge rises — is what must survive the move to trained weights.

Four panels: (A) union coverage vs K_T with the closed-form overlay (recombination *supply*);
(B) analytic teachers — max-merge rises, mean stays flat (the conservation law); (C) the same
on *trained* teacher weights (same signs, compressed/noisier — the smoothing caveat); (D) the
rho=1 control — identical teachers buy nothing under either operator. Reads only the bundle.

Usage:  python figures/plot_recombination.py [results/recombination]
"""

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, mean_ci, savefig  # noqa: E402


def U_closed(K_T, rho, q):
    return rho * q + (1 - rho) * (1 - (1 - q) ** K_T)


def main(results_dir: str = "results/recombination") -> None:
    df, cfg = load_bundle(results_dir)
    q = cfg["coverage"]["q"]
    K_Ts = np.array(sorted(df["K_T"].unique()))

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

    # Panel A: union coverage vs K_T (rho=0 and rho=1) + closed-form overlay.
    ax = axes[0, 0]
    for rho, col in [(0.0, "#1f77b4"), (1.0, "#ff7f0e")]:
        sub = df[df["rho"] == rho]
        _, m, ci = mean_ci(sub, "K_T", "union_coverage")
        ax.errorbar(sorted(sub["K_T"].unique()), m, yerr=ci, fmt="o", color=col, capsize=3,
                    label=fr"union ($\rho$={rho:g})")
        ax.plot(K_Ts, [U_closed(K, rho, q) for K in K_Ts], "-", color=col, lw=1)
    ax.set(xlabel="number of teachers $K_T$", ylabel="union tail coverage",
           title=r"Supply: union matches closed form $U(K_T,\rho,q)$")
    ax.legend(frameon=False, fontsize=9)

    # Panels B/C: mean vs max surviving coverage at rho=0, analytic then trained.
    def operator_panel(ax, mean_col, max_col, title):
        r0 = df[df["rho"] == 0.0]
        for col, style, lab, c in [(max_col, "-o", "max-merge (union)", "#2ca02c"),
                                   (mean_col, "--s", "mean-distill (pool)", "#d62728")]:
            _, m, ci = mean_ci(r0, "K_T", col)
            ax.errorbar(sorted(r0["K_T"].unique()), m, yerr=ci, fmt=style, color=c, capsize=3,
                        label=lab)
        ax.set(xlabel="number of teachers $K_T$", ylabel="surviving tail coverage", title=title)
        ax.legend(frameon=False, fontsize=9)

    operator_panel(axes[0, 1], "surviving_mean_target", "surviving_max_target",
                   r"Analytic teachers ($\rho$=0): max rises, mean flat"
                   "\n(the conservation law)")
    operator_panel(axes[1, 0], "surviving_mean", "surviving_max",
                   r"Trained GRU teachers ($\rho$=0): same signs"
                   "\n(compressed + noisier — smoothing caveat)")

    # Panel D: rho=1 control — identical teachers, both operators flat.
    ax = axes[1, 1]
    r1 = df[df["rho"] == 1.0]
    for col, style, lab, c in [("surviving_max", "-o", "max-merge", "#2ca02c"),
                               ("surviving_mean", "--s", "mean-distill", "#d62728"),
                               ("union_coverage", ":^", "union", "#1f77b4")]:
        _, m, ci = mean_ci(r1, "K_T", col)
        ax.errorbar(sorted(r1["K_T"].unique()), m, yerr=ci, fmt=style, color=c, capsize=3, label=lab)
    ax.set(xlabel="number of teachers $K_T$", ylabel="tail coverage",
           title=r"Control ($\rho$=1, identical teachers):"
                 "\nmore teachers buy nothing")
    ax.legend(frameon=False, fontsize=9)

    fig.suptitle("recombination — 'merge, don't average' holds in real weights: max-merge realises "
                 "the multi-teacher tail benefit, mean-distill conserves collapse", y=1.0, fontsize=12)
    fig.tight_layout()
    savefig(fig, results_dir, "recombination")


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