"""AHD-CMA ablation study on a small set of tasks.

Per CLAUDE.md §5 Phase 10, runs five variants across 3 representative
tasks x 5 seeds (= 75 runs):

* ``no_chaotic``  -- uniform-random init instead of Tent map.
* ``no_adaptive`` -- fixed thresholds (no running-mean update).
* ``no_ruggedness`` -- entropy-only controller (ruggedness ignored).
* ``no_doa_de``  -- DOA branch replaced with Differential Evolution.
* ``full``       -- the released AHD-CMA (control).

We construct each variant by patching the AHD-CMA config with a
suitable override and routing through the same ``cli.run_task``
machinery.

This script intentionally targets the CEC-2022 benchmark functions
(no GPU, fast) so the ablation can run alongside long LoRA jobs
without contention. A separate CLAUDE.md follow-up will swap in
LoRA tasks once Round 1 is approved.
"""

from __future__ import annotations

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

import numpy as np
import yaml
from tqdm import tqdm

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


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


def _make_variant(base: dict[str, Any], variant: str) -> dict[str, Any]:
    """Ablation variants that match the actually-implemented controller.

    * ``full``: the released AHD-CMA -- chaotic init, probe-then-lock,
      stagnation-driven hybrid bursts, elitism.
    * ``no_chaotic``: replace the Tent-map init with IID uniform
      (handled via the _AHDCMAUniformInit subclass, see below).
    * ``no_probe``: probe window collapsed to 0 -- the algorithm jumps
      straight into stagnation-burst mode from generation 1, never
      committing to pure CMA-ES on smooth landscapes.
    * ``no_burst``: disable the hybrid bursts entirely -- after the
      probe the algorithm stays in pure CMA-ES forever even if it
      stagnates.
    * ``no_lock``: probe runs but the lock test is bypassed -- the
      algorithm continues to test for stagnation and enter bursts
      even on smooth landscapes.
    """
    cfg = copy.deepcopy(base)
    if variant == "no_chaotic":
        cfg.setdefault("init", {})["type"] = "uniform"
    elif variant == "no_probe":
        # Make the probe window vanishingly small so the lock test
        # cannot fire and the stagnation-burst path is reachable
        # from the first generation.
        cfg.setdefault("stagnation", {})["window"] = 1
    elif variant == "no_burst":
        # Set hybrid_burst = 0 so the algorithm never switches to
        # hybrid mode -- it runs pure CMA-ES the entire time.
        cfg.setdefault("stagnation", {})["hybrid_burst"] = 0
    elif variant == "no_lock":
        # Run probe but never lock: post-probe always falls through
        # to the stagnation-burst logic.
        cfg.setdefault("stagnation", {})["disable_lock"] = True
    elif variant == "full":
        pass
    else:
        raise ValueError(f"unknown variant {variant!r}")
    return cfg


def _make_uniform_init(cfg: dict[str, Any]) -> AHDCMA:
    """Wrap AHDCMA so init becomes uniform random when variant=no_chaotic."""

    class _AHDCMAUniformInit(AHDCMA):
        def optimize(self) -> Any:
            # Patch the tent-map call by feeding pre-uniform points.

            sp = self.search_space
            rng = np.random.default_rng(self.seed)
            X = sp.lower + rng.uniform(size=(self.n, sp.dim)) * (sp.upper - sp.lower)
            self._init_X_override = X
            # Monkey-patch tent_map_init temporarily by overriding the
            # internal start state.
            return super().optimize()

    return _AHDCMAUniformInit  # type: ignore[return-value]


def run_one(
    variant: str,
    func: str,
    dim: int,
    seed: int,
    *,
    pop: int,
    gens: int,
    output_dir: Path,
) -> dict[str, Any]:
    set_global_seed(seed)
    base = _load_ahdcma_config()
    cfg = _make_variant(base, variant)
    cfg["seed"] = seed
    cfg["population_size"] = pop
    cfg["max_generations"] = gens

    run_id = make_run_id("ahdcma_ab_" + variant, func, seed, dim=dim)
    out_root = output_dir / run_id
    out_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)

    cls: type[AHDCMA] = AHDCMA
    if variant == "no_chaotic":
        cls = _make_uniform_init(cfg)  # type: ignore[assignment]

    opt = cls(cfg, problem, sp, run_id=run_id)
    t0 = time.time()
    result = opt.optimize()
    summary = {
        "run_id": run_id,
        "variant": variant,
        "func": func,
        "dim": dim,
        "seed": seed,
        "best_f": float(result.best_f),
        "wall_time": time.time() - t0,
    }
    (out_root / "result.json").write_text(json.dumps(summary, indent=2))
    return summary


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument(
        "--funcs",
        nargs="+",
        default=["F1_zakharov", "F5_levy", "F9_composition1"],
    )
    p.add_argument("--dim", type=int, default=10)
    p.add_argument("--seeds", nargs="+", type=int, default=list(range(5)))
    p.add_argument(
        "--variants",
        nargs="+",
        default=["full", "no_chaotic", "no_adaptive", "no_ruggedness", "no_doa_de"],
    )
    p.add_argument("--pop", type=int, default=30)
    p.add_argument("--gens", type=int, default=100)
    p.add_argument("--output", default="outputs/runs/ablation")
    args = p.parse_args()

    out_root = Path(args.output)
    out_root.mkdir(parents=True, exist_ok=True)
    combos = list(product(args.variants, args.funcs, args.seeds))
    print(f"Ablation sweep: {len(combos)} runs")
    for variant, func, seed in tqdm(combos, desc="ablation"):
        existing = list(
            out_root.glob(f"ahdcma_ab_{variant}_{func}_d{args.dim}_seed{seed}_*/result.json")
        )
        if existing:
            continue
        run_one(
            variant,
            func,
            args.dim,
            seed,
            pop=args.pop,
            gens=args.gens,
            output_dir=out_root,
        )


if __name__ == "__main__":
    main()