"""Sensitivity sweep of AHD-CMA's three controller hyperparameters.

Sweeps each of the probe window ``W`` (``stagnation.window``), the burst
length ``B`` (``stagnation.hybrid_burst``), and the CMA-ES top fraction
``k_top/N`` around their defaults while holding the others fixed, on a
few representative official CEC-2022 functions. Produces the data the
manuscript's sensitivity figure consumes.

This is a cheap sweep (few functions x few values x few seeds) and is
not a §10 Trigger 1 job. Output: ``outputs/runs/sensitivity/`` with one
``result.json`` per run; the figure is rendered by
``generate_paper_artifacts.py`` (or the inline ``--plot`` flag here).

Usage
-----
    python scripts/run_sensitivity.py --seeds 0 1 2 3 4 --plot
"""

from __future__ import annotations

import argparse
import copy
import json
import time
from pathlib import Path
from typing import Any

import matplotlib.pyplot as plt
import numpy as np
import yaml

from ahdcma.algorithms.ahd_cma import AHDCMA
from ahdcma.algorithms.base import SearchSpace
from ahdcma.fitness.cec2022 import get_problem
from ahdcma.utils.logging import make_run_id
from ahdcma.utils.seed import set_global_seed

PARAMS = {
    "W": [4, 6, 8, 12, 16],  # stagnation.window (probe window)
    "B": [1, 2, 3, 5, 8],  # stagnation.hybrid_burst (burst length)
    "ktop": [0.1, 0.2, 0.3, 0.5, 0.7],  # CMA-ES top fraction in hybrid mode
}
DEFAULTS = {"W": 8, "B": 3, "ktop": 0.3}


def _load_cfg() -> dict[str, Any]:
    path = Path(__file__).resolve().parents[1] / "configs" / "algo" / "ahdcma.yaml"
    with path.open() as f:
        return dict(yaml.safe_load(f))


def _make_ktop_class(fraction: float) -> type[AHDCMA]:
    """Subclass AHDCMA overriding the (hardcoded) 0.3 hybrid top fraction."""

    class _AHDCMAKtop(AHDCMA):
        def __init__(self, *a: Any, **k: Any) -> None:
            super().__init__(*a, **k)
            self._k_hybrid = max(2, int(round(fraction * self.n)))

    return _AHDCMAKtop


def run_one(
    param: str,
    value: float,
    func: str,
    dim: int,
    seed: int,
    *,
    pop: int,
    gens: int,
    out_dir: Path,
) -> dict[str, Any]:
    set_global_seed(seed)
    cfg = copy.deepcopy(_load_cfg())
    cfg["seed"] = seed
    cfg["population_size"] = pop
    cfg["max_generations"] = gens

    cls: type[AHDCMA] = AHDCMA
    if param == "W":
        cfg.setdefault("stagnation", {})["window"] = int(value)
    elif param == "B":
        cfg.setdefault("stagnation", {})["hybrid_burst"] = int(value)
    elif param == "ktop":
        cls = _make_ktop_class(float(value))
    else:
        raise ValueError(f"unknown param {param!r}")

    run_id = make_run_id(f"sens_{param}_{value}", func, seed, dim=dim)
    root = out_dir / run_id
    root.mkdir(parents=True, exist_ok=True)
    sp = SearchSpace(
        dim=dim,
        lower=np.full(dim, -100.0, dtype=np.float64),
        upper=np.full(dim, 100.0, dtype=np.float64),
    )
    problem = get_problem(func, dim)
    opt = cls(cfg, problem, sp, run_id=run_id)
    t0 = time.time()
    res = opt.optimize()
    summary = {
        "run_id": run_id,
        "param": param,
        "value": value,
        "func": func,
        "dim": dim,
        "seed": seed,
        "best_f": float(res.best_f),
        "wall_time": time.time() - t0,
    }
    (root / "result.json").write_text(json.dumps(summary, indent=2))
    return summary


def plot(out_dir: Path, fig_out: Path) -> None:
    rows = [
        json.loads(j.read_text()) for j in out_dir.rglob("result.json")
    ]
    fig, axes = plt.subplots(1, 3, figsize=(12.0, 4.0))
    for ax, param in zip(axes, ["W", "B", "ktop"], strict=True):
        sub = [r for r in rows if r["param"] == param]
        vals = sorted({r["value"] for r in sub})
        # Normalise median final fitness to the default value.
        med = {
            v: float(np.median([r["best_f"] for r in sub if r["value"] == v]))
            for v in vals
        }
        base = med.get(DEFAULTS[param]) or (next(iter(med.values())) if med else 1.0)
        norm = [med[v] / base if base else 1.0 for v in vals]
        ax.plot(vals, norm, "o-")
        ax.axvline(DEFAULTS[param], color="red", ls="--", alpha=0.6, label="default")
        ax.set_title({"W": "(a) probe window W", "B": "(c) burst length B",
                      "ktop": "(b) CMA-ES fraction k_top/N"}[param])
        ax.set_ylabel("Median final fitness (norm.)")
        ax.grid(True, alpha=0.3)
        ax.legend(fontsize=8)
    fig.tight_layout()
    fig_out.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(fig_out.with_suffix(".pdf"), dpi=300)
    fig.savefig(fig_out.with_suffix(".png"), dpi=300)
    plt.close(fig)
    print(f"wrote {fig_out.name}.{{pdf,png}}")


def main() -> None:
    p = argparse.ArgumentParser(description="AHD-CMA sensitivity sweep")
    p.add_argument("--funcs", nargs="+", default=["F5_levy", "F9_composition1"])
    p.add_argument("--dim", type=int, default=20)
    p.add_argument("--seeds", nargs="+", type=int, default=list(range(5)))
    p.add_argument("--pop", type=int, default=30)
    p.add_argument("--gens", type=int, default=200)
    p.add_argument("--out", default="outputs/runs/sensitivity")
    p.add_argument("--fig-out", default="paper/soft_computing/figures/sensitivity")
    p.add_argument("--plot", action="store_true", help="render the figure after running")
    args = p.parse_args()

    out_dir = Path(args.out)
    n_done = 0
    for param, values in PARAMS.items():
        for value in values:
            for func in args.funcs:
                for seed in args.seeds:
                    run_one(param, value, func, args.dim, seed,
                            pop=args.pop, gens=args.gens, out_dir=out_dir)
                    n_done += 1
    print(f"sensitivity sweep done: {n_done} runs -> {out_dir}")
    if args.plot:
        plot(out_dir, Path(args.fig_out))


if __name__ == "__main__":
    main()