"""Tests for adc_noise.py."""

import numpy as np
import pytest
from adc_noise import (jitter_snr, enob_from_sinad, oversampling_gain,
                       sigma_delta_modulate, sigma_delta_decimate, dds_sine,
                        sigma_delta_snr, adc_noise_budget)


class TestJitterSNR:
    def test_higher_frequency_worse_snr(self):
        snr_1k = jitter_snr(1e3, 10e-12)
        snr_10k = jitter_snr(10e3, 10e-12)
        assert snr_10k < snr_1k

    def test_less_jitter_better_snr(self):
        snr_bad = jitter_snr(10e3, 100e-12)
        snr_good = jitter_snr(10e3, 1e-12)
        assert snr_good > snr_bad

    def test_typical_value(self):
        # 10 kHz input, 10 ps jitter -> SNR ~ 124 dB
        snr = jitter_snr(10e3, 10e-12)
        assert 120 < snr < 130


class TestENOB:
    def test_ideal_12bit(self):
        sinad_12 = 6.02 * 12 + 1.76  # ideal 12-bit SINAD
        enob = enob_from_sinad(sinad_12)
        assert abs(enob - 12.0) < 0.1

    def test_degraded(self):
        enob = enob_from_sinad(60.0)
        assert 9.0 < enob < 10.0


class TestOversamplingGain:
    def test_osr_4_gives_6db(self):
        gain = oversampling_gain(4)
        assert abs(gain - 6.02) < 0.1

    def test_osr_16_gives_12db(self):
        gain = oversampling_gain(16)
        assert abs(gain - 12.04) < 0.2


class TestSigmaDelta:
    def test_first_order_gain(self):
        sqnr_osr1 = sigma_delta_snr(1, 4, order=1)
        sqnr_osr4 = sigma_delta_snr(1, 16, order=1)
        assert sqnr_osr4 > sqnr_osr1

    def test_second_order_better(self):
        sqnr_l1 = sigma_delta_snr(1, 64, order=1)
        sqnr_l2 = sigma_delta_snr(1, 64, order=2)
        assert sqnr_l2 > sqnr_l1


class TestNoiseBudget:
    def test_returns_keys(self):
        b = adc_noise_budget()
        expected = {'sqnr_ideal', 'thermal_snr', 'jitter_snr', 'total_snr',
                    'enob', 'limiting_source'}
        assert set(b.keys()) == expected

    def test_enob_less_than_ideal(self):
        b = adc_noise_budget(bits=12, R_source=100e3, sigma_jitter=100e-12)
        assert b['enob'] < 12.0

    def test_low_source_r_thermal_dominant_at_high_bits(self):
        # At 16 bits with 10k source, thermal noise begins to contribute:
        # the fig-noise-budget caption claims ~0.15 dB, so pin that quantity.
        # Quantisation still dominates at 16 bits; thermal becomes limiting past 18 bits.
        b = adc_noise_budget(bits=16, R_source=10000, sigma_jitter=1e-15)
        thermal_contribution_db = b['sqnr_ideal'] - b['total_snr']
        assert 0.1 < thermal_contribution_db < 0.2, (
            f"expected ~0.15 dB thermal contribution, got {thermal_contribution_db:.3f}")


class TestSigmaDeltaModulator:
    def _inband_snr(self, bits, f_sig, osr, fs=1.0):
        from scipy.signal import welch
        f, psd = welch(bits, fs, nperseg=8192)
        band = f <= fs / (2 * osr)
        sig = np.abs(f - f_sig) <= 2 * (f[1] - f[0])
        return 10 * np.log10(psd[sig & band].sum() / psd[band & ~sig].sum())

    def test_output_is_one_bit(self):
        y = sigma_delta_modulate(np.zeros(100))
        assert set(np.unique(y)) <= {-1.0, 1.0}

    def test_dc_tracking_first_order(self):
        # The bitstream mean must converge to the DC input.
        for dc in (-0.5, 0.0, 0.3, 0.7):
            y = sigma_delta_modulate(np.full(20000, dc))
            assert abs(np.mean(y) - dc) < 0.01, f"dc={dc}"

    def test_dc_tracking_second_order(self):
        for dc in (-0.5, 0.3):
            y = sigma_delta_modulate(np.full(20000, dc), order=2)
            assert abs(np.mean(y) - dc) < 0.01, f"dc={dc}"

    def test_noise_is_shaped_high(self):
        # NTF is a differentiator: high-band PSD must dwarf low-band PSD.
        from scipy.signal import welch
        t = np.arange(1 << 16)
        bits = sigma_delta_modulate(0.5 * np.sin(2 * np.pi * t / 1024))
        f, psd = welch(bits, 1.0, nperseg=8192)
        sig = np.abs(f - 1 / 1024) <= 3 * (f[1] - f[0])
        lo = psd[(f > 0) & (f < 0.01) & ~sig].mean()
        hi = psd[f > 0.4].mean()
        assert 10 * np.log10(hi / lo) > 20, "expected >20 dB low-to-high tilt"

    def test_snr_grows_with_osr_first_order(self):
        # Theory: ~9 dB per octave of OSR for L=1 (measured 6.6-8.3).
        t = np.arange(1 << 17)
        bits = sigma_delta_modulate(0.5 * np.sin(2 * np.pi * t / 1024))
        snr32 = self._inband_snr(bits, 1 / 1024, 32)
        snr128 = self._inband_snr(bits, 1 / 1024, 128)
        assert 12 < (snr128 - snr32) < 24, "expected ~9 dB/octave over 2 octaves"

    def test_second_order_beats_first(self):
        t = np.arange(1 << 17)
        x = 0.5 * np.sin(2 * np.pi * t / 1024)
        snr1 = self._inband_snr(sigma_delta_modulate(x, order=1), 1 / 1024, 64)
        snr2 = self._inband_snr(sigma_delta_modulate(x, order=2), 1 / 1024, 64)
        assert snr2 > snr1 + 10, f"L=2 ({snr2:.1f} dB) should beat L=1 ({snr1:.1f} dB) clearly"

    def test_decimated_output_recovers_sine(self):
        osr = 64
        t = np.arange(1 << 16)
        x = 0.5 * np.sin(2 * np.pi * t / 2048)
        y = sigma_delta_decimate(sigma_delta_modulate(x), osr)
        assert len(y) == len(x) // osr
        # Compare against the ideal input decimated the same way (skip edges).
        ref = x[::osr]
        err = np.sqrt(np.mean((y[4:-4] - ref[4:-4]) ** 2))
        assert err < 0.05, f"decimated RMS error {err:.3f}"

    def test_invalid_order_raises(self):
        with pytest.raises(ValueError):
            sigma_delta_modulate(np.zeros(8), order=3)


class TestDDSSine:
    NPS = 8192
    # On-bin tuning word (multiple of 2^24/8192) with active truncated bits.
    F_WORD = (1 << 17) + 2048

    def _spurs(self, mode, seed=1):
        from scipy.signal import welch
        x = dds_sine(1 << 15, self.F_WORD, phase_bits=24, table_bits=8,
                     mode=mode, seed=seed)
        f, psd = welch(x, 1.0, nperseg=self.NPS, window='boxcar')
        f0 = self.F_WORD / 2 ** 24
        sig = np.abs(f - f0) <= 1.5 * (f[1] - f[0])
        carrier = psd[sig].max()
        rest_f = f[~sig & (f > 0)]
        rest = psd[~sig & (f > 0)]
        worst = 10 * np.log10(rest.max() / carrier)
        worst_inband = 10 * np.log10(rest[rest_f < 0.05].max() / carrier)
        return worst, worst_inband

    def test_output_range(self):
        x = dds_sine(1000, self.F_WORD)
        assert np.all(np.abs(x) <= 1.0)

    def test_truncation_spur_near_minus_6dB_per_table_bit(self):
        # Classic worst-case rule: ~ -6.02*W dBc for a W-bit table address
        # (measured -48.1 dBc for W=8; rule predicts -48.2).
        worst, _ = self._spurs('truncate')
        assert -50 < worst < -46, f"expected ~-48 dBc truncation spur, got {worst:.1f}"

    def test_dither_whitens_spurs(self):
        _, inband_trunc = self._spurs('truncate')
        worst_dith, _ = self._spurs('dither')
        assert worst_dith < -65, f"dithered worst spur {worst_dith:.1f} dBc"
        assert worst_dith < inband_trunc - 15, "dither should bury the truncation spurs"

    def test_noise_shaping_clears_the_band(self):
        worst, inband = self._spurs('noise-shaped')
        assert inband < -65, f"in-band spur {inband:.1f} dBc"
        assert worst > inband + 10, "shaped error should sit out of band, not in it"

    def test_invalid_mode_raises(self):
        with pytest.raises(ValueError):
            dds_sine(100, self.F_WORD, mode='wrong')