"""
Dither: deliberately adding noise to improve quantisation.

Undithered truncation creates signal-correlated quantisation error:
harmonics and intermodulation that are perceptually objectionable and
mathematically awkward.  Adding a small amount of noise before quantising
decorrelates the error from the signal, trading a modest increase in the
noise floor for a clean, distortion-free spectrum.

Dither types
------------
* RPDF (rectangular PDF): uniform dither on [-Δ/2, Δ/2].  Eliminates
  distortion but leaves noise modulation (the noise floor varies with
  signal amplitude).
* TPDF (triangular PDF): the sum of two independent RPDF sources.  Eliminates
  both distortion AND noise modulation.  This is the standard for audio.

References
----------
* Vanderkooy, J. & Lipshitz, S. P. (1984). Resolution below the least
  significant bit in digital systems with dither. *J. Audio Eng. Soc.*,
  32(3), 106-115.
* Wannamaker, R. A., Lipshitz, S. P., Vanderkooy, J., & Wright, J. N. (2000).
  A theory of nonsubtractive dither. *IEEE Trans. Signal Process.*, 48(2),
  499-516.
* Pohlmann, K. C. (2011). *Principles of Digital Audio* (6th ed.).
  McGraw-Hill.
"""

import numpy as np


def quantize(x: np.ndarray, bits: int, v_max: float = 1.0) -> np.ndarray:
    """Quantize floating-point samples to *bits* resolution.

    Maps the continuous range [-v_max, v_max] to 2^bits levels via
    mid-tread rounding.  The quantisation step is
        Q = 2 * v_max / (2^bits - 1).

    Parameters
    ----------
    x : ndarray
        Input samples in [-v_max, v_max].
    bits : int
        Number of bits (e.g., 16).
    v_max : float
        Full-scale amplitude.

    Returns
    -------
    xq : ndarray
        Quantised samples, same shape as x.
    """
    Q = 2 * v_max / (2**bits - 1)
    x_clipped = np.clip(x, -v_max, v_max)
    xq = np.round(x_clipped / Q) * Q
    return np.clip(xq, -v_max, v_max)


def rpdf_dither(n: int, Q: float, seed: int | None = None) -> np.ndarray:
    """Generate RPDF (rectangular) dither, uniform on [-Q/2, Q/2].

    Parameters
    ----------
    n : int
        Number of samples.
    Q : float
        Quantisation step (LSB size).
    seed : int or None
        RNG seed.

    Returns
    -------
    d : ndarray of shape (n,)
        RPDF dither samples.
    """
    rng = np.random.default_rng(seed)
    return rng.uniform(-Q/2, Q/2, size=n)


def tpdf_dither(n: int, Q: float, seed: int | None = None) -> np.ndarray:
    """Generate TPDF (triangular) dither, two independent RPDFs summed.

    TPDF dither eliminates both quantisation distortion and noise
    modulation.  It is the standard for professional audio.

    Parameters
    ----------
    n : int
        Number of samples.
    Q : float
        Quantisation step.
    seed : int or None
        RNG seed.

    Returns
    -------
    d : ndarray of shape (n,)
        TPDF dither samples on [-Q, Q].
    """
    rng = np.random.default_rng(seed)
    d1 = rng.uniform(-Q/2, Q/2, size=n)
    d2 = rng.uniform(-Q/2, Q/2, size=n)
    return d1 + d2


def quantize_with_dither(x: np.ndarray, bits: int, v_max: float = 1.0,
                         dither_type: str = 'tpdf',
                         seed: int | None = None) -> np.ndarray:
    """Quantize with dither added before rounding.

    Parameters
    ----------
    x : ndarray
        Input samples.
    bits : int
        Quantisation resolution.
    v_max : float
        Full-scale amplitude.
    dither_type : str
        'rpdf' or 'tpdf'.
    seed : int or None
        RNG seed.

    Returns
    -------
    xq : ndarray
        Dithered and quantised samples.
    """
    Q = 2 * v_max / (2**bits - 1)
    if dither_type == 'tpdf':
        d = tpdf_dither(len(x), Q, seed)
    elif dither_type == 'rpdf':
        d = rpdf_dither(len(x), Q, seed)
    else:
        raise ValueError(f"Unknown dither type: {dither_type}")
    return quantize(x + d, bits, v_max)


def measure_distortion_spectrum(x: np.ndarray, bits: int, v_max: float = 1.0,
                                dither_type: str | None = None,
                                seed: int | None = None) -> dict:
    """Quantize a sinusoid and measure the distortion spectrum.

    Returns the PSD of the quantization error so you can compare the
    harmonic structure of undithered vs dithered quantisation.

    Parameters
    ----------
    x : ndarray
        Input signal (should be a pure sinusoid near full-scale).
    bits : int
        Quantisation resolution.
    v_max : float
        Full-scale amplitude.
    dither_type : str or None
        None = undithered, 'rpdf', 'tpdf'.
    seed : int or None
        RNG seed.

    Returns
    -------
    result : dict
        Keys: 'f' (frequency bins), 'psd_error' (PSD of quantisation error),
        'sqnr_db' (signal-to-quantisation-noise ratio in dB).
    """
    if dither_type:
        xq = quantize_with_dither(x, bits, v_max, dither_type, seed)
    else:
        xq = quantize(x, bits, v_max)

    error = x - xq
    signal_power = np.mean(x**2)
    noise_power = np.mean(error**2)
    sqnr_db = 10 * np.log10(signal_power / noise_power) if noise_power > 0 else np.inf

    # PSD of the error signal.
    from scipy.signal import welch
    f, psd_e = welch(error, 1.0, nperseg=min(1024, len(x)//4))

    return {'f': f, 'psd_error': psd_e, 'sqnr_db': float(sqnr_db)}


def sqnr_with_dither(x: np.ndarray, bits: int, v_max: float = 1.0,
                     dither_type: str = 'tpdf',
                     seed: int | None = None) -> float:
    """Compute SQNR after dithered quantization.

    Parameters
    ----------
    x : ndarray
        Input signal.
    bits : int
        Quantisation resolution.
    v_max : float
        Full-scale amplitude.
    dither_type : str
        'rpdf' or 'tpdf'.
    seed : int or None
        RNG seed.

    Returns
    -------
    sqnr_db : float
        Signal-to-quantisation-noise ratio in dB.
    """
    res = measure_distortion_spectrum(x, bits, v_max, dither_type, seed)
    return res['sqnr_db']