"""llm_merge figure — recombining specialist LLMs (blueprint C2/C4, the real-LLM image of E8).

LoRA specialists on disjoint task families are merged (weight-space) into one deployable model. The
recombined model beats any single specialist overall and — the sharper signature — is competent
across *all* families, which no single parent is. Two panels: (A) per-family accuracy for the base,
each specialist, and the merges (each specialist spikes on its own family; the merges are high
everywhere); (B) overall vs worst-family accuracy (the merges dominate both, especially worst-family).
Reads only the committed bundle.

Usage:  python figures/plot_llm_merge.py [results/llm_merge]
"""

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

_FAMS = ["lists", "strings", "arith"]


def _acc(df, model, metric):
    r = df[(df["model"] == model) & (df["metric"] == metric)]["accuracy"]
    return float(r.iloc[0]) if len(r) else float("nan")


def main(results_dir: str = "results/llm_merge") -> None:
    df, cfg = load_bundle(results_dir)
    specialists = sorted(m for m in df["model"].unique() if m.startswith("spec_"))
    merges = sorted(m for m in df["model"].unique() if m.startswith("merge_"))
    models = ["base"] + specialists + merges
    labels = {"base": "base", **{s: s.replace("spec_", "spec:") for s in specialists},
              **{m: m.replace("merge_", "merge:") for m in merges}}
    colors = {"base": "#7f7f7f"}
    for s in specialists:
        colors[s] = "#1f77b4"
    for m in merges:
        colors[m] = "#2ca02c"

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

    # Panel A: per-family accuracy, grouped by model.
    ax = axes[0]
    x = np.arange(len(_FAMS))
    w = 0.8 / len(models)
    for i, mdl in enumerate(models):
        vals = [_acc(df, mdl, f) for f in _FAMS]
        ax.bar(x + (i - (len(models) - 1) / 2) * w, vals, w, label=labels[mdl],
               color=colors[mdl], alpha=0.9 if mdl.startswith("merge_") else 0.7)
    ax.set_xticks(x); ax.set_xticklabels(_FAMS)
    ax.set(ylabel="accuracy", title="Per-family: each specialist spikes on its own family; the merges\n"
           "(green) are competent everywhere (but averaging dilutes some peaks)")
    ax.legend(frameon=False, fontsize=8, ncol=2)

    # Panel B: overall vs worst-family, per model.
    ax = axes[1]
    x2 = np.arange(len(models))
    for off, metric, hatch, lab in [(-0.2, "overall", "", "overall"),
                                    (0.2, "worst_family", "//", "worst family")]:
        ax.bar(x2 + off, [_acc(df, m, metric) for m in models], 0.38,
               color=[colors[m] for m in models], hatch=hatch, alpha=0.85, label=lab,
               edgecolor="white")
    ax.set_xticks(x2); ax.set_xticklabels([labels[m] for m in models], rotation=25, ha="right",
                                          fontsize=8)
    ax.set(ylabel="accuracy", title="Overall (solid) vs worst-family (hatched): the merge "
           "clearly wins\nworst-family (balance); overall it matches the best specialist")
    ax.legend(frameon=False, fontsize=9)

    fig.suptitle("llm_merge — recombining decorrelated specialist LLMs gives the only model competent "
                 f"across all families (balance); overall parity ({cfg['base_model']})", y=1.0, fontsize=12)
    fig.tight_layout()
    savefig(fig, results_dir, "llm_merge")


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