"""E4 figure: multi-teacher recombination — supply vs realisation.

Three panels tell the honest story: (A) union coverage rises with K_T and decorrelation,
matching the exact closed form (recombination *supplies* the tail); (B) that supply is
realised in the pupil only under a union-preserving merge — mean-mixture distillation
dilutes it away (flat in K_T) while max-merge keeps it; (C) the union-surviving gap.
Usage:  python figures/plot_figS8_multiparent_union.py [results/figS8_multiparent_union]
"""

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, letter_axes  # 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/figS8_multiparent_union") -> None:
    df, cfg = load_bundle(results_dir)
    q = cfg["coverage"]["q"]
    K_Ts = sorted(df["K_T"].unique())
    rhos = sorted(df["rho"].unique())
    g0 = df[df["g"] == 0.0]
    colors = plt.cm.viridis(np.linspace(0, 0.85, len(K_Ts)))

    fig, axes = plt.subplots(1, 3, figsize=(15, 4.3))

    # Panel A: union coverage vs rho per K_T, with closed-form overlay
    ax = axes[0]
    for K, c in zip(K_Ts, colors):
        sub = g0[g0["K_T"] == K].groupby("rho")["union_coverage"].mean()
        ax.plot(sub.index, sub.values, "o", color=c, label=f"$K_T$={K}")
        ax.plot(rhos, [U_closed(K, r, q) for r in rhos], "-", color=c, lw=1)
    ax.set(xlabel=r"parent correlation $\rho$", ylabel="union tail coverage",
           title=r"Supply: union matches $U(K_T,\rho,q)$")
    ax.legend(frameon=False, fontsize=8)

    # Panel B: surviving coverage vs rho per K_T — mean (dashed) vs max (solid)
    ax = axes[1]
    for K, c in zip(K_Ts, colors):
        sub = g0[g0["K_T"] == K].groupby("rho")
        ax.plot(sub["surviving_max"].mean().index, sub["surviving_max"].mean().values,
                "-o", color=c, label=f"$K_T$={K}", ms=4)
        ax.plot(sub["surviving_mean"].mean().index, sub["surviving_mean"].mean().values,
                "--", color=c, lw=1, alpha=0.7)
    ax.set(xlabel=r"parent correlation $\rho$", ylabel="surviving tail coverage",
           title="Realised: max-merge (solid) rises;\nmean-mixture (dashed) stays flat")
    ax.legend(frameon=False, fontsize=8)

    # Panel C: surviving vs K_T at rho=0, both operators — the recombination benefit
    ax = axes[2]
    r0 = g0[g0["rho"] == 0.0]
    mx = r0.groupby("K_T")["surviving_max"].agg(["mean", "sem"])
    mn = r0.groupby("K_T")["surviving_mean"].agg(["mean", "sem"])
    ax.errorbar(mx.index, mx["mean"], yerr=1.96 * mx["sem"], fmt="-o",
                color="#1f77b4", capsize=3, label="max-merge (union-preserving)")
    ax.errorbar(mn.index, mn["mean"], yerr=1.96 * mn["sem"], fmt="--s",
                color="#d62728", capsize=3, label="mean-mixture distillation")
    ax.set(xlabel="number of parents $K_T$", ylabel="surviving tail coverage",
           title=r"Benefit needs a union-preserving merge ($\rho=0$)",
           xticks=K_Ts)
    ax.legend(frameon=False, fontsize=9)

    fig.tight_layout()
    letter_axes(fig)
    savefig(fig, results_dir, "figS8_multiparent_union")


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