"""
ADC noise: beyond the 6.02·B + 1.76 dB formula.

The ideal ADC has only quantisation noise.  Real ADCs add thermal noise,
aperture jitter, INL/DNL non-linearity, and reference noise.  This module
models the major non-ideal noise sources and computes the effective number
of bits (ENOB) from measured data.

Functions
---------
* jitter_snr(f_in, sigma_jitter): SNR limited by aperture jitter
* enob_from_sinad(sinad_db): effective bits from SINAD
* oversampling_gain(osr): SNR improvement from oversampling (white noise)
* sigma_delta_snr(M, osr, order): ideal sigma-delta SQNR
* adc_noise_budget(bits, fs, f_in, v_ref, R_source, T): full noise budget

References
----------
* Kester, W. (2009). *The Data Conversion Handbook*. Analog Devices.
* Maloberti, F. (2007). *Data Converters*. Springer.
"""

import numpy as np


# Physical constants
K_B = 1.380649e-23    # Boltzmann constant [J/K]
Q_E = 1.602176634e-19  # Electron charge [C]


def jitter_snr(f_in: float, sigma_jitter: float) -> float:
    """SNR limited by aperture jitter.

    Aperture jitter causes the sampling instant to vary randomly.  For a
    sinusoidal input at frequency f_in, the resulting SNR is

        SNR_jitter = -20 * log10(2 * pi * f_in * sigma_jitter)   [dB]

    Parameters
    ----------
    f_in : float
        Input frequency [Hz].
    sigma_jitter : float
        RMS aperture jitter [s].

    Returns
    -------
    snr_db : float
        Jitter-limited SNR [dB].
    """
    return -20 * np.log10(2 * np.pi * f_in * sigma_jitter)


def enob_from_sinad(sinad_db: float) -> float:
    """Effective number of bits from SINAD.

        ENOB = (SINAD_dB - 1.76) / 6.02

    SINAD (signal-to-noise-and-distortion) includes quantisation noise,
    thermal noise, jitter, and harmonic distortion.  ENOB is the equivalent
    ideal-ADC resolution that would give the same SINAD.

    Parameters
    ----------
    sinad_db : float
        Measured SINAD [dB].

    Returns
    -------
    enob : float
        Effective number of bits.
    """
    return float((sinad_db - 1.76) / 6.02)


def oversampling_gain(osr: int) -> float:
    """SNR improvement from oversampling with white quantisation noise.

    Each factor-of-4 oversampling buys 1 bit (6.02 dB) of effective
    resolution, because the same total quantisation noise power is spread
    over a wider bandwidth and the in-band portion decreases proportionally.

        gain = 10 * log10(OSR)   [dB]

    This assumes white quantisation noise.  If the noise is coloured
    (e.g., 1/f), the gain is less.

    Parameters
    ----------
    osr : int
        Oversampling ratio (f_sample / (2 * f_signal)).

    Returns
    -------
    gain_db : float
        SNR improvement [dB].
    """
    return 10 * np.log10(osr)


def sigma_delta_snr(M: int, osr: int, order: int = 1) -> float:
    """Ideal signal-to-quantisation-noise ratio for a sigma-delta modulator.

    An L-th order sigma-delta modulator shapes quantisation noise with a
    high-pass transfer function (1 - z^{-1})^L, pushing it out of the
    signal band.  The in-band noise power after decimation is

        SQNR ≈ 6.02·M + 1.76 - 10·log10(pi^(2L)/(2L+1)) + (2L+1)·10·log10(OSR)

    for a full-scale sinusoidal input.  This is the ideal (linearised) model;
    real modulators have additional noise from integrator non-idealities and
    clock jitter.

    Parameters
    ----------
    M : int
        Quantiser resolution in the loop (typically 1 bit: M=1).
    osr : int
        Oversampling ratio.
    order : int
        Loop filter order L (1 or 2 for most practical designs).

    Returns
    -------
    sqnr_db : float
        Ideal SQNR after decimation [dB].
    """
    L = order
    # Ideal SQNR for an L-th order modulator.
    sqnr = 6.02 * M + 1.76
    sqnr -= 10 * np.log10(np.pi**(2*L) / (2*L + 1))
    sqnr += (2*L + 1) * 10 * np.log10(osr)
    return float(sqnr)


def adc_noise_budget(bits: int = 12, fs: float = 100e3, f_in: float = 10e3,
                     v_ref: float = 3.3, R_source: float = 1000.0,
                     T: float = 293.0, sigma_jitter: float = 10e-12) -> dict:
    """Compute the full ADC noise budget.

    Sums the major noise sources (quantisation, thermal, and jitter)
    in quadrature and returns the total SNR, the limiting source, and
    the effective number of bits.

    Parameters
    ----------
    bits : int
        ADC resolution.
    fs : float
        Sample rate [Hz].
    f_in : float
        Input frequency of interest [Hz].
    v_ref : float
        Reference voltage (full-scale peak-to-peak) [V].
    R_source : float
        Source resistance [ohms].
    T : float
        Temperature [K].
    sigma_jitter : float
        RMS aperture jitter [s].

    Returns
    -------
    budget : dict
        Keys: 'sqnr_ideal', 'thermal_snr', 'jitter_snr', 'total_snr',
        'enob', 'limiting_source'.
    """
    Q = v_ref / (2**bits - 1)
    # Ideal quantisation SNR (full-scale sine).
    sqnr_ideal = 6.02 * bits + 1.76

    # Thermal noise in the source resistance over the Nyquist bandwidth.
    vn_thermal = np.sqrt(4 * K_B * T * R_source * fs / 2)
    # Full-scale sine RMS = v_ref/(2*sqrt(2)).
    v_signal_rms = v_ref / (2 * np.sqrt(2))
    thermal_snr = 20 * np.log10(v_signal_rms / vn_thermal) if vn_thermal > 0 else np.inf

    # Jitter SNR.
    jsnr = jitter_snr(f_in, sigma_jitter)

    # Total: noise powers add (variances sum).
    p_signal = v_signal_rms**2
    p_quant = (Q**2 / 12)
    p_thermal = vn_thermal**2
    p_jitter = p_signal * 10**(-jsnr / 10) if jsnr < 300 else 0.0

    p_total_noise = p_quant + p_thermal + p_jitter
    total_snr = float(10 * np.log10(p_signal / p_total_noise)) if p_total_noise > 0 else np.inf

    # Identify the limiting source.
    noises = {'quantisation': p_quant, 'thermal': p_thermal, 'jitter': p_jitter}
    limiting = max(noises, key=noises.get)

    return {
        'sqnr_ideal': float(sqnr_ideal),
        'thermal_snr': float(thermal_snr),
        'jitter_snr': float(jsnr),
        'total_snr': float(total_snr),
        'enob': enob_from_sinad(total_snr),
        'limiting_source': limiting,
    }


# ---------------------------------------------------------------------------
# Sigma-delta modulator (the loop itself, not just the SQNR formula)
# ---------------------------------------------------------------------------

def sigma_delta_modulate(x: np.ndarray, order: int = 1) -> np.ndarray:
    """Run a discrete-time 1-bit sigma-delta modulator over *x*.

    Implements the delaying-integrator loop with the actual nonlinear
    sign() quantiser.  The first-order form is due to Inose, Yasuda &
    Murakami's original delta-sigma paper; the second-order form is the
    standard later generalisation:

    * order 1:  y[n] = sign(s[n]),      s[n+1] = s[n] + x[n] - y[n]
    * order 2:  y[n] = sign(s2[n]),     s2[n+1] = s2[n] + s1[n] - 2 y[n],
                                        s1[n+1] = s1[n] + x[n] - y[n]

    Linearising the quantiser as additive noise e[n] gives

        Y(z) = z^{-L} X(z) + (1 - z^{-1})^L E(z)

    i.e. the signal passes with a pure delay (STF = z^-L) while the
    quantisation noise is differentiated L times (NTF = (1-z^-1)^L),
    pushing its power toward high frequencies where the decimation
    filter removes it.

    Parameters
    ----------
    x : ndarray
        Input signal.  Keep |x| well below 1: the 1-bit feedback can
        only represent +-1, and the second-order loop loses stability
        as the input approaches full scale (keep |x| <= ~0.7).
    order : int
        Loop order, 1 or 2.

    Returns
    -------
    y : ndarray of +-1.0
        The 1-bit output stream at the oversampled rate.
    """
    x = np.asarray(x, dtype=float)
    y = np.empty(len(x))
    if order == 1:
        s = 0.0
        for i, xi in enumerate(x):
            yi = 1.0 if s >= 0.0 else -1.0
            y[i] = yi
            s += xi - yi
    elif order == 2:
        s1 = 0.0
        s2 = 0.0
        for i, xi in enumerate(x):
            yi = 1.0 if s2 >= 0.0 else -1.0
            y[i] = yi
            s2 += s1 - 2.0 * yi     # uses s1[n]: update s2 before s1
            s1 += xi - yi
    else:
        raise ValueError(f'order must be 1 or 2, got {order}')
    return y


def sigma_delta_decimate(bits: np.ndarray, osr: int) -> np.ndarray:
    """Decimate a sigma-delta bitstream by *osr* with a sinc^2 filter.

    Two cascaded length-*osr* moving averages (a CIC-2 without the
    recursive implementation) followed by downsampling.  Good enough to
    demonstrate the resolution recovery; a production decimator would
    use a proper CIC + compensation FIR (see basics/10-multirate).

    Parameters
    ----------
    bits : ndarray of +-1.0
        Modulator output stream.
    osr : int
        Oversampling ratio (decimation factor).

    Returns
    -------
    y : ndarray of length len(bits)//osr
        The decimated, multi-bit output.
    """
    kernel = np.ones(osr) / osr
    smooth = np.convolve(np.convolve(bits, kernel, mode='same'),
                         kernel, mode='same')
    return smooth[::osr]


# ---------------------------------------------------------------------------
# DDS phase truncation: the same noise-shaping trick on the phase path
# ---------------------------------------------------------------------------

def dds_sine(n: int, f_word: int, phase_bits: int = 24, table_bits: int = 8,
             mode: str = 'truncate', seed: int | None = None) -> np.ndarray:
    """Direct digital synthesis of a sine with phase-accumulator truncation.

    A *phase_bits*-wide accumulator advances by *f_word* each sample; the
    top *table_bits* bits address a 2**table_bits-entry sine table.  The
    discarded low bits are a periodic phase error that produces spurs
    (Nicholas & Samueli).  Three strategies for the discarded bits:

    * ``'truncate'``: drop them (periodic error -> discrete spurs).
    * ``'dither'``: add uniform random dither on [0, 2**shift) before
      truncating (spurs whiten into a raised noise floor).
    * ``'noise-shaped'``: first-order error feedback: the truncation
      residual is added back into the next sample's phase before
      truncation, so the phase error is shaped by (1 - z^-1), the same
      differentiator NTF as the first-order sigma-delta.

    Parameters
    ----------
    n : int
        Number of output samples.
    f_word : int
        Frequency tuning word; f_out = f_word / 2**phase_bits * f_s.
    phase_bits, table_bits : int
        Accumulator width and table address width (phase_bits > table_bits).
    mode : str
        'truncate', 'dither', or 'noise-shaped'.
    seed : int or None
        RNG seed (used only for 'dither').

    Returns
    -------
    x : ndarray of shape (n,)
        Sine samples in [-1, 1] from the lookup table.
    """
    if mode not in ('truncate', 'dither', 'noise-shaped'):
        raise ValueError(f"unknown mode {mode!r}")
    table_len = 1 << table_bits
    table = np.sin(2 * np.pi * np.arange(table_len) / table_len)
    shift = phase_bits - table_bits
    mask = (1 << phase_bits) - 1
    rng = np.random.default_rng(seed)

    out = np.empty(n)
    acc = 0
    err = 0
    for i in range(n):
        p = acc
        if mode == 'dither':
            p = (acc + int(rng.integers(0, 1 << shift))) & mask
        elif mode == 'noise-shaped':
            p = (acc + err) & mask
        idx = p >> shift
        if mode == 'noise-shaped':
            err = p & ((1 << shift) - 1)   # residual fed back next sample
        out[i] = table[idx]
        acc = (acc + f_word) & mask
    return out
