"""llm_directed figure — directed sex in weight space (E10): breed offspring, select the fittest.

A population of recombinant offspring (weighted merges of the specialists) is scored on a held-out
validation split by the verifier; the winners (best-overall, best-worst-family) are reported on a
fresh test split against the uniform-soup blend and the best single specialist. Two panels: (A)
per-family accuracy — directed offspring (green) vs soup (orange) vs specialists (blue); (B) overall
vs worst-family, with the best-specialist bar as the parent ceiling. The suptitle reports whether
directed selection beat the single a-priori soup. Reads only the committed bundle.

Usage:  python figures/plot_llm_directed.py [results/llm_directed]
"""

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"]
_DIRECTED = {"directed_overall": "directed:overall", "directed_balanced": "directed:balanced"}


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_directed") -> None:
    df, cfg = load_bundle(results_dir)
    present = set(df["model"].unique())
    specialists = sorted(m for m in present if m.startswith("spec_"))
    directed = [m for m in _DIRECTED if m in present]
    soup = ["merge_soup"] if "merge_soup" in present else []
    models = ["base"] + specialists + soup + directed
    labels = {"base": "base", **{s: s.replace("spec_", "spec:") for s in specialists},
              "merge_soup": "soup (uniform)", **_DIRECTED}
    colors = {"base": "#7f7f7f", **{s: "#1f77b4" for s in specialists},
              "merge_soup": "#ff7f0e", **{m: "#2ca02c" for m in directed}}

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

    # Panel A: per-family accuracy.
    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 in directed else 0.65)
    ax.set_xticks(x); ax.set_xticklabels(_FAMS)
    ax.set(ylabel="accuracy", title="Per-family: directed offspring (green), selected on the\n"
           "verifier, vs the single uniform soup (orange) and the parents")
    ax.legend(frameon=False, fontsize=8, ncol=2)

    # Panel B: overall vs worst-family, with the best-specialist ceiling.
    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)
    if specialists:
        ceil = max(_acc(df, s, "overall") for s in specialists)
        ax.axhline(ceil, ls=":", c="#1f77b4", lw=1, alpha=0.7)      # best-parent ceiling
    ax.set(ylabel="accuracy", title="Overall (solid) vs worst-family (hatched):\n"
           "directed offspring vs soup vs the best parent (dotted)")
    ax.legend(frameon=False, fontsize=9)

    best_dir = max([_acc(df, m, "overall") for m in directed], default=float("nan"))
    soup_o = _acc(df, "merge_soup", "overall")
    best_spec = max([_acc(df, s, "overall") for s in specialists], default=float("nan"))
    if best_dir > soup_o + 0.005:
        verdict = f"directed {best_dir:.2f} > soup {soup_o:.2f} overall"
    elif best_dir > soup_o - 0.005:
        verdict = f"directed {best_dir:.2f} ≈ soup {soup_o:.2f} overall"
    else:
        verdict = f"directed {best_dir:.2f} < soup {soup_o:.2f} overall"
    verdict += f" (best parent {best_spec:.2f})"
    fig.suptitle(f"llm_directed — directed sex (breed offspring + select on the verifier): {verdict} "
                 f"({cfg['base_model'].split('/')[-1]})", y=1.0, fontsize=12)
    fig.tight_layout()
    savefig(fig, results_dir, "llm_directed")


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