"""Tests for adaptive filtering algorithms."""

import numpy as np
import pytest

from adaptive import LMS, NLMS, RLS, identify_system


# --- Construction ---

def test_lms_requires_positive_taps():
    with pytest.raises(ValueError):
        LMS(0)

def test_lms_requires_positive_mu():
    with pytest.raises(ValueError):
        LMS(4, mu=-0.01)

def test_nlms_requires_positive_taps():
    with pytest.raises(ValueError):
        NLMS(0)

def test_rls_requires_valid_lambda():
    with pytest.raises(ValueError):
        RLS(4, lam=0.0)
    with pytest.raises(ValueError):
        RLS(4, lam=1.5)


# --- Basic convergence ---

def _converges(cls, n_taps=4, **kwargs):
    """Helper: verify that a filter converges on a known system."""
    rng = np.random.default_rng(0)
    true_w = rng.standard_normal(n_taps)
    x = rng.standard_normal(2000)
    d = np.convolve(x, true_w)[:2000]
    filt = cls(n_taps, **kwargs)
    _, e = filt.run(x, d)
    # Error should decrease significantly
    early = np.mean(e[:200] ** 2)
    late = np.mean(e[-200:] ** 2)
    assert late < early * 0.01, f"Filter did not converge: early MSE={early:.4f}, late MSE={late:.4f}"
    # Coefficients should be close to true weights
    np.testing.assert_allclose(filt.w, true_w, atol=0.05)

def test_lms_converges():
    _converges(LMS, mu=0.01)

def test_nlms_converges():
    _converges(NLMS, mu=0.5)

def test_rls_converges():
    _converges(RLS, lam=0.99)


# --- RLS converges faster than LMS ---

def test_rls_converges_faster_than_lms():
    rng = np.random.default_rng(1)
    true_w = rng.standard_normal(8)
    x = rng.standard_normal(1000)
    d = np.convolve(x, true_w)[:1000]

    lms = LMS(8, mu=0.005)
    _, e_lms = lms.run(x, d)

    rls = RLS(8, lam=0.99)
    _, e_rls = rls.run(x, d)

    # RLS should have lower MSE in the first 200 samples
    mse_lms = np.mean(e_lms[:200] ** 2)
    mse_rls = np.mean(e_rls[:200] ** 2)
    assert mse_rls < mse_lms


# --- identify_system convenience function ---

def test_identify_system_nlms():
    true_ir = np.array([1.0, -0.5, 0.25, -0.1])
    w, e, w_hist = identify_system(true_ir, n_samples=3000, algorithm="nlms")
    np.testing.assert_allclose(w, true_ir, atol=0.1)
    assert w_hist.shape == (3000, 4)

def test_identify_system_rls():
    true_ir = np.array([0.5, 0.3, -0.2])
    w, e, _ = identify_system(true_ir, n_samples=2000, algorithm="rls")
    np.testing.assert_allclose(w, true_ir, atol=0.05)

def test_identify_system_unknown_algorithm():
    with pytest.raises(ValueError, match="Unknown algorithm"):
        identify_system(np.array([1.0]), algorithm="foo")


# --- Noise cancellation scenario ---

def test_noise_cancellation():
    """Adaptive filter should remove correlated noise from a signal."""
    rng = np.random.default_rng(2)
    n = 3000

    # Clean signal (low-frequency sine)
    t = np.arange(n)
    clean = np.sin(2 * np.pi * 0.01 * t)

    # Noise source and its filtered version reaching the primary mic
    noise_ref = rng.standard_normal(n)
    noise_path = np.array([0.8, -0.4, 0.2])
    noise_at_mic = np.convolve(noise_ref, noise_path)[:n]

    # Primary signal = clean + noise
    primary = clean + noise_at_mic

    # Use NLMS to cancel noise; longer run for convergence
    n_run = 5000
    noise_ref = rng.standard_normal(n_run)
    noise_at_mic = np.convolve(noise_ref, noise_path)[:n_run]
    t = np.arange(n_run)
    clean = np.sin(2 * np.pi * 0.01 * t)
    primary = clean + noise_at_mic

    filt = NLMS(8, mu=0.5)
    _, e = filt.run(noise_ref, primary)

    # After convergence, error should approximate the clean signal
    # (some residual expected due to finite step size)
    residual_noise = np.std(e[-1000:] - clean[-1000:])
    assert residual_noise < 0.5, f"Noise cancellation residual too high: {residual_noise:.3f}"


# --- Edge cases ---

def test_single_tap():
    """A single-tap adaptive filter should learn a scalar gain."""
    filt = NLMS(1, mu=0.5)
    x = np.ones(500)
    d = 3.0 * x
    _, e = filt.run(x, d)
    assert abs(filt.w[0] - 3.0) < 0.01

def test_zero_input():
    """Filter should handle zero input without crashing (NLMS eps prevents div by zero)."""
    filt = NLMS(4, mu=0.5)
    y, e = filt.update(0.0, 1.0)
    assert np.isfinite(y)
    assert np.isfinite(e)