"""Tests for statistical_features.py.

Cross-checks the moments against scipy.stats (the reference implementation),
pins closed forms for signals whose statistics are known analytically, guards
the kurtosis mixture trap from the gotcha log, and checks the zero-crossing-rate
convention against the zero-crossing chapter's count.
"""

import numpy as np
import pytest
from scipy import stats

from statistical_features import (
    mean, variance, std, rms, energy, peak_to_peak,
    skewness, kurtosis, zero_crossing_rate, frame_features,
)


# ---------------------------------------------------------------- level / spread

def test_mean_variance_match_numpy():
    rng = np.random.default_rng(0)
    x = rng.standard_normal(1000)
    assert np.isclose(mean(x), np.mean(x))
    assert np.isclose(variance(x), np.var(x))          # population (ddof=0)
    assert np.isclose(std(x), np.std(x))


def test_variance_is_population_not_sample():
    x = np.array([1.0, 2.0, 3.0, 4.0])
    # population variance = 1.25, sample variance (ddof=1) = 1.6667
    assert np.isclose(variance(x), 1.25)
    assert not np.isclose(variance(x), np.var(x, ddof=1))


def test_rms_energy_relation():
    rng = np.random.default_rng(1)
    x = rng.standard_normal(512)
    assert np.isclose(energy(x), np.sum(x ** 2))
    assert np.isclose(rms(x), np.sqrt(energy(x) / x.size))


def test_rms_squared_is_mean_squared_plus_variance():
    # a window with a DC offset: rms**2 = mean**2 + variance
    rng = np.random.default_rng(2)
    x = 3.0 + rng.standard_normal(2000)
    assert np.isclose(rms(x) ** 2, mean(x) ** 2 + variance(x), rtol=1e-6)


def test_peak_to_peak():
    x = np.array([-2.0, 0.5, 1.0, 4.0, -1.0])
    assert peak_to_peak(x) == 6.0                       # 4 - (-2)


def test_sine_rms_and_ptp():
    # a unit sine has rms = 1/sqrt(2) and peak-to-peak = 2
    t = np.linspace(0, 1, 10000, endpoint=False)
    x = np.sin(2 * np.pi * 5 * t)
    assert np.isclose(rms(x), 1 / np.sqrt(2), atol=1e-3)
    assert np.isclose(peak_to_peak(x), 2.0, atol=1e-3)


# ---------------------------------------------------------------------- moments

def test_skewness_matches_scipy():
    rng = np.random.default_rng(3)
    x = rng.exponential(size=4000)                      # right-skewed
    assert np.isclose(skewness(x), stats.skew(x), rtol=1e-9)
    assert skewness(x) > 0                              # long tail on the high side


def test_kurtosis_matches_scipy_excess():
    rng = np.random.default_rng(4)
    x = rng.standard_normal(5000)
    # scipy default is Fisher (excess) + biased, matching our definition
    assert np.isclose(kurtosis(x), stats.kurtosis(x), rtol=1e-9)


def test_symmetric_signal_has_zero_skew():
    t = np.linspace(0, 1, 8000, endpoint=False)
    x = np.sin(2 * np.pi * 3 * t)
    assert abs(skewness(x)) < 1e-6


def test_sine_is_platykurtic():
    # a pure sine has excess kurtosis -1.5 exactly
    t = np.linspace(0, 1, 20000, endpoint=False)
    x = np.sin(2 * np.pi * 7 * t)
    assert np.isclose(kurtosis(x), -1.5, atol=1e-2)


def test_uniform_excess_kurtosis():
    # uniform noise has excess kurtosis -1.2 in the limit
    rng = np.random.default_rng(5)
    x = rng.uniform(-1, 1, size=200000)
    assert np.isclose(kurtosis(x), -1.2, atol=0.05)


def test_spiky_signal_is_leptokurtic():
    # mostly quiet with a few large spikes -> heavy tails -> positive excess
    x = np.zeros(1000)
    x[::200] = 10.0
    assert kurtosis(x) > 3.0


def test_kurtosis_mixture_trap():
    """Gotcha pin: a narrow Gaussian + a WIDE uniform spike is leptokurtic.

    The uniform component is platykurtic in isolation, but because it is much
    wider than the Gaussian it puts mass in the tails, so the *mixture* has
    positive excess kurtosis. Reason about the mixture, not the parts.
    """
    rng = np.random.default_rng(6)
    narrow = rng.standard_normal(9000)                 # sigma = 1
    wide = rng.uniform(-5, 5, size=1000)               # sigma ~ 2.9, the tails
    mixture = np.concatenate([narrow, wide])
    assert kurtosis(mixture) > 0                        # leptokurtic, not platykurtic


def test_moments_zero_on_constant_window():
    x = np.full(64, 2.5)
    assert variance(x) == 0.0
    assert skewness(x) == 0.0                           # guarded, not nan
    assert kurtosis(x) == 0.0


# --------------------------------------------------------------- zero crossings

def test_zcr_sine_matches_frequency_count():
    # a sine at f0 over duration T has ~2*f0*T crossings; ZCR ~ that / (n-1)
    fs, f0, T = 8000, 100, 1.0
    t = np.arange(int(fs * T)) / fs
    x = np.sin(2 * np.pi * f0 * t)
    n = x.size
    expected_count = 2 * f0 * T
    assert np.isclose(zero_crossing_rate(x) * (n - 1), expected_count, rtol=0.02)


def test_zcr_recovers_frequency():
    # invert the rate to a frequency estimate, the zero-crossing chapter's formula
    fs, f0 = 4000, 60
    t = np.arange(fs) / fs
    x = np.sin(2 * np.pi * f0 * t)
    n = x.size
    f_hat = zero_crossing_rate(x) * (n - 1) * fs / (2 * (n - 1))
    assert np.isclose(f_hat, f0, rtol=0.02)


def test_zcr_noise_higher_than_tone():
    rng = np.random.default_rng(7)
    t = np.arange(4000) / 4000
    tone = np.sin(2 * np.pi * 20 * t)                  # slow tone, few crossings
    noise = rng.standard_normal(4000)                  # white-ish, many crossings
    assert zero_crossing_rate(noise) > zero_crossing_rate(tone)


def test_zcr_ignores_exact_zeros():
    # a run of zeros between two same-sign values is not a crossing
    x = np.array([1.0, 0.0, 0.0, 2.0])
    assert zero_crossing_rate(x) == 0.0
    # but a genuine sign change is counted
    y = np.array([1.0, -1.0, 1.0])
    assert np.isclose(zero_crossing_rate(y), 1.0)       # both pairs cross


def test_zcr_in_unit_interval():
    rng = np.random.default_rng(8)
    x = rng.standard_normal(500)
    z = zero_crossing_rate(x)
    assert 0.0 <= z <= 1.0


# ----------------------------------------------------------------- broadcasting

def test_features_reduce_along_last_axis():
    rng = np.random.default_rng(9)
    frames = rng.standard_normal((7, 256))             # 7 windows of 256 samples
    v = variance(frames)
    assert v.shape == (7,)
    # each entry equals the per-window variance computed independently
    for i in range(7):
        assert np.isclose(v[i], variance(frames[i]))


def test_frame_features_bundle_keys_and_shapes():
    rng = np.random.default_rng(10)
    frames = rng.standard_normal((5, 128))
    feats = frame_features(frames)
    expected = {"mean", "variance", "rms", "energy", "peak_to_peak",
                "skewness", "kurtosis", "zero_crossing_rate"}
    assert set(feats) == expected
    for name, val in feats.items():
        assert val.shape == (5,), name


def test_frame_features_single_window_scalars():
    rng = np.random.default_rng(11)
    x = rng.standard_normal(300)
    feats = frame_features(x)
    for name, val in feats.items():
        assert np.ndim(val) == 0, name


# ------------------------------------------------------ a discrimination check

def test_still_vs_moving_separable_by_variance():
    """The embedded-classifier premise: variance separates still from moving."""
    rng = np.random.default_rng(12)
    still = 1.0 + 0.01 * rng.standard_normal((20, 128))    # tiny jitter about 1g
    moving = 1.0 + 0.5 * rng.standard_normal((20, 128))    # large swings
    assert variance(still).max() < variance(moving).min()  # cleanly separable