"""Unit tests for particle swarm optimization and PSO filter design."""

import numpy as np
import pytest

from pso import (
    pso,
    gradient_descent,
    iir_is_stable,
    iir_magnitude_fitness,
    design_iir_pso,
)
from scipy import signal


def sphere(x):
    return float(np.sum(np.asarray(x) ** 2))


def rastrigin(x):
    x = np.asarray(x)
    return float(10 * len(x) + np.sum(x ** 2 - 10 * np.cos(2 * np.pi * x)))


# ---------------------------------------------------------------------------
# Core PSO
# ---------------------------------------------------------------------------

class TestPSO:
    def test_converges_on_convex_sphere(self):
        rng = np.random.default_rng(0)
        best, fbest, hist = pso(sphere, [(-5, 5)] * 3, rng=rng)
        assert fbest < 1e-6
        assert np.allclose(best, 0, atol=1e-2)

    def test_history_is_non_increasing(self):
        # gbest is updated only on strict improvement, so the best-so-far
        # history is exactly monotone non-increasing.
        rng = np.random.default_rng(1)
        _, _, hist = pso(sphere, [(-5, 5)] * 2, rng=rng)
        assert np.all(np.diff(hist) <= 0)

    def test_default_rng_runs(self):
        # The unseeded path should execute; check only shapes, not values.
        best, fbest, hist = pso(sphere, [(-1, 1)] * 2, n_particles=5,
                                n_iters=10)
        assert best.shape == (2,)
        assert hist.shape == (11,)
        assert isinstance(fbest, float)

    def test_reproducible_with_seed(self):
        a = pso(sphere, [(-5, 5)] * 2, rng=np.random.default_rng(42))
        b = pso(sphere, [(-5, 5)] * 2, rng=np.random.default_rng(42))
        assert a[1] == b[1]
        np.testing.assert_array_equal(a[0], b[0])

    def test_finds_global_basin_of_multimodal(self):
        # Rastrigin's global minimum is 0 at the origin, surrounded by many
        # local minima. PSO should reach near it.
        rng = np.random.default_rng(3)
        best, fbest, _ = pso(rastrigin, [(-5.12, 5.12)] * 2,
                             n_particles=40, n_iters=200, rng=rng)
        assert fbest < 1.0

    def test_respects_bounds(self):
        rng = np.random.default_rng(5)
        best, _, _ = pso(sphere, [(2, 5), (2, 5)], rng=rng)
        assert np.all(best >= 2) and np.all(best <= 5)


# ---------------------------------------------------------------------------
# Gradient-descent baseline (the contrast)
# ---------------------------------------------------------------------------

class TestGradientContrast:
    def test_gradient_descent_solves_convex(self):
        x, fx, _ = gradient_descent(sphere, [3.0, -2.0], [(-5, 5)] * 2,
                                    lr=0.1, n_iters=500)
        assert fx < 1e-3

    def test_gradient_descent_traps_in_local_minimum(self):
        # From a bad start on Rastrigin, gradient descent stalls in a nearby
        # local minimum, while PSO reaches the global basin. This is the
        # whole motivation for the topic.
        bounds = [(-5.12, 5.12)] * 2
        x_gd, f_gd, _ = gradient_descent(rastrigin, [4.0, -3.0], bounds)
        _, f_pso, _ = pso(rastrigin, bounds, n_particles=40, n_iters=200,
                          rng=np.random.default_rng(7))
        assert f_gd > 5.0          # trapped
        assert f_pso < 1.0         # escaped
        assert f_pso < f_gd


# ---------------------------------------------------------------------------
# IIR design
# ---------------------------------------------------------------------------

class TestStability:
    def test_known_stable(self):
        # Butterworth poles are inside the unit circle.
        sos = signal.butter(2, 1000, fs=8000, output='sos')
        assert iir_is_stable(sos[0, 4], sos[0, 5])

    def test_known_unstable(self):
        assert not iir_is_stable(0.0, 1.5)   # |a2| > 1
        assert not iir_is_stable(2.5, 0.5)   # |a1| > 1 + a2

    def test_unstable_coeffs_get_penalised(self):
        freqs = np.linspace(0, 4000, 50)
        target = np.ones_like(freqs)
        bad = iir_magnitude_fitness([1, 0, 0, 0.0, 1.5], freqs, target, 8000)
        assert bad == pytest.approx(10.0)


class TestDesignIIR:
    def test_designs_stable_filter_matching_target(self):
        fs = 8000
        freqs = np.linspace(0, fs / 2, 200)
        sos_t = signal.butter(2, [800, 1600], btype='band', fs=fs,
                              output='sos')
        _, ht = signal.sosfreqz(sos_t, worN=freqs, fs=fs)
        target = np.abs(ht)

        sos, err, hist = design_iir_pso(freqs, target, fs,
                                        rng=np.random.default_rng(2))
        assert iir_is_stable(sos[0, 4], sos[0, 5])
        assert err < 0.05
        assert sos.shape == (1, 6)