"""Differential reproduction in the six-generation population (SI figure).

Left: best-lineage all-families accuracy per generation (mean over seeds, 95% CI) for the four arms:
never merge, declinable merge, and each with culling (the lowest-scoring lineage re-founded from
the highest-scoring one after every generation). Middle: population MEAN accuracy over the three
lineages, same arms (culling acts on the mean first). Right: the number of cull events per
generation in each culled arm (mean over seeds), with the fraction of merges declined in the
culled declinable arm.

Reads the committed curriculum bundles through stats_llm_curriculum (no re-simulation).
Usage:  python figures/plot_curriculum_cull.py [out_dir=results/llm_curriculum_v5_cull]
"""

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 mean_ci, savefig, letter_axes  # noqa: E402
from stats_llm_curriculum import best_lineage, load_curriculum  # noqa: E402

ARMS = {"isolated": ("#2c7fb8", "-", "never merge"), "veto": ("#2ca02c", "-", "declinable merge"),
        "cull_isolated": ("#2c7fb8", "--", "never merge + culling"),
        "cull_veto": ("#2ca02c", "--", "declinable merge + culling")}


def main(out_dir: str = "results/llm_curriculum_v5_cull") -> None:
    df = load_curriculum()
    df = df[df["curriculum"] == "latin"]
    best = best_lineage(df)
    lin = df[(df["metric"] == "all_families") & df["model"].str.startswith("lineage") & (df["generation"] >= 0)]
    pop_mean = lin.groupby(["arm", "seed", "generation"])["value"].mean().reset_index()
    fig, (a1, a2, a3) = plt.subplots(1, 3, figsize=(13, 3.8))
    for arm, (color, ls, label) in ARMS.items():
        for ax, src in ((a1, best), (a2, pop_mean)):
            sub = src[src["arm"] == arm]
            if len(sub):
                x, m, h = mean_ci(sub, "generation", "value")
                ax.errorbar(x + 1, m, yerr=np.nan_to_num(h), fmt="o", ls=ls, color=color, capsize=3,
                            label=label)
    for arm, color in (("cull_isolated", "#2c7fb8"), ("cull_veto", "#2ca02c")):
        c = df[(df["arm"] == arm) & (df["metric"] == "culled")]
        if len(c):
            ev = c.groupby(["seed", "generation"])["value"].sum().reset_index()
            x, m, _ = mean_ci(ev, "generation", "value")
            a3.plot(x + 1, m, "o--", color=color, label=f"{ARMS[arm][2]}: culls")
    v = df[(df["arm"] == "cull_veto") & (df["metric"] == "veto_used")]
    if len(v):
        x, m, h = mean_ci(v, "generation", "value")
        a3.errorbar(x + 1, m, yerr=np.nan_to_num(h), fmt="s-", color="#d62728", capsize=3,
                    label="declined merges (culled arm)")
    a1.set(xlabel="generation", ylabel="best-lineage accuracy, all families", ylim=(0.3, 0.9),
           title="best lineage")
    a2.set(xlabel="generation", ylabel="population mean accuracy", ylim=(0.3, 0.9), title="population mean")
    a3.set(xlabel="generation", ylabel="events per generation / fraction", ylim=(-0.05, 1.1),
           title="culls and declines")
    for a in (a1, a2, a3):
        a.set_xticks(range(1, 7)); a.legend(frameon=False, fontsize=7)
    fig.tight_layout()
    letter_axes(fig)
    savefig(fig, out_dir, "curriculum_cull")


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