embedded-dsp 0.5.1

//! Frequency-response, group-delay, and pole-based stability analysis for filters
//! produced by [`crate::filter_design`] or hand-written FIR/biquad coefficients.
//!
//! These routines evaluate the DTFT `H(e^{jω})` of a coefficient set directly (no FFT
//! required), so a filter design can be inspected at arbitrary frequencies before it is
//! deployed to a real-time processing path.

#[allow(unused_imports)]
use crate::math::FloatMath;
use crate::types::Complex;

// --- Frequency Response (DTFT evaluation, H(e^{jw})) ---

/// Evaluates the frequency response `H(e^{jω}) = Σ h[k] e^{-jkω}` of an FIR filter (or any
/// raw coefficient sequence) at a single normalized frequency `freq_norm` (cycles/sample,
/// `0.0..=0.5`, where `0.5` is Nyquist).
pub fn fir_frequency_response(taps: &[f32], freq_norm: f32) -> Complex<f32> {
    let omega = 2.0 * core::f32::consts::PI * freq_norm;
    let mut real = 0.0f32;
    let mut imag = 0.0f32;
    for (k, &tap) in taps.iter().enumerate() {
        let angle = omega * k as f32;
        real += tap * angle.cos();
        imag -= tap * angle.sin();
    }
    Complex::new(real, imag)
}

/// Evaluates the frequency response `H(e^{jω})` of a single Direct Form I biquad section
/// `[b0, b1, b2, a1, a2]` (as produced by [`crate::filter_design`] and consumed by
/// [`crate::filtering::biquad_cascade_df1_f32`], where `y(n) = b0 x(n) + b1 x(n-1) + b2 x(n-2)
/// + a1 y(n-1) + a2 y(n-2)`) at a single normalized frequency `freq_norm` (cycles/sample,
/// `0.0..=0.5`).
pub fn biquad_frequency_response(coeffs: &[f32; 5], freq_norm: f32) -> Complex<f32> {
    let omega = 2.0 * core::f32::consts::PI * freq_norm;
    let cos1 = omega.cos();
    let sin1 = omega.sin();
    let cos2 = (2.0 * omega).cos();
    let sin2 = (2.0 * omega).sin();

    let num = Complex::new(
        coeffs[0] + coeffs[1] * cos1 + coeffs[2] * cos2,
        -coeffs[1] * sin1 - coeffs[2] * sin2,
    );
    // Denominator of H(z) = 1 - a1*z^-1 - a2*z^-2, matching the recurrence's sign convention.
    let den = Complex::new(
        1.0 - coeffs[3] * cos1 - coeffs[4] * cos2,
        coeffs[3] * sin1 + coeffs[4] * sin2,
    );
    complex_divide(num, den)
}

/// Evaluates the combined frequency response of a cascade of Direct Form I biquad sections
/// (`coeffs.len()` must be a multiple of 5, as produced by e.g.
/// [`crate::filter_design::butterworth_lowpass_biquads`]) at a single normalized frequency
/// `freq_norm` (cycles/sample, `0.0..=0.5`).
pub fn biquad_cascade_frequency_response(coeffs: &[f32], freq_norm: f32) -> Complex<f32> {
    let mut total = Complex::new(1.0f32, 0.0f32);
    for stage in coeffs.chunks_exact(5) {
        let section: [f32; 5] = [stage[0], stage[1], stage[2], stage[3], stage[4]];
        total = complex_multiply(total, biquad_frequency_response(&section, freq_norm));
    }
    total
}

/// Returns the linear magnitude `|H(e^{jω})|` of a complex frequency-response value.
pub fn response_magnitude(h: Complex<f32>) -> f32 {
    (h.real * h.real + h.imag * h.imag).sqrt()
}

/// Returns the magnitude of a complex frequency-response value in decibels: `20 * log10(|H|)`.
pub fn response_magnitude_db(h: Complex<f32>) -> f32 {
    20.0 * response_magnitude(h).max(1e-20).log10()
}

/// Returns the phase (argument) of a complex frequency-response value, in radians, wrapped to
/// `(-π, π]`.
pub fn response_phase(h: Complex<f32>) -> f32 {
    h.imag.atan2(h.real)
}

fn complex_multiply(a: Complex<f32>, b: Complex<f32>) -> Complex<f32> {
    Complex::new(
        a.real * b.real - a.imag * b.imag,
        a.real * b.imag + a.imag * b.real,
    )
}

fn complex_divide(a: Complex<f32>, b: Complex<f32>) -> Complex<f32> {
    let denom = b.real * b.real + b.imag * b.imag;
    if denom < 1e-20 {
        return Complex::new(0.0, 0.0);
    }
    let inv_denom = 1.0 / denom;
    Complex::new(
        (a.real * b.real + a.imag * b.imag) * inv_denom,
        (a.imag * b.real - a.real * b.imag) * inv_denom,
    )
}

// --- Group Delay (Discrete-Time Fourier Transform, linear-phase analysis) ---

/// Computes the group delay (in samples), `τ(ω) = Re[B(e^{jω}) / H(e^{jω})]` where
/// `B(e^{jω}) = Σ k·h[k]·e^{-jkω}`, of an FIR filter at a single normalized frequency
/// `freq_norm` (cycles/sample, `0.0..=0.5`). For a linear-phase (symmetric) FIR of length `M`,
/// this is constant and equal to `(M - 1) / 2` at every frequency.
pub fn fir_group_delay(taps: &[f32], freq_norm: f32) -> f32 {
    let omega = 2.0 * core::f32::consts::PI * freq_norm;
    let mut h_re = 0.0f32;
    let mut h_im = 0.0f32;
    let mut b_re = 0.0f32;
    let mut b_im = 0.0f32;
    for (k, &tap) in taps.iter().enumerate() {
        let n = k as f32;
        let angle = omega * n;
        let c = angle.cos();
        let s = angle.sin();
        h_re += tap * c;
        h_im -= tap * s;
        b_re += n * tap * c;
        b_im -= n * tap * s;
    }
    let denom = h_re * h_re + h_im * h_im;
    if denom < 1e-20 {
        return 0.0;
    }
    (b_re * h_re + b_im * h_im) / denom
}

// --- Pole-Based Stability Analysis (Z-transform) ---

/// Computes the pole radius (largest pole magnitude on the z-plane) of a single Direct Form I
/// biquad section `[b0, b1, b2, a1, a2]`, whose poles are the roots of
/// `z^2 - a1*z - a2 = 0`. A causal LTI system is stable if and only if all poles lie strictly
/// inside the unit circle (`pole_radius < 1.0`).
pub fn biquad_pole_radius(coeffs: &[f32; 5]) -> f32 {
    let a1 = coeffs[3];
    let a2 = coeffs[4];
    let discriminant = a1 * a1 + 4.0 * a2;
    if discriminant >= 0.0 {
        let sqrt_d = discriminant.sqrt();
        let p1 = (a1 + sqrt_d) / 2.0;
        let p2 = (a1 - sqrt_d) / 2.0;
        p1.abs().max(p2.abs())
    } else {
        // Complex-conjugate pole pair: |pole|^2 equals the product of the roots, -a2.
        (-a2).sqrt()
    }
}

/// Returns `true` if the single biquad section `[b0, b1, b2, a1, a2]` is stable, i.e. both
/// poles lie strictly inside the unit circle.
pub fn biquad_is_stable(coeffs: &[f32; 5]) -> bool {
    biquad_pole_radius(coeffs) < 1.0
}

/// Returns `true` if every stage of a biquad cascade (`coeffs.len()` a multiple of 5) is
/// stable.
pub fn biquad_cascade_is_stable(coeffs: &[f32]) -> bool {
    coeffs
        .chunks_exact(5)
        .all(|stage| biquad_is_stable(&[stage[0], stage[1], stage[2], stage[3], stage[4]]))
}

// ─────────────────────────────────────────────────────────────────────────────
// Quantization, Headroom, and SQNR Analysis
// ─────────────────────────────────────────────────────────────────────────────

use crate::types::q15;

/// Computes the peak frequency response gain `||H(e^{jω})||_∞` of a biquad section.
pub fn biquad_peak_gain(coeffs: &[f32; 5], num_points: usize) -> f32 {
    let pts = num_points.max(16);
    let mut max_mag = 0.0f32;
    for i in 0..=pts {
        let f = (i as f32) * 0.5 / (pts as f32);
        let resp = biquad_frequency_response(coeffs, f);
        let mag = (resp.real * resp.real + resp.imag * resp.imag).sqrt();
        if mag > max_mag {
            max_mag = mag;
        }
    }
    max_mag
}

/// Computes the L2-norm energy `||H(e^{jω})||_2` of a biquad section.
pub fn biquad_l2_norm(coeffs: &[f32; 5], num_points: usize) -> f32 {
    let pts = num_points.max(16);
    let mut sum_sq = 0.0f32;
    for i in 0..=pts {
        let f = (i as f32) * 0.5 / (pts as f32);
        let resp = biquad_frequency_response(coeffs, f);
        sum_sq += resp.real * resp.real + resp.imag * resp.imag;
    }
    (sum_sq / (pts as f32 + 1.0)).sqrt()
}

/// Estimates required integer headroom bits and peak gain for a biquad section.
///
/// Returns `(headroom_bits, peak_gain)`.
/// `headroom_bits` is the number of bits required above unity (`ceil(log2(max(1.0, peak_gain)))`).
pub fn estimate_biquad_headroom_bits(coeffs: &[f32; 5]) -> (u8, f32) {
    let peak = biquad_peak_gain(coeffs, 64);
    if peak <= 1.0 {
        (0, peak)
    } else {
        // Calculate ceil(log2(peak))
        let mut bits = 0u8;
        let mut threshold = 1.0f32;
        while threshold < peak && bits < 14 {
            bits += 1;
            threshold *= 2.0;
        }
        (bits, peak)
    }
}

/// Evaluates the frequency response of a Q15 quantized biquad section.
pub fn biquad_q15_frequency_response(
    coeffs_q15: &[q15; 5],
    post_shift: u8,
    freq_norm: f32,
) -> Complex<f32> {
    let scale = (1u32 << post_shift.min(14)) as f32;
    let b0 = coeffs_q15[0].to_num::<f32>() * scale;
    let b1 = coeffs_q15[1].to_num::<f32>() * scale;
    let b2 = coeffs_q15[2].to_num::<f32>() * scale;
    let a1 = coeffs_q15[3].to_num::<f32>() * scale;
    let a2 = coeffs_q15[4].to_num::<f32>() * scale;

    let float_coeffs = [b0, b1, b2, a1, a2];
    biquad_frequency_response(&float_coeffs, freq_norm)
}

/// Computes the Signal-to-Quantization-Noise Ratio (SQNR in dB) between an ideal floating-point
/// biquad cascade and its Q15 quantized equivalent.
pub fn biquad_quantization_snr_db(
    sos_f32: &[f32],
    sos_q15: &[q15],
    post_shift: u8,
    num_points: usize,
) -> f32 {
    if sos_f32.len() != sos_q15.len() || sos_f32.is_empty() || sos_f32.len() % 5 != 0 {
        return 0.0;
    }

    let num_stages = sos_f32.len() / 5;
    let pts = num_points.max(32);
    let mut sig_pow = 0.0f32;
    let mut err_pow = 0.0f32;

    for i in 0..=pts {
        let f = (i as f32) * 0.5 / (pts as f32);

        // Ideal response
        let mut h_ideal = Complex::new(1.0f32, 0.0f32);
        for stage in 0..num_stages {
            let idx = stage * 5;
            let section = [
                sos_f32[idx],
                sos_f32[idx + 1],
                sos_f32[idx + 2],
                sos_f32[idx + 3],
                sos_f32[idx + 4],
            ];
            h_ideal = h_ideal * biquad_frequency_response(&section, f);
        }

        // Quantized response
        let mut h_quant = Complex::new(1.0f32, 0.0f32);
        for stage in 0..num_stages {
            let idx = stage * 5;
            let section = [
                sos_q15[idx],
                sos_q15[idx + 1],
                sos_q15[idx + 2],
                sos_q15[idx + 3],
                sos_q15[idx + 4],
            ];
            h_quant = h_quant * biquad_q15_frequency_response(&section, post_shift, f);
        }

        let mag_sq = h_ideal.real * h_ideal.real + h_ideal.imag * h_ideal.imag;
        let diff_re = h_ideal.real - h_quant.real;
        let diff_im = h_ideal.imag - h_quant.imag;
        let err_sq = diff_re * diff_re + diff_im * diff_im;

        sig_pow += mag_sq;
        err_pow += err_sq;
    }

    if err_pow < 1e-20 {
        return 120.0; // Near perfect representation
    }
    10.0 * (sig_pow / err_pow).log10()
}

/// Computes the SQNR (in dB) between an ideal floating-point FIR filter and its Q15 quantized version.
pub fn fir_quantization_snr_db(taps_f32: &[f32], taps_q15: &[q15], num_points: usize) -> f32 {
    if taps_f32.len() != taps_q15.len() || taps_f32.is_empty() {
        return 0.0;
    }

    let pts = num_points.max(32);
    let mut sig_pow = 0.0f32;
    let mut err_pow = 0.0f32;

    for i in 0..=pts {
        let f = (i as f32) * 0.5 / (pts as f32);
        let h_ideal = fir_frequency_response(taps_f32, f);

        // Quantized FIR response (scale taps by 1/32768)
        let omega = 2.0 * core::f32::consts::PI * f;
        let mut q_re = 0.0f32;
        let mut q_im = 0.0f32;
        for (k, &tap) in taps_q15.iter().enumerate() {
            let tap_f = tap.to_num::<f32>();
            let angle = omega * k as f32;
            q_re += tap_f * angle.cos();
            q_im -= tap_f * angle.sin();
        }

        let mag_sq = h_ideal.real * h_ideal.real + h_ideal.imag * h_ideal.imag;
        let diff_re = h_ideal.real - q_re;
        let diff_im = h_ideal.imag - q_im;
        let err_sq = diff_re * diff_re + diff_im * diff_im;

        sig_pow += mag_sq;
        err_pow += err_sq;
    }

    if err_pow < 1e-20 {
        return 120.0;
    }
    10.0 * (sig_pow / err_pow).log10()
}