"""Second base lineage (SI figure): the Fisher-Muller and headroom results on SmolLM2-1.7B-Instruct
beside the Qwen2.5-0.5B-Instruct originals, mean ± 95% CI over seeds.

(A) merged specialists vs the best single specialist, easy benchmark (5 seeds per lineage);
(B) union (routing) vs fusion (soup, ties) on the hard benchmark (3 seeds per lineage).
Skips silently when the SmolLM2 bundles are not present yet (``make figures`` runs every script).

Usage:  python figures/plot_llm_smol.py [out_dir=results/llm_merge_seeds_smol]
"""

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_seed_bundles, savefig, letter_axes  # noqa: E402
from plot_llm_seeds import _agg, _best_spec  # noqa: E402

LINEAGES = {"Qwen2.5-0.5B": ("results/llm_merge_seeds", "results/llm_moe_hard_seeds", "#9ecae1", "#2c7fb8"),
            "SmolLM2-1.7B": ("results/llm_merge_seeds_smol", "results/llm_moe_hard_seeds_smol", "#fdae6b", "#d62728")}
PANELS = {"merge": (["best_specialist", "merge_soup", "merge_ties"],
                    ["best\nspecialist", "merge\n(soup)", "merge\n(ties)"],
                    "Fisher–Muller, easy benchmark"),
          "moe": (["best_specialist", "merge_soup", "merge_ties", "moe_oracle", "moe_learned"],
                  ["best\nspecialist", "fusion\n(soup)", "fusion\n(ties)", "union\n(oracle)", "union\n(learned)"],
                  "union vs fusion, hard benchmark")}


def main(out_dir: str = "results/llm_merge_seeds_smol") -> None:
    if not all(Path(d).exists() for d in LINEAGES["SmolLM2-1.7B"][:2]):
        print("plot_llm_smol: SmolLM2 bundles not present yet; skipping"); return
    fig, axes = plt.subplots(1, 2, figsize=(12, 4.4))
    for ax, (key, (models, labels, title)) in zip(axes, PANELS.items()):
        x = np.arange(len(models)); n_l = len(LINEAGES); w = 0.8 / (2 * n_l)
        for li, (lineage, (dm, dmo, c_over, c_worst)) in enumerate(LINEAGES.items()):
            df = _best_spec(load_seed_bundles(dm if key == "merge" else dmo)[0])
            n = df["seed"].nunique()
            for mi, (metric, color) in enumerate((("overall", c_over), ("worst_family", c_worst))):
                vals = _agg(df, models, metric)
                off = (li * 2 + mi - (2 * n_l - 1) / 2) * w
                ax.bar(x + off, [v for v, _ in vals], w, yerr=[e for _, e in vals], capsize=2,
                       color=color, label=f"{lineage}, {metric.replace('_', ' ')} ({n} seeds)")
        ax.set_xticks(x); ax.set_xticklabels(labels, fontsize=8)
        ax.set(ylabel="verifier accuracy", ylim=(0, 1.0), title=title)
        ax.legend(frameon=False, fontsize=7)
    fig.tight_layout()
    letter_axes(fig)
    savefig(fig, out_dir, "llm_smol")


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