"""CLI entry point for a single CEC-2022 run.

Usage::

    python -m ahdcma.cli.run_benchmark \\
        --algo ahdcma --func F5_rastrigin --dim 10 --seed 0 \\
        --output outputs/runs/cec2022/

Writes ``{output}/{run_id}/result.json`` with the final best fitness,
config snapshot, history summary and wall time.
"""

from __future__ import annotations

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

import numpy as np
import yaml

from ahdcma.algorithms.ahd_cma import AHDCMA
from ahdcma.algorithms.base import Optimizer, SearchSpace
from ahdcma.algorithms.bohb import BOHB
from ahdcma.algorithms.cmaes_restart import BIPOPCMAES, IPOPCMAES
from ahdcma.algorithms.cmaes_wrapper import CMAES
from ahdcma.algorithms.doa import DOA
from ahdcma.algorithms.gego import GEGO
from ahdcma.algorithms.grid_search import GridSearch
from ahdcma.algorithms.gwo import GWO
from ahdcma.algorithms.hyperband import Hyperband
from ahdcma.algorithms.ihaho import IHAHO
from ahdcma.algorithms.optuna_baseline import OptunaTPE
from ahdcma.algorithms.pso import PSO
from ahdcma.algorithms.random_search import RandomSearch
from ahdcma.algorithms.scso import SCSO
from ahdcma.algorithms.woa import WOA
from ahdcma.fitness.cec2022 import ALL_NAMES, get_problem
from ahdcma.utils.logging import make_run_id, setup_logging
from ahdcma.utils.seed import set_global_seed

ALGO_REGISTRY: dict[str, type[Optimizer]] = {
    "ahdcma": AHDCMA,
    "doa": DOA,
    "cmaes": CMAES,
    "ipop": IPOPCMAES,
    "bipop": BIPOPCMAES,
    "pso": PSO,
    "gwo": GWO,
    "woa": WOA,
    "scso": SCSO,
    "optuna": OptunaTPE,
    "random": RandomSearch,
    "grid": GridSearch,
    "hyperband": Hyperband,
    "bohb": BOHB,
    "gego": GEGO,
    "ihaho": IHAHO,
}


def _load_algo_config(algo: str, *, seed: int, max_evals: int, pop: int) -> dict[str, Any]:
    cfg_path = Path(__file__).resolve().parents[3] / "configs" / "algo" / f"{algo}.yaml"
    with cfg_path.open() as f:
        cfg = yaml.safe_load(f)
    # Override the run-time controllable knobs
    cfg["seed"] = seed
    cfg["population_size"] = pop
    cfg["max_generations"] = max(1, max_evals // pop)
    return dict(cfg)


def run_single(
    algo: str,
    func: str,
    dim: int,
    seed: int,
    *,
    max_evals: int = 30000,
    pop: int = 30,
    output_dir: Path | str = "outputs/runs/cec2022",
) -> dict[str, Any]:
    """Execute a single algo/func/dim/seed run and write result.json."""
    if algo not in ALGO_REGISTRY:
        raise KeyError(f"unknown algorithm {algo!r}; choose from {sorted(ALGO_REGISTRY)}")
    if func not in ALL_NAMES:
        raise KeyError(f"unknown CEC-2022 function {func!r}; choose from {ALL_NAMES}")

    set_global_seed(seed)
    run_id = make_run_id(algo, func, seed, dim=dim)
    out_root = Path(output_dir) / run_id
    out_root.mkdir(parents=True, exist_ok=True)

    log_dir = out_root / "logs"
    setup_logging(run_id, log_dir=log_dir)

    problem = get_problem(func, dim)
    sp = SearchSpace(
        dim=dim,
        lower=np.full(dim, problem.bound[0], dtype=np.float64),
        upper=np.full(dim, problem.bound[1], dtype=np.float64),
    )
    cfg = _load_algo_config(algo, seed=seed, max_evals=max_evals, pop=pop)
    cls = ALGO_REGISTRY[algo]
    opt = cls(cfg, problem, sp, run_id=run_id)

    t0 = time.time()
    result = opt.optimize()
    wall = time.time() - t0

    summary = {
        "run_id": run_id,
        "algo": algo,
        "func": func,
        "dim": dim,
        "seed": seed,
        "best_f": float(result.best_f),
        "best_x": result.best_x.tolist(),
        "wall_time": wall,
        "n_generations": len(result.history),
        "config": cfg,
        "best_fitness_curve": list(result.history.best_fitness),
        "mode_curve": list(result.history.mode_per_gen),
    }
    (out_root / "result.json").write_text(json.dumps(summary, indent=2))
    return summary


def main() -> None:
    p = argparse.ArgumentParser(description="Run a single CEC-2022 evaluation.")
    p.add_argument("--algo", required=True, choices=sorted(ALGO_REGISTRY))
    p.add_argument("--func", required=True, choices=ALL_NAMES)
    p.add_argument("--dim", type=int, default=10)
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--max-evals", type=int, default=30000)
    p.add_argument("--pop", type=int, default=30)
    p.add_argument("--output", default="outputs/runs/cec2022")
    args = p.parse_args()
    summary = run_single(
        args.algo,
        args.func,
        args.dim,
        args.seed,
        max_evals=args.max_evals,
        pop=args.pop,
        output_dir=args.output,
    )
    print(json.dumps({"run_id": summary["run_id"], "best_f": summary["best_f"]}))


if __name__ == "__main__":
    main()