"""Tests for random_processes.py: generators and characterisation."""

import numpy as np
import pytest
from random_processes import (
    gaussian_white, gaussian_ar1, poisson_process, ar_process,
    levy_impulse_noise, cyclostationary_noise, telegraph_noise,
    characterize,
)


class TestGaussianWhite:
    def test_shape(self):
        x = gaussian_white(1000)
        assert x.shape == (1000,)

    def test_approx_zero_mean(self):
        x = gaussian_white(100000, sigma=2.0, seed=42)
        assert abs(np.mean(x)) < 0.05  # should be close to zero

    def test_std_close_to_sigma(self):
        x = gaussian_white(50000, sigma=3.0, seed=7)
        assert 2.9 < np.std(x) < 3.1

    def test_acf_zero_at_nonzero_lag(self):
        x = gaussian_white(10000, seed=99)
        xc = x - np.mean(x)
        acf = np.correlate(xc, xc, mode='full')
        mid = len(xc) - 1
        r0 = acf[mid]
        r1 = acf[mid + 1]
        # White noise: |r1| / r0 should be tiny.
        assert abs(r1 / r0) < 0.03

    def test_reproducible_with_seed(self):
        a = gaussian_white(100, seed=42)
        b = gaussian_white(100, seed=42)
        assert np.array_equal(a, b)


class TestGaussianAR1:
    def test_shape(self):
        x = gaussian_ar1(500, phi=0.9)
        assert x.shape == (500,)

    def test_positive_acf_lag1(self):
        x = gaussian_ar1(50000, phi=0.9, seed=1)
        info = characterize(x)
        assert info['acf_lag1'] > 0.5  # strong positive correlation

    def test_negative_phi_gives_alternating_acf(self):
        x = gaussian_ar1(50000, phi=-0.8, seed=1)
        info = characterize(x)
        assert info['acf_lag1'] < -0.3
        assert info['acf_lag2'] > 0.1   # alternating sign => lag2 positive

    def test_phi_zero_is_white(self):
        x = gaussian_ar1(50000, phi=0.0, seed=1)
        info = characterize(x)
        assert abs(info['acf_lag1']) < 0.03

    def test_sigma_scales_output(self):
        # Regression: sigma used to be cancelled by an output renormalisation.
        a = gaussian_ar1(50000, phi=0.9, sigma=1.0, seed=1)
        b = gaussian_ar1(50000, phi=0.9, sigma=5.0, seed=1)
        assert np.allclose(b, 5.0 * a)

    def test_variance_matches_theory(self):
        # AR(1): var = sigma^2 / (1 - phi^2)
        phi, sigma = 0.9, 2.0
        x = gaussian_ar1(200000, phi=phi, sigma=sigma, seed=3)
        expected = sigma**2 / (1 - phi**2)
        assert abs(np.var(x) / expected - 1) < 0.1


class TestPoissonProcess:
    def test_shape(self):
        x = poisson_process(200, rate=5.0)
        assert x.shape == (200,)

    def test_mean_close_to_rate(self):
        x = poisson_process(100000, rate=15.0, seed=42)
        assert 14.5 < np.mean(x) < 15.5

    def test_variance_close_to_mean(self):
        # For Poisson, variance == mean.
        x = poisson_process(100000, rate=20.0, seed=42)
        mu = np.mean(x)
        var = np.var(x)
        assert 0.9 < var / mu < 1.1

    def test_low_rate_is_discrete(self):
        x = poisson_process(10000, rate=0.3, seed=7)
        # Poisson with rate 0.3 should have > 99% of values in {0, 1, 2}.
        frac_low = np.mean((x == 0) | (x == 1) | (x == 2))
        assert frac_low > 0.99


class TestARProcess:
    def test_shape(self):
        x = ar_process(1000, coeffs=[-0.5])
        assert x.shape == (1000,)

    def test_ar2_has_two_nonzero_acf_lags(self):
        # AR(2) with mild coefficients should have structured ACF.
        x = ar_process(50000, coeffs=[-0.7, 0.2], seed=1)
        info = characterize(x)
        assert abs(info['acf_lag1']) > 0.2
        assert abs(info['acf_lag2']) > 0.05


class TestLevyImpulseNoise:
    def test_shape(self):
        x = levy_impulse_noise(500)
        assert x.shape == (500,)

    def test_excess_kurtosis_positive(self):
        x = levy_impulse_noise(100000, scale=10.0, outlier_fraction=0.05, seed=42)
        info = characterize(x)
        # Heavy-tailed mixture should have excess kurtosis > 0.
        assert info['excess_kurtosis'] > 1.0

    def test_large_excursions_exist(self):
        x = levy_impulse_noise(100000, scale=20.0, outlier_fraction=0.02, seed=1)
        # Should have samples exceeding 10 sigma of the core distribution.
        assert np.max(np.abs(x)) > 8.0

    def test_no_outliers_gives_gaussian(self):
        x = levy_impulse_noise(50000, outlier_fraction=0.0, seed=1)
        info = characterize(x)
        assert abs(info['excess_kurtosis']) < 0.3


class TestCyclostationary:
    def test_shape(self):
        x = cyclostationary_noise(1000)
        assert x.shape == (1000,)

    def test_variance_varies_with_period(self):
        n = 100000
        period = 50
        x = cyclostationary_noise(n, period=period, modulation_depth=0.5, seed=1)
        # Compute running variance over windows of length period.
        n_windows = n // period
        var_per_window = np.array(
            [np.var(x[i*period:(i+1)*period]) for i in range(n_windows)]
        )
        # With modulation depth 0.5 the variance should vary noticeably.
        assert np.std(var_per_window) / np.mean(var_per_window) > 0.05

    def test_zero_modulation_is_stationary_white(self):
        x = cyclostationary_noise(50000, modulation_depth=0.0, seed=1)
        info = characterize(x)
        assert abs(info['acf_lag1']) < 0.03


class TestTelegraphNoise:
    def test_shape(self):
        x = telegraph_noise(500)
        assert x.shape == (500,)

    def test_only_two_levels(self):
        x = telegraph_noise(10000, levels=(-1.0, 1.0), seed=42)
        assert set(np.unique(x)) <= {-1.0, 1.0}

    def test_custom_levels(self):
        x = telegraph_noise(5000, levels=(0.0, 3.0), seed=99)
        assert set(np.unique(x)) <= {0.0, 3.0}

    def test_acf_decays_with_lag(self):
        x = telegraph_noise(50000, p_switch=0.01, seed=1)
        info = characterize(x)
        # Telegraph signal has exponential ACF decay: |r1| > |r5| (in expectation).
        assert abs(info['acf_lag1']) > abs(info['acf_lag5'])

    def test_fast_switching_approaches_white(self):
        x = telegraph_noise(50000, p_switch=0.5, seed=1)
        info = characterize(x)
        # Near-50% switching probability -> nearly uncorrelated.
        assert abs(info['acf_lag1']) < 0.1


class TestCharacterize:
    def test_returns_expected_keys(self):
        info = characterize(gaussian_white(1000, seed=1))
        expected = {'mean', 'var', 'std', 'skewness', 'excess_kurtosis',
                    'acf_lag1', 'acf_lag2', 'acf_lag5',
                    'psd_slope', 'psd_r2', 'is_zero_mean', 'is_symmetric'}
        assert set(info.keys()) == expected

    def test_white_noise_psd_flat(self):
        x = gaussian_white(10000, seed=42)
        info = characterize(x, fs=1000)
        assert abs(info['psd_slope']) < 0.3  # flat spectrum → slope near zero

    def test_brownian_psd_slope(self):
        # Brownian noise (cumulative sum of white) has 1/f^2 → slope ≈ -2.
        x = np.cumsum(gaussian_white(20000, seed=7))
        info = characterize(x, fs=1000)
        assert -2.5 < info['psd_slope'] < -1.5

    def test_short_signal_handles_gracefully(self):
        info = characterize(gaussian_white(20, seed=1))
        # Should not crash on short input.
        assert isinstance(info, dict)
        assert 'mean' in info