"""Conflict-arrival curricula (SI figure): does the declinable merge's decline rate, and the obligate
merge's collapse, follow the generation at which conflicting conventions arrive?

Left: fraction of proposed merges declined per generation (mean over seeds, 95% CI) for each
curriculum; a filled marker on the curve marks the first generation at which both conflicting
families (boolq, winogrande) are present in every lineage. Right: best-lineage all-families accuracy
of the OBLIGATE society arm per curriculum, same marker. Curricula: Latin square (conflict from
generation 5), decorrelated (from 4), conflict-early (from 2), conflict-late (from 6).

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

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 CONFLICT_FROM, best_lineage, decline_table, load_curriculum  # noqa: E402

STYLE = {"latin": ("#7f7f7f", "Latin square"), "decor": ("#2c7fb8", "decorrelated"),
         "early": ("#d62728", "conflict-early"), "late": ("#2ca02c", "conflict-late")}
OBLIGATE = {"latin": "society", "early": "early_society", "late": "late_society"}


def main(out_dir: str = "results/llm_curriculum_v5_early") -> None:
    df = load_curriculum()
    tab = decline_table(df)
    best = best_lineage(df)
    fig, (a1, a2) = plt.subplots(1, 2, figsize=(10, 3.8))
    for cur, (color, label) in STYLE.items():
        sub = tab[tab["curriculum"] == cur]
        if len(sub):
            x, m, h = mean_ci(sub, "generation", "declined")
            a1.errorbar(x + 1, m, yerr=np.nan_to_num(h), fmt="-o", color=color, capsize=3, label=label,
                        markerfacecolor="white")
            g0 = CONFLICT_FROM[cur]
            if g0 in set(x):
                a1.plot(g0 + 1, m[list(x).index(g0)], "o", color=color, ms=12, markeredgecolor="black", markeredgewidth=1.2)
        arm = OBLIGATE.get(cur)
        ob = best[(best["curriculum"] == cur) & (best["arm"] == arm)] if arm else best.iloc[0:0]
        if len(ob):
            x, m, h = mean_ci(ob, "generation", "value")
            a2.errorbar(x + 1, m, yerr=np.nan_to_num(h), fmt="-o", color=color, capsize=3, label=label,
                        markerfacecolor="white")
            g0 = CONFLICT_FROM[cur]
            if g0 in set(x):
                a2.plot(g0 + 1, m[list(x).index(g0)], "o", color=color, ms=12, markeredgecolor="black", markeredgewidth=1.2)
    a1.set(xlabel="generation", ylabel="fraction of merges declined", ylim=(-0.02, 1.05),
           title="declinable merge: decline rate")
    a2.set(xlabel="generation", ylabel="best-lineage accuracy, all families", ylim=(0.1, 0.9),
           title="obligate merge: accuracy")
    a1.plot([], [], "o", color="white", ms=10, markeredgecolor="black", markeredgewidth=1.2,
            label="first generation with both\nconflicting conventions")
    a1.legend(frameon=False, fontsize=8); a2.legend(frameon=False, fontsize=8)
    for a in (a1, a2):
        a.set_xticks(range(1, 7))
    fig.tight_layout()
    letter_axes(fig)
    savefig(fig, out_dir, "curriculum_timing")


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