melpe-rs 0.1.2

/// All-pole IIR synthesis filter for MELPe
///
/// Implements H(z) = 1/A(z) where A(z) = 1 + a1*z^-1 + a2*z^-2 + ... + ap*z^-p
/// This is the core of LPC speech synthesis: excitation in → speech out.

use crate::core_types::LPC_ORDER;

/// All-pole synthesis filter with internal delay memory.
#[derive(Debug, Clone)]
pub struct SynthesisFilter {
    /// Filter delay line (past output samples)
    state: [f32; LPC_ORDER],
}

impl SynthesisFilter {
    pub fn new() -> Self {
        Self {
            state: [0.0; LPC_ORDER],
        }
    }

    /// Process a single sample through H(z) = 1/A(z).
    ///
    /// `excitation`: input sample (from mixed excitation generator)
    /// `a`: LPC coefficients a[0..LPC_ORDER] where a[i] corresponds to z^{-(i+1)}
    ///
    /// Output: y[n] = x[n] - a[0]*y[n-1] - a[1]*y[n-2] - ... - a[p-1]*y[n-p]
    #[inline]
    pub fn process_sample(&mut self, excitation: f32, a: &[f32; LPC_ORDER]) -> f32 {
        let mut y = excitation;
        for i in 0..LPC_ORDER {
            y -= a[i] * self.state[i];
        }

        // Shift delay line
        for i in (1..LPC_ORDER).rev() {
            self.state[i] = self.state[i - 1];
        }
        self.state[0] = y;

        y
    }

    /// Process a buffer of excitation samples with fixed LPC coefficients.
    /// Output is written to `output` (must be same length as `excitation`).
    pub fn process_buffer(
        &mut self,
        excitation: &[f32],
        a: &[f32; LPC_ORDER],
        output: &mut [f32],
    ) {
        let n = excitation.len().min(output.len());
        for i in 0..n {
            output[i] = self.process_sample(excitation[i], a);
        }
    }

    /// Process a buffer with linearly interpolated LPC coefficients.
    /// `a_start` and `a_end` are the LPC coefficients at the beginning and end
    /// of the buffer. Interpolation avoids clicks at frame boundaries.
    pub fn process_buffer_interp(
        &mut self,
        excitation: &[f32],
        a_start: &[f32; LPC_ORDER],
        a_end: &[f32; LPC_ORDER],
        output: &mut [f32],
    ) {
        let n = excitation.len().min(output.len());
        if n == 0 {
            return;
        }

        let mut a_interp = [0.0f32; LPC_ORDER];
        for i in 0..n {
            let t = i as f32 / (n - 1).max(1) as f32;
            for k in 0..LPC_ORDER {
                a_interp[k] = (1.0 - t) * a_start[k] + t * a_end[k];
            }
            output[i] = self.process_sample(excitation[i], &a_interp);
        }
    }

    /// Reset filter state to zero.
    pub fn reset(&mut self) {
        self.state = [0.0; LPC_ORDER];
    }

    /// Read current filter state (for debugging/testing).
    pub fn state(&self) -> &[f32; LPC_ORDER] {
        &self.state
    }
}

/// First-order de-emphasis filter: y[n] = x[n] + coeff * y[n-1]
/// Inverse of the pre-emphasis filter in the analysis path.
#[derive(Debug, Clone)]
pub struct DeEmphasis {
    prev: f32,
    coeff: f32,
}

impl DeEmphasis {
    pub fn new(coeff: f32) -> Self {
        Self { prev: 0.0, coeff }
    }

    /// Process a single sample.
    #[inline]
    pub fn process_sample(&mut self, x: f32) -> f32 {
        let y = x + self.coeff * self.prev;
        self.prev = y;
        y
    }

    /// Process a buffer in-place.
    pub fn process_inplace(&mut self, buf: &mut [f32]) {
        for sample in buf.iter_mut() {
            *sample = self.process_sample(*sample);
        }
    }

    pub fn reset(&mut self) {
        self.prev = 0.0;
    }
}

// ── Frame-level synthesis processor ──

use crate::core_types::{FrameParams, FRAME_SAMPLES, NUM_LSF};
use crate::lpc::lsf_to_lpc;
use crate::math::{powf, log10f, fabsf, fmaxf};

/// Pre-emphasis coefficient (must match the analysis side)
pub const PRE_EMPHASIS_COEFF: f32 = 0.97;

/// Frame-level synthesis processor.
///
/// Wraps the all-pole filter with:
/// - LSF → LPC conversion
/// - LPC coefficient interpolation across frames
/// - Gain scaling of excitation
/// - De-emphasis (inverse of analysis pre-emphasis)
///
/// The decoder calls `process_frame()` once per 22.5ms frame.
pub struct SynthesisFrameProcessor {
    filter: SynthesisFilter,
    deemph: DeEmphasis,
    /// LPC coefficients from the previous frame (for interpolation)
    prev_lpc: [f32; LPC_ORDER],
}

impl SynthesisFrameProcessor {
    pub fn new() -> Self {
        Self {
            filter: SynthesisFilter::new(),
            deemph: DeEmphasis::new(PRE_EMPHASIS_COEFF),
            prev_lpc: [0.0; LPC_ORDER],
        }
    }

    /// Synthesize one frame of audio.
    ///
    /// - `params`: decoded frame parameters (LSFs, gain, etc.)
    /// - `excitation`: mixed excitation signal for this frame (FRAME_SAMPLES long)
    /// - `output`: buffer to write synthesized audio (FRAME_SAMPLES long)
    pub fn process_frame(
        &mut self,
        params: &FrameParams,
        excitation: &[f32],
        output: &mut [f32],
    ) {
        let n = excitation.len().min(output.len()).min(FRAME_SAMPLES);

        // Convert LSFs to LPC coefficients
        let lpc = lsf_to_lpc(&params.lsf);

        // Apply gain to excitation
        let gain_linear = db_to_linear(params.gain);
        let mut scaled_excitation = [0.0f32; FRAME_SAMPLES];
        for i in 0..n {
            scaled_excitation[i] = excitation[i] * gain_linear;
        }

        // Synthesis filter with LPC interpolation from previous frame
        self.filter.process_buffer_interp(
            &scaled_excitation[..n],
            &self.prev_lpc,
            &lpc,
            &mut output[..n],
        );

        // De-emphasis
        self.deemph.process_inplace(&mut output[..n]);

        // Store current LPC for next frame's interpolation
        self.prev_lpc = lpc;
    }

    /// Process a frame without LPC interpolation (uses current frame's LPC only).
    /// Useful for the first frame or after a reset.
    pub fn process_frame_no_interp(
        &mut self,
        params: &FrameParams,
        excitation: &[f32],
        output: &mut [f32],
    ) {
        let n = excitation.len().min(output.len()).min(FRAME_SAMPLES);

        let lpc = lsf_to_lpc(&params.lsf);
        let gain_linear = db_to_linear(params.gain);

        let mut scaled_excitation = [0.0f32; FRAME_SAMPLES];
        for i in 0..n {
            scaled_excitation[i] = excitation[i] * gain_linear;
        }

        self.filter.process_buffer(
            &scaled_excitation[..n],
            &lpc,
            &mut output[..n],
        );

        self.deemph.process_inplace(&mut output[..n]);
        self.prev_lpc = lpc;
    }

    pub fn reset(&mut self) {
        self.filter.reset();
        self.deemph.reset();
        self.prev_lpc = [0.0; LPC_ORDER];
    }

    /// Access the previous frame's LPC coefficients (for testing/debug).
    pub fn prev_lpc(&self) -> &[f32; LPC_ORDER] {
        &self.prev_lpc
    }
}

/// Convert dB gain to linear amplitude.
#[inline]
pub fn db_to_linear(db: f32) -> f32 {
    // 10^(dB/20)
    powf(10.0, db / 20.0)
}

/// Convert linear amplitude to dB gain.
#[inline]
pub fn linear_to_db(linear: f32) -> f32 {
    20.0 * log10f(fmaxf(fabsf(linear), 1e-10))
}

#[cfg(test)]
mod tests {
    use super::*;

    // ── SynthesisFilter basics ──

    #[test]
    fn test_identity_filter() {
        // All-zero LPC coefficients → H(z) = 1, output = input
        let mut filt = SynthesisFilter::new();
        let a = [0.0f32; LPC_ORDER];

        let input = [1.0, 2.0, 3.0, -1.0, 0.5];
        for &x in &input {
            let y = filt.process_sample(x, &a);
            assert!(
                (y - x).abs() < 1e-6,
                "Identity filter: in={}, out={}",
                x, y
            );
        }
    }

    #[test]
    fn test_impulse_response() {
        // Single-tap filter: a[0] = 0.9, rest = 0
        // H(z) = 1 / (1 + 0.9*z^-1)
        // Impulse response: y[0] = 1.0, y[n] = -0.9 * y[n-1]
        let mut filt = SynthesisFilter::new();
        let mut a = [0.0f32; LPC_ORDER];
        a[0] = 0.9;

        // Feed impulse
        let y0 = filt.process_sample(1.0, &a);
        assert!((y0 - 1.0).abs() < 1e-6);

        let y1 = filt.process_sample(0.0, &a);
        assert!((y1 - (-0.9)).abs() < 1e-6, "y[1] should be -0.9, got {}", y1);

        let y2 = filt.process_sample(0.0, &a);
        assert!((y2 - 0.81).abs() < 1e-5, "y[2] should be 0.81, got {}", y2);

        let y3 = filt.process_sample(0.0, &a);
        assert!((y3 - (-0.729)).abs() < 1e-4, "y[3] should be -0.729, got {}", y3);
    }

    #[test]
    fn test_impulse_response_decays() {
        // Stable filter (all poles inside unit circle) → impulse response must decay
        let mut filt = SynthesisFilter::new();
        let mut a = [0.0f32; LPC_ORDER];
        a[0] = 0.5;
        a[1] = 0.2;

        let _y0 = filt.process_sample(1.0, &a);

        let mut decayed = false;
        for _ in 0..100 {
            let y = filt.process_sample(0.0, &a);
            if y.abs() < 0.001 {
                decayed = true;
                break;
            }
        }
        assert!(decayed, "Stable filter impulse response should decay");
    }

    #[test]
    fn test_process_buffer() {
        let mut filt1 = SynthesisFilter::new();
        let mut filt2 = SynthesisFilter::new();
        let a = [0.0f32; LPC_ORDER]; // identity

        let excitation = [1.0, 2.0, 3.0, 4.0, 5.0];
        let mut output = [0.0f32; 5];

        filt1.process_buffer(&excitation, &a, &mut output);

        // Compare with sample-by-sample
        for i in 0..5 {
            let y = filt2.process_sample(excitation[i], &a);
            assert!(
                (output[i] - y).abs() < 1e-6,
                "Buffer vs sample mismatch at [{}]",
                i
            );
        }
    }

    #[test]
    fn test_state_carries_across_calls() {
        let mut filt = SynthesisFilter::new();
        let mut a = [0.0f32; LPC_ORDER];
        a[0] = 0.5;

        // Process two samples
        let y0 = filt.process_sample(1.0, &a);
        let y1 = filt.process_sample(0.0, &a);

        // State[0] should be the most recent output
        assert!(
            (filt.state()[0] - y1).abs() < 1e-6,
            "State[0] should be last output"
        );

        // Reset and verify
        filt.reset();
        assert!(
            filt.state().iter().all(|&s| s == 0.0),
            "State should be zero after reset"
        );
    }

    #[test]
    fn test_dc_gain() {
        // DC gain of H(z) = 1/A(z) at z=1 is 1 / (1 + sum(a_i))
        // For a = [0.5, 0, ...]: DC gain = 1 / 1.5 = 0.6667
        let mut filt = SynthesisFilter::new();
        let mut a = [0.0f32; LPC_ORDER];
        a[0] = 0.5;

        // Feed constant input, let it settle
        let mut y = 0.0f32;
        for _ in 0..200 {
            y = filt.process_sample(1.0, &a);
        }

        let expected_dc = 1.0 / (1.0 + 0.5);
        assert!(
            (y - expected_dc).abs() < 0.01,
            "DC gain: expected {}, got {}",
            expected_dc, y
        );
    }

    // ── Interpolated processing ──

    #[test]
    fn test_interp_same_coefficients() {
        // When a_start == a_end, interpolated should match fixed
        let mut filt1 = SynthesisFilter::new();
        let mut filt2 = SynthesisFilter::new();
        let mut a = [0.0f32; LPC_ORDER];
        a[0] = 0.7;
        a[1] = -0.2;

        let excitation = [1.0, 0.5, -0.3, 0.8, 0.0, -0.5, 0.2, 0.1];
        let mut out1 = [0.0f32; 8];
        let mut out2 = [0.0f32; 8];

        filt1.process_buffer(&excitation, &a, &mut out1);
        filt2.process_buffer_interp(&excitation, &a, &a, &mut out2);

        for i in 0..8 {
            assert!(
                (out1[i] - out2[i]).abs() < 1e-5,
                "Interp with same coeffs should match fixed at [{}]: {} vs {}",
                i, out1[i], out2[i]
            );
        }
    }

    #[test]
    fn test_interp_produces_output() {
        let mut filt = SynthesisFilter::new();
        let mut a_start = [0.0f32; LPC_ORDER];
        let mut a_end = [0.0f32; LPC_ORDER];
        a_start[0] = 0.3;
        a_end[0] = 0.8;

        let excitation = [1.0; 32];
        let mut output = [0.0f32; 32];
        filt.process_buffer_interp(&excitation, &a_start, &a_end, &mut output);

        // Output should have non-trivial values
        let energy: f32 = output.iter().map(|x| x * x).sum::<f32>() / 32.0;
        assert!(energy > 0.1, "Interpolated output should have energy: {}", energy);
    }

    // ── DeEmphasis ──

    #[test]
    fn test_deemphasis_identity() {
        // coeff = 0 → pass-through
        let mut de = DeEmphasis::new(0.0);
        assert!((de.process_sample(1.0) - 1.0).abs() < 1e-6);
        assert!((de.process_sample(-0.5) - (-0.5)).abs() < 1e-6);
    }

    #[test]
    fn test_deemphasis_accumulates() {
        // y[n] = x[n] + 0.97 * y[n-1]
        let mut de = DeEmphasis::new(0.97);
        let y0 = de.process_sample(1.0);
        assert!((y0 - 1.0).abs() < 1e-6);

        let y1 = de.process_sample(0.0);
        assert!((y1 - 0.97).abs() < 1e-6, "y[1] should be 0.97, got {}", y1);

        let y2 = de.process_sample(0.0);
        assert!((y2 - 0.9409).abs() < 1e-4, "y[2] should be ~0.9409, got {}", y2);
    }

    #[test]
    fn test_pre_emphasis_deemphasis_inverse() {
        // Pre-emphasis then de-emphasis should roughly recover the original
        let original = [0.5, 0.8, -0.3, 0.1, 0.9, -0.7, 0.4, 0.2];
        let coeff = 0.97;

        // Pre-emphasis: y[n] = x[n] - coeff * x[n-1]
        let mut pre = [0.0f32; 8];
        pre[0] = original[0]; // first sample unchanged
        for i in 1..8 {
            pre[i] = original[i] - coeff * original[i - 1];
        }

        // De-emphasis
        let mut de = DeEmphasis::new(coeff);
        let mut recovered = [0.0f32; 8];
        for i in 0..8 {
            recovered[i] = de.process_sample(pre[i]);
        }

        // Should match original
        for i in 0..8 {
            assert!(
                (recovered[i] - original[i]).abs() < 1e-5,
                "Pre/de-emphasis inverse failed at [{}]: {} vs {}",
                i, original[i], recovered[i]
            );
        }
    }

    #[test]
    fn test_deemphasis_inplace() {
        let mut de1 = DeEmphasis::new(0.97);
        let mut de2 = DeEmphasis::new(0.97);

        let input = [1.0, 0.5, -0.3, 0.8];

        // Sample-by-sample
        let mut expected = [0.0f32; 4];
        for i in 0..4 {
            expected[i] = de1.process_sample(input[i]);
        }

        // In-place
        let mut buf = input;
        de2.process_inplace(&mut buf);

        for i in 0..4 {
            assert!(
                (buf[i] - expected[i]).abs() < 1e-6,
                "Inplace mismatch at [{}]",
                i
            );
        }
    }

    // ── SynthesisFrameProcessor tests ──

    #[test]
    fn test_db_to_linear() {
        assert!((db_to_linear(0.0) - 1.0).abs() < 1e-6);
        assert!((db_to_linear(-20.0) - 0.1).abs() < 1e-4);
        assert!((db_to_linear(-6.0) - 0.5012).abs() < 0.01);
        assert!((db_to_linear(-60.0) - 0.001).abs() < 1e-4);
    }

    #[test]
    fn test_linear_to_db_roundtrip() {
        let values = [-3.0, -10.0, -20.0, -40.0, 0.0];
        for &db in &values {
            let lin = db_to_linear(db);
            let recovered = linear_to_db(lin);
            assert!(
                (recovered - db).abs() < 0.01,
                "dB roundtrip failed: {} → {} → {}",
                db, lin, recovered
            );
        }
    }

    #[test]
    fn test_frame_processor_silence() {
        // Zero excitation → output should be near-silent
        let mut proc = SynthesisFrameProcessor::new();
        let params = FrameParams::default(); // gain = -60 dB
        let excitation = [0.0f32; FRAME_SAMPLES];
        let mut output = [0.0f32; FRAME_SAMPLES];

        proc.process_frame(&params, &excitation, &mut output);

        let peak = output.iter().fold(0.0f32, |m, &s| m.max(s.abs()));
        assert!(peak < 0.001, "Silence in should give silence out, peak={}", peak);
    }

    #[test]
    fn test_frame_processor_produces_output() {
        // Non-zero excitation + reasonable gain → non-trivial output
        let mut proc = SynthesisFrameProcessor::new();
        let mut params = FrameParams::default();
        params.gain = -10.0; // modest gain

        // Pulse-like excitation
        let mut excitation = [0.0f32; FRAME_SAMPLES];
        excitation[0] = 1.0;
        excitation[80] = 1.0;

        let mut output = [0.0f32; FRAME_SAMPLES];
        proc.process_frame(&params, &excitation, &mut output);

        let energy: f32 = output.iter().map(|x| x * x).sum::<f32>() / FRAME_SAMPLES as f32;
        assert!(energy > 1e-6, "Should produce non-trivial output, energy={}", energy);
    }

    #[test]
    fn test_frame_processor_gain_scaling() {
        // Higher gain should produce louder output
        let excitation = [0.1f32; FRAME_SAMPLES];

        let mut proc_quiet = SynthesisFrameProcessor::new();
        let mut params_quiet = FrameParams::default();
        params_quiet.gain = -40.0;
        let mut out_quiet = [0.0f32; FRAME_SAMPLES];
        proc_quiet.process_frame_no_interp(&params_quiet, &excitation, &mut out_quiet);

        let mut proc_loud = SynthesisFrameProcessor::new();
        let mut params_loud = FrameParams::default();
        params_loud.gain = -10.0;
        let mut out_loud = [0.0f32; FRAME_SAMPLES];
        proc_loud.process_frame_no_interp(&params_loud, &excitation, &mut out_loud);

        let energy_quiet: f32 = out_quiet.iter().map(|x| x * x).sum();
        let energy_loud: f32 = out_loud.iter().map(|x| x * x).sum();
        assert!(
            energy_loud > energy_quiet,
            "Louder gain should produce more energy: loud={}, quiet={}",
            energy_loud, energy_quiet
        );
    }

    #[test]
    fn test_frame_processor_stores_prev_lpc() {
        let mut proc = SynthesisFrameProcessor::new();

        // Initially zero
        assert!(proc.prev_lpc().iter().all(|&x| x == 0.0));

        let params = FrameParams::default();
        let excitation = [0.0f32; FRAME_SAMPLES];
        let mut output = [0.0f32; FRAME_SAMPLES];

        proc.process_frame(&params, &excitation, &mut output);

        // After one frame, prev_lpc should be set from default LSFs
        let lpc = lsf_to_lpc(&params.lsf);
        for i in 0..LPC_ORDER {
            assert!(
                (proc.prev_lpc()[i] - lpc[i]).abs() < 1e-5,
                "prev_lpc[{}] mismatch: {} vs {}",
                i, proc.prev_lpc()[i], lpc[i]
            );
        }
    }

    #[test]
    fn test_frame_processor_reset() {
        let mut proc = SynthesisFrameProcessor::new();

        // Run a frame to populate state
        let mut params = FrameParams::default();
        params.gain = -10.0;
        let excitation = [0.5f32; FRAME_SAMPLES];
        let mut output = [0.0f32; FRAME_SAMPLES];
        proc.process_frame(&params, &excitation, &mut output);

        // Reset
        proc.reset();
        assert!(proc.prev_lpc().iter().all(|&x| x == 0.0));
    }

    #[test]
    fn test_frame_processor_output_finite() {
        // Process multiple frames, verify output stays finite
        let mut proc = SynthesisFrameProcessor::new();
        let mut params = FrameParams::default();
        params.gain = -6.0;

        let excitation = [0.2f32; FRAME_SAMPLES];
        let mut output = [0.0f32; FRAME_SAMPLES];

        for _ in 0..10 {
            proc.process_frame(&params, &excitation, &mut output);
            assert!(
                output.iter().all(|x| x.is_finite()),
                "Output must be finite after multiple frames"
            );
        }
    }

    #[test]
    fn test_no_interp_vs_interp_first_frame() {
        // On a fresh processor (prev_lpc = zeros), no_interp and interp
        // should differ since interp blends from zeros
        let mut proc1 = SynthesisFrameProcessor::new();
        let mut proc2 = SynthesisFrameProcessor::new();
        let mut params = FrameParams::default();
        params.gain = -10.0;

        let excitation = [0.3f32; FRAME_SAMPLES];
        let mut out1 = [0.0f32; FRAME_SAMPLES];
        let mut out2 = [0.0f32; FRAME_SAMPLES];

        proc1.process_frame_no_interp(&params, &excitation, &mut out1);
        proc2.process_frame(&params, &excitation, &mut out2);

        // Both should produce output
        let e1: f32 = out1.iter().map(|x| x * x).sum();
        let e2: f32 = out2.iter().map(|x| x * x).sum();
        assert!(e1 > 1e-6, "no_interp should produce output");
        assert!(e2 > 1e-6, "interp should produce output");

        // After the first frame, prev_lpc should be the same for both
        for i in 0..LPC_ORDER {
            assert!(
                (proc1.prev_lpc()[i] - proc2.prev_lpc()[i]).abs() < 1e-5,
                "prev_lpc should match after first frame"
            );
        }
    }
}