"""Tests for dither.py: quantisation with and without dither."""

import numpy as np
import pytest
from dither import (quantize, rpdf_dither, tpdf_dither,
                     quantize_with_dither, measure_distortion_spectrum,
                     sqnr_with_dither)


class TestQuantize:
    def test_shape(self):
        x = np.random.default_rng(0).standard_normal(200)
        xq = quantize(x, bits=8)
        assert xq.shape == x.shape

    def test_fewer_unique_values(self):
        x = np.random.default_rng(1).standard_normal(1000)
        xq = quantize(x, bits=4)
        assert len(np.unique(xq)) < 30  # 4-bit = 16 levels max

    def test_clips_to_vmax(self):
        x = np.array([-2.0, 0.0, 2.0])
        xq = quantize(x, bits=8, v_max=1.0)
        assert np.max(np.abs(xq)) <= 1.0

    def test_mid_tread_zero_is_preserved(self):
        x = np.array([0.0])
        xq = quantize(x, bits=8)
        assert xq[0] == 0.0


class TestDitherGenerators:
    def test_rpdf_range(self):
        d = rpdf_dither(10000, Q=1.0, seed=42)
        assert np.min(d) >= -0.5
        assert np.max(d) <= 0.5

    def test_rpdf_uniform(self):
        d = rpdf_dither(100000, Q=1.0, seed=7)
        assert abs(np.mean(d)) < 0.01
        # Variance of uniform[-0.5, 0.5] = 1/12 ≈ 0.0833
        assert abs(np.var(d) - 1/12) < 0.005

    def test_tpdf_range(self):
        d = tpdf_dither(10000, Q=1.0, seed=42)
        assert np.min(d) >= -1.0
        assert np.max(d) <= 1.0

    def test_tpdf_triangular(self):
        d = tpdf_dither(100000, Q=1.0, seed=7)
        assert abs(np.mean(d)) < 0.02
        # Variance of TPDF (sum of two uniforms) = 2 * (1/12) = 1/6 ≈ 0.1667
        assert abs(np.var(d) - 1/6) < 0.01


class TestQuantizeWithDither:
    def test_shape(self):
        x = np.sin(2 * np.pi * 0.01 * np.arange(2000))
        xq = quantize_with_dither(x, bits=8, dither_type='tpdf', seed=1)
        assert xq.shape == x.shape

    def test_tpdf_gives_different_from_undithered(self):
        x = np.sin(2 * np.pi * 0.01 * np.arange(200))
        xq_undithered = quantize(x, bits=6)
        xq_dithered = quantize_with_dither(x, bits=6, dither_type='tpdf', seed=42)
        assert not np.array_equal(xq_undithered, xq_dithered)

    def test_reproducible_with_seed(self):
        x = np.sin(2 * np.pi * 0.01 * np.arange(100))
        a = quantize_with_dither(x, bits=6, dither_type='tpdf', seed=42)
        b = quantize_with_dither(x, bits=6, dither_type='tpdf', seed=42)
        assert np.array_equal(a, b)


class TestMeasureDistortion:
    def test_returns_keys(self):
        x = 0.9 * np.sin(2 * np.pi * 100 * np.arange(8192) / 8000)
        res = measure_distortion_spectrum(x, bits=8, dither_type='tpdf', seed=1)
        assert 'f' in res and 'psd_error' in res and 'sqnr_db' in res

    def test_undithered_low_bits_distorted(self):
        """At 4 bits, undithered quantisation creates strong harmonics."""
        n = 8192
        x = 0.9 * np.sin(2 * np.pi * 313 * np.arange(n) / 8000)  # non-coherent freq
        res = measure_distortion_spectrum(x, bits=4, dither_type=None)
        # SQNR should be well below the 6.02*4+1.76 = 25.8 dB theoretical
        # (the sine is not full-scale, and 4-bit undithered has strong harmonics)
        assert res['sqnr_db'] < 30

    def test_tpdf_reduces_psd_peaks(self):
        """TPDF dither should flatten the error spectrum vs undithered."""
        n = 8192
        x = 0.9 * np.sin(2 * np.pi * 313 * np.arange(n) / 8000)
        res_u = measure_distortion_spectrum(x, bits=6, dither_type=None, seed=1)
        res_d = measure_distortion_spectrum(x, bits=6, dither_type='tpdf', seed=1)
        # Peak-to-median ratio of error PSD: dithered should be flatter.
        pmr_u = np.max(res_u['psd_error'][10:]) / np.median(res_u['psd_error'][10:])
        pmr_d = np.max(res_d['psd_error'][10:]) / np.median(res_d['psd_error'][10:])
        assert pmr_d < pmr_u


class TestSQNR:
    def test_sqnr_approx_theory_high_bits(self):
        x = 0.99 * np.sin(2 * np.pi * 0.01 * np.arange(8192))
        sqnr = sqnr_with_dither(x, bits=12, dither_type='tpdf', seed=1)
        # TPDF adds ~3 dB extra noise vs undithered SQNR.
        # Undithered theory: 6.02*12+1.76 = 74.0 dB.
        # TPDF should be within ~10 dB of that.
        assert sqnr > 60