tests/unit/test_optimizers.py

"""Per-Phase-2 optimizer acceptance tests.

For every standalone optimizer (DOA, CMA-ES, PSO, GWO, WOA, SCSO):

* Sphere 5D, ~2000 evals: final f < 1e-2 (lenient because metaheuristics
  are population-based and don't hit machine precision).
* Rastrigin 10D, ~2000 evals: final f < 50.
* Reproducibility: same seed -> identical history.

CMA-ES gets a tighter sphere bound because it is locally-optimal on
quadratic problems.
"""

from __future__ import annotations

from collections.abc import Callable
from typing import Any

import numpy as np
import pytest

from ahdcma.algorithms.base import OptimizationResult, Optimizer, SearchSpace
from ahdcma.algorithms.cmaes_wrapper import CMAES
from ahdcma.algorithms.doa import DOA
from ahdcma.algorithms.gwo import GWO
from ahdcma.algorithms.pso import PSO
from ahdcma.algorithms.scso import SCSO
from ahdcma.algorithms.woa import WOA
from tests.conftest import rastrigin_05, sphere_05


def _run(
    cls: type[Optimizer],
    *,
    dim: int,
    fitness: Callable[[np.ndarray], float],
    pop: int = 20,
    gens: int = 100,
    seed: int = 0,
    extra: dict[str, Any] | None = None,
) -> OptimizationResult:
    config: dict[str, Any] = {
        "population_size": pop,
        "max_generations": gens,
        "seed": seed,
    }
    if extra:
        config.update(extra)
    sp = SearchSpace.unit_cube(dim)
    opt = cls(config, fitness, sp, run_id=f"{cls.__name__}_test")
    return opt.optimize()


ALGOS: list[tuple[type[Optimizer], float]] = [
    (DOA, 1e-2),
    (CMAES, 1e-6),
    (PSO, 1e-2),
    (GWO, 1e-2),
    (WOA, 1e-2),
    (SCSO, 1e-2),
]


@pytest.mark.parametrize("cls,sphere_threshold", ALGOS)
def test_sphere_5d(cls: type[Optimizer], sphere_threshold: float) -> None:
    res = _run(cls, dim=5, fitness=sphere_05, pop=20, gens=100, seed=0)
    assert (
        res.best_f < sphere_threshold
    ), f"{cls.__name__} sphere f={res.best_f:.3g} >= {sphere_threshold}"


@pytest.mark.parametrize("cls,_", ALGOS)
def test_rastrigin_10d(cls: type[Optimizer], _: float) -> None:
    res = _run(cls, dim=10, fitness=rastrigin_05, pop=30, gens=100, seed=0)
    assert res.best_f < 50.0, f"{cls.__name__} rastrigin f={res.best_f:.3g}"


@pytest.mark.parametrize("cls,_", ALGOS)
def test_reproducibility(cls: type[Optimizer], _: float) -> None:
    a = _run(cls, dim=5, fitness=sphere_05, pop=10, gens=20, seed=7)
    b = _run(cls, dim=5, fitness=sphere_05, pop=10, gens=20, seed=7)
    assert a.best_f == b.best_f
    np.testing.assert_array_equal(a.best_x, b.best_x)
    assert len(a.history) == len(b.history)
    for fa, fb in zip(a.history.fitnesses, b.history.fitnesses, strict=True):
        np.testing.assert_array_equal(fa, fb)


def test_doa_history_modes_all_set() -> None:
    res = _run(DOA, dim=4, fitness=sphere_05, pop=10, gens=15, seed=0)
    assert all(m == "explore" for m in res.history.mode_per_gen)
    assert len(res.history) == 16  # initial + 15 iterations


def test_cmaes_rejects_dim_one() -> None:
    sp = SearchSpace.unit_cube(1)
    with pytest.raises(ValueError):
        CMAES(
            {"population_size": 8, "max_generations": 5, "seed": 1},
            sphere_05,
            sp,
        )