"""Thin wrapper around :mod:`cma` (pycma) implementing our :class:`Optimizer` API.

We keep the wrapper deliberately minimal: pycma manages step-size,
covariance, and the population schedule, while we record the per-
generation history, clip into the box, and expose the same interface as
the other optimizers in this package.

Notes
-----
- pycma operates without explicit bounds by default. We pass them via
  ``BoundaryHandler`` and clip the asked points before evaluation, which
  is robust on the unit cube used throughout the project.
- pycma's "generation" corresponds to one ``ask``/``tell`` round.
- pycma needs ``d >= 2``; raise a clear error on ``d == 1``.
"""

from __future__ import annotations

import time
from typing import Any

import cma
import numpy as np
from numpy.typing import NDArray

from ahdcma.algorithms.base import (
    FitnessFn,
    OptimizationResult,
    Optimizer,
    SearchSpace,
)


class CMAES(Optimizer):
    """CMA-ES via pycma."""

    def __init__(
        self,
        config: dict[str, Any],
        fitness_fn: FitnessFn,
        search_space: SearchSpace,
        *,
        run_id: str = "anonymous",
    ) -> None:
        super().__init__(config, fitness_fn, search_space, run_id=run_id)
        if search_space.dim < 2:
            raise ValueError("CMA-ES requires dim >= 2; got dim=1")
        self.population_size: int = int(config["population_size"])
        self.t_max: int = int(config["max_generations"])
        self.sigma: float = float(config.get("initial_sigma", 0.3))
        self.seed: int = int(config["seed"])
        if self.population_size < 4:
            raise ValueError("CMA-ES population_size must be >= 4")
        if self.t_max < 1:
            raise ValueError("max_generations must be >= 1")

    def _evaluate(self, X: NDArray[np.float64]) -> NDArray[np.float64]:
        return np.array([float(self.fitness_fn(x)) for x in X], dtype=np.float64)

    def optimize(self) -> OptimizationResult:
        sp = self.search_space
        x0 = 0.5 * (sp.lower + sp.upper)
        opts = {
            "popsize": self.population_size,
            "bounds": [sp.lower.tolist(), sp.upper.tolist()],
            "seed": self.seed if self.seed > 0 else 1,  # pycma forbids seed==0
            "verbose": -9,
            "verb_log": 0,
            "verb_disp": 0,
            "maxiter": self.t_max,
        }
        es = cma.CMAEvolutionStrategy(x0.tolist(), self.sigma, opts)

        wall_t0 = time.time()
        gen = 0
        while not es.stop() and gen < self.t_max:
            asked = es.ask()
            X = sp.clip(np.asarray(asked, dtype=np.float64))
            F = self._evaluate(X)
            es.tell([X[i].tolist() for i in range(X.shape[0])], F.tolist())
            self._history.append(gen, X, F, mode="exploit")
            gen += 1

        best = es.best.x
        best_f = float(es.best.f) if es.best.f is not None else float("inf")
        return OptimizationResult(
            best_x=np.asarray(best, dtype=np.float64),
            best_f=best_f,
            history=self._history,
            config_snapshot=self.config,
            run_id=self.run_id,
            wall_time=time.time() - wall_t0,
        )