melpe-rs 0.1.2

/// MELPe 600 bps Decoder
///
/// Accepts 6 bytes of bitstream (1 superframe),
/// dequantizes, interpolates parameters across 3 frames,
/// generates mixed excitation, and synthesizes 540 samples.

use crate::core_types::{
    FrameParams, FRAME_SAMPLES, SUPERFRAME_FRAMES,
    SUPERFRAME_SAMPLES, SUPERFRAME_BYTES_600, NUM_LSF, NUM_BANDS,
};
use crate::bitstream::unpack_superframe;
use crate::quantize::dequantize_superframe;
use crate::voicing::{NoiseGen, mixed_excitation};
use crate::synthesis::SynthesisFrameProcessor;

/// MELPe 600 bps decoder.
pub struct Decoder {
    /// Frame-level synthesis processor (filter + de-emphasis + LPC interp)
    synth: SynthesisFrameProcessor,
    /// Previous superframe's anchor LSFs (needed for LSF interpolation)
    prev_anchor_lsf: [f32; NUM_LSF],
    /// Noise generator for mixed excitation
    noise: NoiseGen,
    /// Pulse train phase accumulator (carries across frames)
    pulse_phase: f32,
    /// Whether we've decoded at least one superframe
    initialized: bool,
}

impl Decoder {
    pub fn new() -> Self {
        Self {
            synth: SynthesisFrameProcessor::new(),
            prev_anchor_lsf: crate::core_types::default_lsf(),
            noise: NoiseGen::new(12345),
            pulse_phase: 0.0,
            initialized: false,
        }
    }

    /// Decode one superframe: 6 bytes → 540 samples.
    ///
    /// `bitstream` must be exactly SUPERFRAME_BYTES_600 (6) bytes.
    /// Writes SUPERFRAME_SAMPLES (540) samples to `output`.
    pub fn decode(
        &mut self,
        bitstream: &[u8; SUPERFRAME_BYTES_600],
        output: &mut [f32; SUPERFRAME_SAMPLES],
    ) {
        // Unpack bitstream → quantized indices
        let (quantized, _bits_read) = unpack_superframe(bitstream);

        // Dequantize → 3 FrameParams (with LSF interpolation)
        let frames = dequantize_superframe(&quantized, &self.prev_anchor_lsf);

        // Update previous anchor LSFs (frame 2 is the anchor)
        self.prev_anchor_lsf = frames[SUPERFRAME_FRAMES - 1].lsf;

        // Synthesize each of the 3 frames
        for f in 0..SUPERFRAME_FRAMES {
            let start = f * FRAME_SAMPLES;
            let end = start + FRAME_SAMPLES;

            self.synthesize_frame(
                &frames[f],
                &mut output[start..end],
                f == 0 && !self.initialized,
            );
        }

        self.initialized = true;
    }

    /// Synthesize a single frame: generate excitation, run through synthesis filter.
    fn synthesize_frame(
        &mut self,
        params: &FrameParams,
        output: &mut [f32],
        first_frame: bool,
    ) {
        let n = output.len().min(FRAME_SAMPLES);

        // Generate mixed excitation
        let mut excitation = [0.0f32; FRAME_SAMPLES];

        // Collapse bandpass voicing to the format mixed_excitation expects
        let mut voicing = [0.0f32; NUM_BANDS];
        voicing.copy_from_slice(&params.bandpass_voicing);

        mixed_excitation(
            params.pitch,
            &voicing,
            params.jitter,
            &mut excitation[..n],
            &mut self.pulse_phase,
            &mut self.noise,
        );

        // Run synthesis (LSF→LPC + gain + filter + de-emphasis)
        if first_frame {
            self.synth.process_frame_no_interp(params, &excitation[..n], &mut output[..n]);
        } else {
            self.synth.process_frame(params, &excitation[..n], &mut output[..n]);
        }
    }

    /// Reset decoder state.
    pub fn reset(&mut self) {
        self.synth.reset();
        self.prev_anchor_lsf = crate::core_types::default_lsf();
        self.noise = NoiseGen::new(12345);
        self.pulse_phase = 0.0;
        self.initialized = false;
    }

    /// Whether at least one superframe has been decoded.
    pub fn is_initialized(&self) -> bool {
        self.initialized
    }

    /// Previous anchor LSFs (for debugging/testing).
    pub fn prev_anchor_lsf(&self) -> &[f32; NUM_LSF] {
        &self.prev_anchor_lsf
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core_types::{SAMPLE_RATE, SuperFrameQuantized};
    use crate::bitstream::pack_superframe;
    use crate::quantize;

    /// Pack a known SuperFrameQuantized into bytes for decoder testing.
    fn pack_known(sq: &SuperFrameQuantized) -> [u8; SUPERFRAME_BYTES_600] {
        let mut buf = [0u8; SUPERFRAME_BYTES_600];
        pack_superframe(sq, &mut buf);
        buf
    }

    #[test]
    fn test_decode_output_length() {
        let mut dec = Decoder::new();
        let bitstream = [0u8; SUPERFRAME_BYTES_600];
        let mut output = [0.0f32; SUPERFRAME_SAMPLES];

        dec.decode(&bitstream, &mut output);

        // All 540 samples should be written (finite values)
        assert!(
            output.iter().all(|x| x.is_finite()),
            "All output samples must be finite"
        );
    }

    #[test]
    fn test_decode_zeros_quiet() {
        // All-zero bitstream → should produce relatively quiet output
        let mut dec = Decoder::new();
        let bitstream = [0u8; SUPERFRAME_BYTES_600];
        let mut output = [0.0f32; SUPERFRAME_SAMPLES];

        dec.decode(&bitstream, &mut output);

        let rms = (output.iter().map(|x| x * x).sum::<f32>() / SUPERFRAME_SAMPLES as f32).sqrt();
        // The zero bitstream dequantizes to low gain, so output should be quiet
        // (not necessarily silent due to noise excitation and filter transients)
        assert!(
            rms < 1.0,
            "Zero bitstream should produce quiet output, rms={}",
            rms
        );
    }

    #[test]
    fn test_decode_nonzero_bitstream() {
        // A bitstream with voiced excitation and moderate gain
        let sq = SuperFrameQuantized {
            lsf_indices: [4, 4, 4, 4, 2, 2, 2, 2, 2, 2],
            pitch_index: 30,        // mid-range pitch
            gain_index: 20,         // moderate gain
            voicing_bits: 0b0111,   // all 3 frames voiced
            jitter_bit: 0,
        };
        let bitstream = pack_known(&sq);

        let mut dec = Decoder::new();
        let mut output = [0.0f32; SUPERFRAME_SAMPLES];
        dec.decode(&bitstream, &mut output);

        // Voiced with moderate gain → should have real energy
        let energy: f32 = output.iter().map(|x| x * x).sum::<f32>() / SUPERFRAME_SAMPLES as f32;
        assert!(
            energy > 1e-8,
            "Voiced frame with gain should produce energy, got {}",
            energy
        );
        assert!(
            output.iter().all(|x| x.is_finite()),
            "Output must be finite"
        );
    }

    #[test]
    fn test_decode_state_updates() {
        let mut dec = Decoder::new();
        assert!(!dec.is_initialized());

        let bitstream = [0u8; SUPERFRAME_BYTES_600];
        let mut output = [0.0f32; SUPERFRAME_SAMPLES];

        dec.decode(&bitstream, &mut output);
        assert!(dec.is_initialized());

        // prev_anchor_lsf should be ordered
        let lsf = dec.prev_anchor_lsf();
        for i in 1..NUM_LSF {
            assert!(
                lsf[i] > lsf[i - 1],
                "Anchor LSFs should be ordered after decode"
            );
        }
    }

    #[test]
    fn test_decode_reset() {
        let mut dec = Decoder::new();
        let bitstream = [0u8; SUPERFRAME_BYTES_600];
        let mut output = [0.0f32; SUPERFRAME_SAMPLES];

        dec.decode(&bitstream, &mut output);
        assert!(dec.is_initialized());

        dec.reset();
        assert!(!dec.is_initialized());
    }

    #[test]
    fn test_decode_multi_superframe() {
        // Decode 5 consecutive superframes, verify stability
        let mut dec = Decoder::new();
        let sq = SuperFrameQuantized {
            lsf_indices: [3, 3, 3, 3, 1, 1, 1, 1, 1, 1],
            pitch_index: 25,
            gain_index: 15,
            voicing_bits: 0b0101, // frames 0 and 2 voiced
            jitter_bit: 0,
        };
        let bitstream = pack_known(&sq);

        let mut output = [0.0f32; SUPERFRAME_SAMPLES];
        for i in 0..5 {
            dec.decode(&bitstream, &mut output);
            assert!(
                output.iter().all(|x| x.is_finite()),
                "Output not finite at superframe {}",
                i
            );
            let peak = output.iter().fold(0.0f32, |m, &s| m.max(s.abs()));
            assert!(
                peak < 100.0,
                "Output blowing up at superframe {}: peak={}",
                i, peak
            );
        }
    }

    #[test]
    fn test_decode_deterministic() {
        let sq = SuperFrameQuantized {
            lsf_indices: [5, 3, 6, 2, 3, 1, 2, 3, 1, 2],
            pitch_index: 35,
            gain_index: 18,
            voicing_bits: 0b0011,
            jitter_bit: 1,
        };
        let bitstream = pack_known(&sq);

        let mut dec1 = Decoder::new();
        let mut dec2 = Decoder::new();
        let mut out1 = [0.0f32; SUPERFRAME_SAMPLES];
        let mut out2 = [0.0f32; SUPERFRAME_SAMPLES];

        dec1.decode(&bitstream, &mut out1);
        dec2.decode(&bitstream, &mut out2);

        for i in 0..SUPERFRAME_SAMPLES {
            assert!(
                (out1[i] - out2[i]).abs() < 1e-6,
                "Determinism failed at sample [{}]: {} vs {}",
                i, out1[i], out2[i]
            );
        }
    }

    #[test]
    fn test_voiced_vs_unvoiced_differ() {
        // Same parameters except voicing
        let sq_voiced = SuperFrameQuantized {
            lsf_indices: [4, 4, 4, 4, 2, 2, 2, 2, 2, 2],
            pitch_index: 30,
            gain_index: 20,
            voicing_bits: 0b0111, // all voiced
            jitter_bit: 0,
        };
        let sq_unvoiced = SuperFrameQuantized {
            lsf_indices: [4, 4, 4, 4, 2, 2, 2, 2, 2, 2],
            pitch_index: 30,
            gain_index: 20,
            voicing_bits: 0b0000, // all unvoiced
            jitter_bit: 0,
        };

        let mut dec = Decoder::new();
        let mut out_v = [0.0f32; SUPERFRAME_SAMPLES];
        dec.decode(&pack_known(&sq_voiced), &mut out_v);

        dec.reset();
        let mut out_uv = [0.0f32; SUPERFRAME_SAMPLES];
        dec.decode(&pack_known(&sq_unvoiced), &mut out_uv);

        // Outputs should differ
        let diff: f32 = out_v.iter().zip(out_uv.iter())
            .map(|(a, b)| (a - b).abs())
            .sum::<f32>() / SUPERFRAME_SAMPLES as f32;
        assert!(
            diff > 1e-6,
            "Voiced and unvoiced should produce different output, avg_diff={}",
            diff
        );
    }

    #[test]
    fn test_different_bitstreams_differ() {
        let sq1 = SuperFrameQuantized {
            lsf_indices: [1, 1, 1, 1, 1, 1, 1, 1, 1, 1],
            pitch_index: 10,
            gain_index: 10,
            voicing_bits: 0b0111,
            jitter_bit: 0,
        };
        let sq2 = SuperFrameQuantized {
            lsf_indices: [6, 6, 6, 6, 3, 3, 3, 3, 3, 3],
            pitch_index: 50,
            gain_index: 25,
            voicing_bits: 0b0000,
            jitter_bit: 1,
        };

        let mut dec = Decoder::new();
        let mut out1 = [0.0f32; SUPERFRAME_SAMPLES];
        dec.decode(&pack_known(&sq1), &mut out1);

        dec.reset();
        let mut out2 = [0.0f32; SUPERFRAME_SAMPLES];
        dec.decode(&pack_known(&sq2), &mut out2);

        assert_ne!(
            out1.as_slice(), out2.as_slice(),
            "Different bitstreams should produce different output"
        );
    }
}