scripts/generate_paper_artifacts.py

"""Generate every CEC-2022-derived paper figure from sweep result files.

This is the single committed generator for the paper's data-driven
figures, so the manuscript is fully reproducible from the raw
``result.json`` corpus. It reuses the viz helpers in
``ahdcma.viz`` where they exist and adds the bar/heatmap renderers that
were previously produced ad-hoc.

Inputs
------
* ``--cec-root``      CEC-2022 sweep results (default
  ``outputs/runs/cec2022_rotated``).
* ``--ablation-root`` ablation sweep results (default
  ``outputs/runs/ablation_rotated``).
* ``--lora-root``     LoRA sweep results (default ``outputs/runs/lora``)
  -- optional; skipped if absent (LoRA is unaffected by CEC rotations).
* ``--out``           figure output dir (default
  ``paper/soft_computing/figures``).

Outputs (PDF + PNG, 300 dpi) matching the manuscript's figure names:
  friedman_d10, friedman_d20, per_func_rank, runtime, wilcoxon_d20,
  conv_d20_F1, conv_d20_F5, conv_d20_F9, ablation, mode_distribution,
  sensitivity (only if a sensitivity sweep is present),
  lora_boxplots, lora_convergence (only if LoRA results present).

Usage
-----
    python scripts/generate_paper_artifacts.py \
        --cec-root outputs/runs/cec2022_rotated \
        --ablation-root outputs/runs/ablation_rotated \
        --out paper/soft_computing/figures
"""

from __future__ import annotations

import argparse
import json
from collections import defaultdict
from pathlib import Path
from typing import Any

import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import rankdata

from ahdcma.fitness.cec2022 import ALL_NAMES
from ahdcma.stats.tests import friedman_test
from ahdcma.viz.convergence import plot_convergence

OURS = "ahdcma"
DPI = 300
# Representative functions for the convergence panels (unimodal,
# multimodal, composition) using the official identities.
HEADLINE = ["F1_zakharov", "F5_levy", "F9_composition1"]

# Human-readable algorithm labels for figure legends/axes.
_DISP = {
    "ahdcma": "AHD-CMA", "cmaes": "CMA-ES", "ipop": "IPOP-CMA-ES",
    "bipop": "BIPOP-CMA-ES", "doa": "DOA", "gego": "GEGO", "ihaho": "I-HAHO",
    "pso": "PSO", "gwo": "GWO", "woa": "WOA", "scso": "SCSO",
    "optuna": "Optuna-TPE", "bohb": "BOHB", "hyperband": "Hyperband",
    "random": "Random", "grid": "Grid",
}


def _disp(algo: str) -> str:
    """Map an internal algo key to its publication label."""
    return _DISP.get(algo, algo)


# --------------------------------------------------------------------------
# Result loading / pivoting (mirrors ahdcma.cli.make_figures helpers)
# --------------------------------------------------------------------------
def load_results(root: Path) -> list[dict[str, Any]]:
    out: list[dict[str, Any]] = []
    if not root.exists():
        return out
    for jf in root.rglob("result.json"):
        try:
            out.append(json.loads(jf.read_text()))
        except json.JSONDecodeError:
            continue
    return out


def pivot_best_f(
    results: list[dict[str, Any]],
) -> dict[int, dict[tuple[str, str], list[float]]]:
    by_dim: dict[int, dict[tuple[str, str], list[float]]] = defaultdict(
        lambda: defaultdict(list)
    )
    for r in results:
        by_dim[int(r["dim"])][r["algo"], r["func"]].append(float(r["best_f"]))
    return by_dim


def mean_rank_table(
    by_func_algo: dict[str, dict[str, float]],
) -> dict[str, float]:
    """Mean Friedman rank per algorithm over functions (lower = better).

    ``by_func_algo[func][algo]`` is that algo's median best_f on the
    function. Ranks are computed per function (1 = best) then averaged.
    """
    algos = sorted({a for d in by_func_algo.values() for a in d})
    per_func_ranks: dict[str, dict[str, float]] = {}
    for func, scores in by_func_algo.items():
        vals = np.array([scores.get(a, np.inf) for a in algos], dtype=np.float64)
        ranks = rankdata(vals, method="average")  # 1 = smallest = best
        per_func_ranks[func] = dict(zip(algos, ranks, strict=True))
    mean_rank = {
        a: float(np.mean([per_func_ranks[f][a] for f in per_func_ranks])) for a in algos
    }
    return mean_rank


def median_by_func_algo(
    cell: dict[tuple[str, str], list[float]],
) -> dict[str, dict[str, float]]:
    out: dict[str, dict[str, float]] = defaultdict(dict)
    for (algo, func), vals in cell.items():
        if vals:
            out[func][algo] = float(np.median(vals))
    return out


# --------------------------------------------------------------------------
# Figure renderers
# --------------------------------------------------------------------------
def fig_friedman_bar(mean_rank: dict[str, float], dim: int, out: Path) -> None:
    order = sorted(mean_rank, key=lambda a: mean_rank[a])
    ranks = [mean_rank[a] for a in order]
    colors = ["tab:orange" if a == OURS else "tab:blue" for a in order]
    fig, ax = plt.subplots(figsize=(6.5, 5.0))
    ax.barh(range(len(order)), ranks, color=colors)
    ax.set_yticks(range(len(order)))
    ax.set_yticklabels([_disp(a) for a in order])
    ax.invert_yaxis()
    for i, v in enumerate(ranks):
        ax.text(v + 0.05, i, f"{v:.2f}", va="center", fontsize=8)
    ax.set_xlabel("Mean Friedman rank (lower = better)")
    ax.set_title(f"CEC-2022 dim={dim}: 12 functions x 30 seeds, {len(order)} algorithms")
    fig.tight_layout()
    _save(fig, out / f"friedman_d{dim}")


def fig_per_func_rank(
    by_func_algo: dict[str, dict[str, float]],
    mean_rank: dict[str, float],
    out: Path,
    *,
    top_k: int = 5,
) -> None:
    top = sorted(mean_rank, key=lambda a: mean_rank[a])[:top_k]
    funcs = [n for n in ALL_NAMES if n in by_func_algo]
    algos_all = sorted({a for d in by_func_algo.values() for a in d})
    fig, ax = plt.subplots(figsize=(9.0, 4.5))
    width = 0.8 / len(top)
    # Distinct, non-orange palette for the baselines so AHD-CMA's orange
    # stands out unambiguously.
    base_colors = ["tab:blue", "tab:green", "tab:purple", "tab:brown",
                   "tab:cyan", "tab:olive", "tab:pink", "tab:gray"]
    bi = 0
    for i, algo in enumerate(top):
        rks = []
        for f in funcs:
            vals = np.array(
                [by_func_algo[f].get(a, np.inf) for a in algos_all], dtype=np.float64
            )
            r = rankdata(vals, method="average")
            rks.append(r[algos_all.index(algo)])
        xs = np.arange(len(funcs)) + i * width
        # Highlight AHD-CMA in orange to match the other figures.
        if algo == OURS:
            color = "tab:orange"
        else:
            color = base_colors[bi % len(base_colors)]
            bi += 1
        ax.bar(xs, rks, width=width, label=_disp(algo), color=color)
    ax.set_xticks(np.arange(len(funcs)) + 0.4)
    ax.set_xticklabels([f"F{i + 1}" for i in range(len(funcs))])
    ax.set_ylabel("Friedman rank (1 = best)")
    ax.set_title(f"Per-function ranking, official CEC-2022 (top-{top_k} algorithms)")
    ax.legend(ncol=len(top), fontsize=8, loc="upper center")
    fig.tight_layout()
    _save(fig, out / "per_func_rank")


def fig_runtime(results: list[dict[str, Any]], dim: int, out: Path) -> None:
    by_algo: dict[str, list[float]] = defaultdict(list)
    for r in results:
        if int(r["dim"]) == dim and r.get("wall_time"):
            by_algo[r["algo"]].append(float(r["wall_time"]))
    order = sorted(by_algo, key=lambda a: np.median(by_algo[a]))
    meds = [float(np.median(by_algo[a])) for a in order]
    p95 = [float(np.percentile(by_algo[a], 95)) for a in order]
    colors = ["tab:orange" if a == OURS else "tab:blue" for a in order]
    fig, ax = plt.subplots(figsize=(7.0, 4.5))
    ax.bar(range(len(order)), meds, color=colors,
           yerr=[np.zeros(len(order)), np.array(p95) - np.array(meds)], capsize=3)
    ax.set_yscale("log")
    ax.set_xticks(range(len(order)))
    ax.set_xticklabels([a.upper() if a == OURS else a for a in order],
                       rotation=45, ha="right")
    ax.set_ylabel("Per-run wall time (s, log scale)")
    ax.set_title(f"Computational cost, CEC-2022 dim={dim} (median + 95th pct)")
    fig.tight_layout()
    _save(fig, out / "runtime")


def fig_wilcoxon(cell: dict[tuple[str, str], list[float]], dim: int, out: Path) -> None:
    """AHD-CMA vs each baseline: a horizontal -log10(p) bar chart with
    bars coloured by which algorithm wins (per-function median best_f).
    Clearer than a single-column heatmap.
    """
    from scipy.stats import wilcoxon as _wilcoxon

    by_func_algo = median_by_func_algo(cell)
    funcs = [n for n in ALL_NAMES if n in by_func_algo]
    algos = [a for a in sorted({a for d in by_func_algo.values() for a in d}) if a != OURS]
    ours = np.array([by_func_algo[f].get(OURS, np.nan) for f in funcs])

    rows = []  # (algo, neglog10p, ahd_wins)
    for a in algos:
        bb = np.array([by_func_algo[f].get(a, np.nan) for f in funcs])
        try:
            _, p = _wilcoxon(ours, bb)
        except ValueError:
            p = 1.0
        wins = int(np.sum(ours < bb))
        rows.append((a, -np.log10(max(p, 1e-12)), wins > len(funcs) / 2, p))
    # Order: AHD-CMA's strongest wins at the top, losses at the bottom.
    rows.sort(key=lambda r: (r[2], r[1]))
    names = [_disp(r[0]) for r in rows]
    vals = [r[1] for r in rows]
    colors = ["tab:blue" if r[2] else "tab:red" for r in rows]

    fig, ax = plt.subplots(figsize=(7.0, 5.0))
    ax.barh(range(len(rows)), vals, color=colors)
    ax.set_yticks(range(len(rows)))
    ax.set_yticklabels(names, fontsize=9)
    ax.axvline(-np.log10(0.05), color="black", ls="--", lw=1,
               label=r"$p=0.05$ significance")
    for i, r in enumerate(rows):
        ax.text(r[1] + 0.05, i, f"$p$={r[3]:.3f}", va="center", fontsize=7)
    ax.set_xlabel(r"$-\log_{10} p$ (Wilcoxon, per-function medians)")
    ax.set_title(f"AHD-CMA vs. each baseline, official CEC-2022 dim={dim}")
    # Legend explaining colour direction.
    from matplotlib.patches import Patch
    ax.legend(handles=[
        Patch(color="tab:blue", label="AHD-CMA wins (majority of functions)"),
        Patch(color="tab:red", label="baseline wins"),
        plt.Line2D([0], [0], color="black", ls="--", label=r"$p=0.05$"),
    ], fontsize=8, loc="lower right")
    ax.grid(True, axis="x", alpha=0.3)
    fig.tight_layout()
    _save(fig, out / f"wilcoxon_d{dim}")


def fig_convergence(results: list[dict[str, Any]], dim: int, out: Path) -> None:
    for idx, func in enumerate(HEADLINE, start=1):
        curves: dict[str, list[list[float]]] = defaultdict(list)
        for r in results:
            if int(r["dim"]) == dim and r["func"] == func:
                curves[r["algo"]].append(list(r["best_fitness_curve"]))
        # Only show a few headline algos to keep the plot legible.
        keep = {OURS, "cmaes", "ipop", "bipop", "doa"}
        sub = {k: v for k, v in curves.items() if k in keep and v}
        if not sub:
            continue
        plot_convergence(
            sub, out / f"conv_d{dim}_F{idx}",
            title=f"F{idx} ({func})", ylabel="Error (best-so-far)", log_y=True,
        )


def fig_ablation(ab_results: list[dict[str, Any]], out: Path) -> None:
    """log10(variant median / full median) heatmap over functions."""
    by_var_func: dict[str, dict[str, list[float]]] = defaultdict(lambda: defaultdict(list))
    for r in ab_results:
        var = r.get("variant") or r.get("algo", "?")
        by_var_func[var][r["func"]].append(float(r["best_f"]))
    if "full" not in by_var_func:
        print("[ablation] no 'full' control found; skipping")
        return
    funcs = [n for n in ALL_NAMES if n in by_var_func["full"]]
    variants = [v for v in by_var_func if v != "full"]
    mat = np.full((len(variants), len(funcs)), np.nan)
    for i, v in enumerate(variants):
        for j, f in enumerate(funcs):
            fv = by_var_func["full"].get(f)
            vv = by_var_func[v].get(f)
            if fv and vv:
                fm, vm = np.median(fv), np.median(vv)
                if fm > 0 and vm > 0:
                    mat[i, j] = np.log10(vm / fm)
    fig, ax = plt.subplots(figsize=(8.0, 3.5))
    vmax = np.nanmax(np.abs(mat)) if np.isfinite(mat).any() else 1.0
    im = ax.imshow(mat, cmap="coolwarm", aspect="auto", vmin=-vmax, vmax=vmax)
    ax.set_xticks(range(len(funcs)))
    ax.set_xticklabels([f"F{i + 1}" for i in range(len(funcs))])
    ax.set_yticks(range(len(variants)))
    ax.set_yticklabels(variants)
    fig.colorbar(im, ax=ax, label="log10(variant / full)")
    ax.set_title("Ablation: log10 median-fitness ratio vs full (blue = variant better)")
    fig.tight_layout()
    _save(fig, out / "ablation")


def fig_mode_distribution(results: list[dict[str, Any]], dim: int, out: Path) -> None:
    """Controller behaviour per function: when the first HYBRID burst
    fires (for runs that burst) and how many runs instead LOCK into pure
    CMA-ES (the graceful-fallback path). Runs that never burst are not
    plotted as a spurious "generation 1000" outlier --- they are counted
    separately as locked runs.
    """
    burst_gen: dict[str, list[int]] = defaultdict(list)
    n_locked: dict[str, int] = defaultdict(int)
    n_total: dict[str, int] = defaultdict(int)
    probe_w = 8
    for r in results:
        if r["algo"] != OURS or int(r["dim"]) != dim:
            continue
        n_total[r["func"]] += 1
        modes = r.get("mode_curve", [])
        gen = next(
            (i for i, m in enumerate(modes) if str(m).lower() in {"hybrid", "burst"}),
            None,
        )
        if gen is None:
            n_locked[r["func"]] += 1
        else:
            burst_gen[r["func"]].append(gen)
    funcs = [n for n in ALL_NAMES if n in n_total]
    if not funcs:
        print("[mode_dist] no AHD-CMA mode curves found; skipping")
        return
    labels = [f"F{i + 1}" for i in range(len(funcs))]

    fig, (ax1, ax2) = plt.subplots(
        1, 2, figsize=(9.5, 3.6), gridspec_kw={"width_ratios": [3, 2]}
    )

    # Panel (a): first-burst generation for runs that DO burst (zoomed).
    positions = list(range(len(funcs)))
    burst_data = [burst_gen[f] for f in funcs]
    has_data = [i for i, d in enumerate(burst_data) if d]
    if has_data:
        ax1.boxplot(
            [burst_data[i] for i in has_data],
            positions=[positions[i] for i in has_data],
            widths=0.6,
            patch_artist=True,
            boxprops={"facecolor": "tab:blue", "alpha": 0.6},
            medianprops={"color": "black"},
            flierprops={"marker": ".", "markersize": 3},
        )
    ax1.axhline(probe_w, color="red", ls="--", lw=1.2, label=f"probe window $W={probe_w}$")
    all_bursts = [g for f in funcs for g in burst_gen[f]]
    if all_bursts:
        ax1.set_ylim(0, max(probe_w + 4, max(all_bursts) + 2))
    ax1.set_xticks(positions)
    ax1.set_xticklabels(labels, fontsize=8)
    ax1.set_ylabel("First HYBRID-burst generation")
    ax1.set_title("(a) When the burst fires (burst runs only)")
    ax1.legend(fontsize=8, loc="upper left")
    ax1.grid(True, axis="y", alpha=0.3)

    # Panel (b): fraction of seeds that lock vs burst, per function.
    lock_frac = [n_locked[f] / n_total[f] for f in funcs]
    burst_frac = [1 - lf for lf in lock_frac]
    ax2.bar(positions, burst_frac, color="tab:blue", label="burst (stagnation)")
    ax2.bar(positions, lock_frac, bottom=burst_frac, color="tab:orange",
            label="lock (graceful fallback)")
    ax2.set_xticks(positions)
    ax2.set_xticklabels(labels, fontsize=8)
    ax2.set_ylim(0, 1)
    ax2.set_ylabel("Fraction of 30 seeds")
    ax2.set_title("(b) Burst vs. lock outcome")
    ax2.legend(fontsize=8, loc="lower right")

    fig.suptitle(f"AHD-CMA controller behaviour on official CEC-2022 dim={dim}",
                 fontsize=11)
    fig.tight_layout(rect=(0, 0, 1, 0.96))
    _save(fig, out / "mode_distribution")


def fig_lora(lora_root: Path, out: Path) -> None:
    results = load_results(lora_root)
    if not results:
        print(f"[lora] no results under {lora_root}; skipping LoRA figures")
        return
    by_task_algo: dict[str, dict[str, list[float]]] = defaultdict(lambda: defaultdict(list))
    curves: dict[str, dict[str, list[list[float]]]] = defaultdict(lambda: defaultdict(list))
    for r in results:
        task = r.get("task") or r.get("func", "?")
        acc = -float(r["best_f"]) if float(r["best_f"]) < 0 else float(r["best_f"])
        by_task_algo[task][r["algo"]].append(acc)
        if r.get("best_fitness_curve"):
            curves[task][r["algo"]].append([-c for c in r["best_fitness_curve"]])
    # Box plots: one group per task.
    tasks = sorted(by_task_algo)
    algos = sorted({a for d in by_task_algo.values() for a in d})
    fig, axes = plt.subplots(1, len(tasks), figsize=(4.0 * len(tasks), 4.0), squeeze=False)
    for ax, task in zip(axes[0], tasks, strict=True):
        ax.boxplot([by_task_algo[task].get(a, []) for a in algos], tick_labels=algos)
        ax.set_title(task)
        ax.set_ylabel("Validation accuracy")
    fig.tight_layout()
    _save(fig, out / "lora_boxplots")
    # Convergence per task (median + IQR), first task only into one fig grid.
    for task in tasks:
        sub = {a: c for a, c in curves[task].items() if c}
        if sub:
            plot_convergence(
                sub, out / f"lora_convergence_{task}",
                title=f"LoRA convergence: {task}", ylabel="Best-so-far accuracy",
            )


def _save(fig: Any, stem: Path) -> None:
    stem.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(stem.with_suffix(".pdf"), dpi=DPI)
    fig.savefig(stem.with_suffix(".png"), dpi=DPI)
    plt.close(fig)
    print(f"  wrote {stem.name}.{{pdf,png}}")


# --------------------------------------------------------------------------
def main() -> None:
    p = argparse.ArgumentParser(description="Generate CEC-derived paper figures")
    p.add_argument("--cec-root", default="outputs/runs/cec2022_rotated")
    p.add_argument("--ablation-root", default="outputs/runs/ablation_rotated")
    p.add_argument("--lora-root", default="outputs/runs/lora")
    p.add_argument("--out", default="paper/soft_computing/figures")
    p.add_argument("--dims", nargs="+", type=int, default=[10, 20])
    args = p.parse_args()

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)

    cec = load_results(Path(args.cec_root))
    print(f"loaded {len(cec)} CEC result files from {args.cec_root}")
    if cec:
        by_dim = pivot_best_f(cec)
        for dim in args.dims:
            if dim not in by_dim:
                print(f"[dim {dim}] no results; skipping")
                continue
            cell = by_dim[dim]
            bfa = median_by_func_algo(cell)
            mr = mean_rank_table(bfa)
            stat, pval = friedman_test(
                {
                    a: np.array(
                        [bfa[f].get(a, np.nan) for f in bfa], dtype=np.float64
                    )
                    for a in sorted({x for d in bfa.values() for x in d})
                }
            )
            print(f"[dim {dim}] Friedman chi2={stat:.2f} p={pval:.2e}; "
                  f"AHD-CMA mean rank={mr.get(OURS, float('nan')):.3f}")
            fig_friedman_bar(mr, dim, out)
            fig_per_func_rank(bfa, mr, out)
            fig_runtime(cec, dim, out)
            fig_wilcoxon(cell, dim, out)
            fig_convergence(cec, dim, out)
            fig_mode_distribution(cec, dim, out)

    ab = load_results(Path(args.ablation_root))
    print(f"loaded {len(ab)} ablation result files from {args.ablation_root}")
    if ab:
        fig_ablation(ab, out)

    fig_lora(Path(args.lora_root), out)
    print("done.")


if __name__ == "__main__":
    main()