"""Tests for the sinusoid parameter estimators and their bounds."""

import numpy as np
import pytest

from sinusoid import (coherent_dft, crlb_sinusoid, crlb_sinusoid_exact,
                      crlb_known_frequency, estimate_frequency_fft,
                      estimate_frequency_phase_diff)


def tone(A=1.0, f0=0.2334, phi=0.7, N=256, sigma=0.0, fs=1.0, rng=None):
    """A real tone, optionally in white Gaussian noise."""
    x = A * np.cos(2 * np.pi * f0 * np.arange(N) / fs + phi)
    if sigma > 0:
        x = x + (rng or np.random.default_rng(0)).normal(0, sigma, N)
    return x


# ---------------------------------------------------------------------------
# coherent_dft
# ---------------------------------------------------------------------------

class TestCoherentDFT:
    def test_recovers_amplitude_and_phase_without_noise(self):
        A, phi, f0 = 2.5, -1.1, 0.13
        a, p = coherent_dft(tone(A=A, f0=f0, phi=phi, N=1000), f0)
        assert a == pytest.approx(A, rel=1e-3)
        assert p == pytest.approx(phi, abs=1e-3)

    def test_works_in_hz(self):
        fs, f0, A = 8000.0, 440.0, 1.7
        a, _ = coherent_dft(tone(A=A, f0=f0, N=4000, fs=fs), f0, fs=fs)
        assert a == pytest.approx(A, rel=1e-2)

    def test_phase_wraps_to_pi_interval(self):
        for phi in (-3.0, -0.2, 0.0, 2.9):
            _, p = coherent_dft(tone(phi=phi, N=2000), 0.2334)
            assert -np.pi < p <= np.pi
            assert np.cos(p) == pytest.approx(np.cos(phi), abs=1e-2)

    def test_rejects_wrong_window_length(self):
        with pytest.raises(ValueError):
            coherent_dft(tone(N=64), 0.2, window=np.hanning(32))

    def test_rejects_empty_input(self):
        with pytest.raises(ValueError):
            coherent_dft(np.array([]), 0.2)

    def test_window_reduces_leakage_from_a_second_tone(self):
        """A window costs noise performance but buys interferer rejection."""
        N = 512
        f_want, f_other = 0.10, 0.13
        x = tone(A=1.0, f0=f_want, N=N) + tone(A=50.0, f0=f_other, phi=0.3, N=N)
        a_rect, _ = coherent_dft(x, f_want)
        a_win, _ = coherent_dft(x, f_want, window=np.blackman(N))
        assert abs(a_win - 1.0) < abs(a_rect - 1.0)

    def test_attains_the_known_frequency_bound(self):
        """Coherent detection is the MLE for known f0, so it is efficient."""
        rng = np.random.default_rng(4)
        A, sigma, N, f0 = 1.0, 0.1, 256, 0.2334
        amps, phases = [], []
        for _ in range(4000):
            a, p = coherent_dft(tone(A=A, f0=f0, N=N, sigma=sigma, rng=rng), f0)
            amps.append(a)
            phases.append(p)
        bound = crlb_known_frequency(A, sigma, N)
        # 4000 trials give ~2.2% standard error on a variance.
        assert np.var(amps) == pytest.approx(bound['amplitude'], rel=0.10)
        assert np.var(phases) == pytest.approx(bound['phase'], rel=0.10)


# ---------------------------------------------------------------------------
# Cramer-Rao bounds
# ---------------------------------------------------------------------------

class TestCRLB:
    @pytest.mark.parametrize("f0", [0.10, 0.23, 0.37])
    def test_closed_forms_match_the_exact_fisher_information(self, f0):
        """The large-N approximations must agree away from DC/Nyquist."""
        closed = crlb_sinusoid(1.0, 0.1, 256)
        exact = crlb_sinusoid_exact(1.0, 0.7, f0, 0.1, 256)
        for key in ('amplitude', 'phase', 'frequency'):
            assert closed[key] == pytest.approx(exact[key], rel=0.05), key

    def test_closed_form_breaks_near_dc(self):
        """Honest limit: the approximation drops 2*f0 terms that stop
        averaging away when the tone is near DC."""
        closed = crlb_sinusoid(1.0, 0.1, 256)
        exact = crlb_sinusoid_exact(1.0, 0.7, 0.002, 0.1, 256)
        assert closed['frequency'] / exact['frequency'] < 0.8

    def test_frequency_bound_falls_as_n_cubed(self):
        prev = None
        for N in (128, 256, 512, 1024):
            v = crlb_sinusoid(1.0, 0.1, N)['frequency']
            if prev is not None:
                assert prev / v == pytest.approx(8.0, rel=0.05)
            prev = v

    def test_amplitude_bound_falls_as_one_over_n(self):
        b = crlb_sinusoid(1.0, 0.1, 100)['amplitude']
        assert crlb_sinusoid(1.0, 0.1, 400)['amplitude'] == pytest.approx(b / 4)

    def test_snr_parameterisation_trap(self):
        """The frequency bound is 12/(eta ...) with eta = A^2/(2 sigma^2).

        Writing it as 12 sigma^2/A^2 instead of 12/eta silently halves
        the bound: the two differ by exactly the factor of 2 that eta
        carries for a real sinusoid.  Pin the correct one.
        """
        A, sigma, N = 1.0, 0.1, 256
        eta = A**2 / (2 * sigma**2)
        got = crlb_sinusoid(A, sigma, N)['frequency']

        right = 12.0 / ((2 * np.pi) ** 2 * eta * N * (N**2 - 1))
        wrong = 12.0 * sigma**2 / ((2 * np.pi) ** 2 * A**2 * N * (N**2 - 1))
        assert got == pytest.approx(right, rel=1e-12)
        assert got == pytest.approx(2 * wrong, rel=1e-12)
        # ... and the correct form is the one the exact FIM agrees with.
        exact = crlb_sinusoid_exact(A, 0.7, 0.23, sigma, N)['frequency']
        assert got == pytest.approx(exact, rel=0.05)

    def test_known_frequency_improves_phase_fourfold(self):
        unknown = crlb_sinusoid(1.0, 0.1, 256)
        known = crlb_known_frequency(1.0, 0.1, 256)
        assert known['amplitude'] == pytest.approx(unknown['amplitude'])
        assert known['phase'] == pytest.approx(unknown['phase'] / 4)

    def test_bound_scales_with_sample_rate_squared(self):
        b1 = crlb_sinusoid(1.0, 0.1, 256, fs=1.0)['frequency']
        b2 = crlb_sinusoid(1.0, 0.1, 256, fs=1000.0)['frequency']
        assert b2 == pytest.approx(b1 * 1000.0**2)

    def test_rejects_bad_input(self):
        with pytest.raises(ValueError):
            crlb_sinusoid(0.0, 0.1, 100)
        with pytest.raises(ValueError):
            crlb_sinusoid(1.0, 0.0, 100)
        with pytest.raises(ValueError):
            crlb_sinusoid(1.0, 0.1, 1)
        with pytest.raises(ValueError):
            crlb_known_frequency(1.0, -1.0, 100)


# ---------------------------------------------------------------------------
# Frequency estimators
# ---------------------------------------------------------------------------

class TestFrequencyEstimators:
    def test_bare_peak_is_quantised_to_the_bin_grid(self):
        N = 256
        f = estimate_frequency_fft(tone(f0=0.2334, N=N), interpolation='none')
        assert f * N == pytest.approx(round(f * N))       # lands on a bin

    def test_interpolation_beats_the_bare_peak_off_bin(self):
        N, f0 = 256, (20 + 0.4) / 256                     # deliberately off-bin
        x = tone(f0=f0, N=N)
        bare = abs(estimate_frequency_fft(x, interpolation='none') - f0)
        jac = abs(estimate_frequency_fft(x, interpolation='jacobsen') - f0)
        assert jac < bare / 10

    def test_on_bin_tone_needs_no_interpolation(self):
        N, f0 = 256, 20 / 256
        for mode in ('none', 'quadratic', 'jacobsen'):
            f = estimate_frequency_fft(tone(f0=f0, N=N), interpolation=mode)
            assert f == pytest.approx(f0, abs=1e-6), mode

    def test_jacobsen_suits_the_rectangular_window(self):
        """Interpolator/window pairing: the mismatched choice is worse."""
        rng = np.random.default_rng(1)
        N, sigma = 256, 0.05
        errs = {'quadratic': [], 'jacobsen': []}
        for i in range(200):
            f0 = (20 + (i % 97) / 97.0) / N
            x = tone(f0=f0, N=N, sigma=sigma, rng=rng)
            for mode in errs:
                errs[mode].append(estimate_frequency_fft(x, interpolation=mode) - f0)
        rms = {m: np.sqrt(np.mean(np.square(e))) for m, e in errs.items()}
        assert rms['jacobsen'] < rms['quadratic'] / 5

    def test_quadratic_suits_a_tapered_window(self):
        """The same comparison flips once the window is tapered."""
        rng = np.random.default_rng(1)
        N, sigma = 256, 0.05
        w = np.blackman(N)
        errs = {'quadratic': [], 'jacobsen': []}
        for i in range(200):
            f0 = (20 + (i % 97) / 97.0) / N
            x = tone(f0=f0, N=N, sigma=sigma, rng=rng) * w
            for mode in errs:
                errs[mode].append(estimate_frequency_fft(x, interpolation=mode) - f0)
        rms = {m: np.sqrt(np.mean(np.square(e))) for m, e in errs.items()}
        assert rms['quadratic'] < rms['jacobsen'] / 5

    def test_works_in_hz(self):
        fs, f0, N = 8000.0, 437.3, 2048
        f = estimate_frequency_fft(tone(f0=f0, N=N, fs=fs), fs=fs,
                                   interpolation='jacobsen')
        assert f == pytest.approx(f0, rel=1e-3)

    def test_phase_difference_estimator_is_accurate(self):
        N, f0 = 512, (40 + 0.37) / 512
        f = estimate_frequency_phase_diff(tone(f0=f0, N=N))
        assert f == pytest.approx(f0, rel=1e-4)

    def test_phase_difference_within_a_small_multiple_of_the_bound(self):
        rng = np.random.default_rng(2)
        A, sigma, N, f0 = 1.0, 0.2, 256, 0.2334
        est = np.array([estimate_frequency_phase_diff(
            tone(A=A, f0=f0, N=N, sigma=sigma, rng=rng)) for _ in range(600)])
        rmse = np.sqrt(np.mean((est - f0) ** 2))
        bound_std = np.sqrt(crlb_sinusoid(A, sigma, N)['frequency'])
        assert rmse < 2.0 * bound_std

    def test_rejects_bad_input(self):
        with pytest.raises(ValueError):
            estimate_frequency_fft(np.ones(2))
        with pytest.raises(ValueError):
            estimate_frequency_fft(tone(N=64), interpolation='spline')
        with pytest.raises(ValueError):
            estimate_frequency_phase_diff(np.ones(3))