"""Unit tests for noise generation, characterisation, and whitening."""

import numpy as np
import pytest

from whitening import (
    acf,
    kasdin_coefficients,
    generate_power_law_noise_psd,
    generate_power_law_noise_ar,
    whiten,
    build_whitening_filter,
)


# ---------------------------------------------------------------------------
# acf()
# ---------------------------------------------------------------------------

class TestACF:
    def test_lag_zero_is_one(self):
        x = np.random.randn(512)
        r = acf(x, max_lag=10)
        assert r[0] == pytest.approx(1.0)

    def test_all_zero_input_returns_zeros(self):
        x = np.zeros(128)
        r = acf(x, max_lag=5)
        np.testing.assert_array_equal(r, np.zeros(6))

    def test_white_noise_low_autocorrelation(self):
        np.random.seed(42)
        x = np.random.randn(10_000)
        r = acf(x, max_lag=20)
        # Lag > 0 autocorrelation should be small for white noise
        assert all(abs(r[k]) < 0.05 for k in range(1, 21))

    def test_max_lag_respected(self):
        x = np.random.randn(256)
        r = acf(x, max_lag=7)
        assert len(r) == 8  # lags 0 through 7


# ---------------------------------------------------------------------------
# kasdin_coefficients()
# ---------------------------------------------------------------------------

class TestKasdinCoefficients:
    def test_white_noise_short_list(self):
        coeffs = kasdin_coefficients(alpha=0.0)
        # alpha=0 means all subsequent coefficients are zero → just [1.0]
        assert len(coeffs) == 1
        assert coeffs[0] == 1.0

    def test_brownian_first_coefficients(self):
        coeffs = kasdin_coefficients(alpha=2.0)
        assert coeffs[0] == pytest.approx(1.0)
        assert coeffs[1] == pytest.approx(-1.0)

    def test_starts_with_one(self):
        for alpha in [0.5, 1.0, 1.5]:
            coeffs = kasdin_coefficients(alpha=alpha)
            assert coeffs[0] == 1.0


# ---------------------------------------------------------------------------
# generate_power_law_noise_psd()
# ---------------------------------------------------------------------------

class TestGeneratePowerLawNoisePSD:
    def test_output_length(self):
        n = 256
        y = generate_power_law_noise_psd(n, alpha=1.0)
        assert len(y) == n

    def test_output_is_real(self):
        y = generate_power_law_noise_psd(512, alpha=1.0)
        assert np.isrealobj(y)


# ---------------------------------------------------------------------------
# generate_power_law_noise_ar()
# ---------------------------------------------------------------------------

class TestGeneratePowerLawNoiseAR:
    def test_output_length(self):
        n = 300
        y = generate_power_law_noise_ar(n, alpha=1.0)
        assert len(y) == n


# ---------------------------------------------------------------------------
# whiten()
# ---------------------------------------------------------------------------

class TestWhiten:
    def test_whitening_reduces_lag1_autocorrelation(self):
        np.random.seed(0)
        alpha = 2.0
        x = generate_power_law_noise_ar(5000, alpha=alpha)
        r_before = acf(x, max_lag=1)[1]
        y = whiten(x, alpha=alpha)
        r_after = acf(y, max_lag=1)[1]
        assert abs(r_after) < abs(r_before)
        # After whitening alpha=2 noise, lag-1 correlation should be near zero
        assert abs(r_after) < 0.1


# ---------------------------------------------------------------------------
# build_whitening_filter()
# ---------------------------------------------------------------------------

class TestBuildWhiteningFilter:
    def test_returns_numpy_array(self):
        h = build_whitening_filter(alpha=1.0)
        assert isinstance(h, np.ndarray)

    def test_starts_with_one(self):
        h = build_whitening_filter(alpha=1.5)
        assert h[0] == pytest.approx(1.0)