melpe-rs 0.1.2

/// MELPe 600 bps Encoder
///
/// Accepts 540 samples (1 superframe = 3 × 22.5ms frames),
/// runs the full analysis chain, and produces 6 bytes of bitstream.

use crate::core_types::{
    FrameParams, SuperFrame, FRAME_SAMPLES, SUPERFRAME_FRAMES,
    SUPERFRAME_SAMPLES, SUPERFRAME_BYTES_600, NUM_LSF, NUM_BANDS,
};
use crate::lpc::{autocorrelation, levinson_durbin, lpc_to_lsf, hamming_window, pre_emphasis};
use crate::pitch::{estimate_pitch, adaptive_clip_threshold};
use crate::voicing::{FilterBank, bandpass_voicing_strength};
use crate::quantize::quantize_superframe;
use crate::bitstream::pack_superframe;
use crate::core_types::LPC_ORDER;
use crate::math::{log10f, sqrtf, clampf};

/// Pre-emphasis coefficient (matches synthesis side)
const PRE_EMPHASIS_COEFF: f32 = 0.97;

/// Center-clipping fraction for pitch detection
const CLIP_FRACTION: f32 = 0.3;

/// MELPe 600 bps encoder.
pub struct Encoder {
    /// Bandpass filter bank for voicing analysis
    filter_bank: FilterBank,
    /// Previous superframe's anchor LSFs (for decoder-side interpolation context)
    prev_anchor_lsf: [f32; NUM_LSF],
    /// Whether we've processed at least one superframe
    initialized: bool,
}

impl Encoder {
    pub fn new() -> Self {
        Self {
            filter_bank: FilterBank::new(),
            prev_anchor_lsf: crate::core_types::default_lsf(),
            initialized: false,
        }
    }

    /// Encode one superframe: 540 samples → 6 bytes.
    ///
    /// `samples` must be exactly SUPERFRAME_SAMPLES (540) long.
    /// Returns the packed bitstream in `output` (6 bytes).
    pub fn encode(
        &mut self,
        samples: &[f32; SUPERFRAME_SAMPLES],
        output: &mut [u8; SUPERFRAME_BYTES_600],
    ) {
        // Analyze each of the 3 constituent frames
        let mut superframe = SuperFrame::default();

        for f in 0..SUPERFRAME_FRAMES {
            let start = f * FRAME_SAMPLES;
            let end = start + FRAME_SAMPLES;
            superframe.frames[f] = self.analyze_frame(&samples[start..end]);
        }

        // Quantize the superframe
        let quantized = quantize_superframe(&superframe);

        // Pack into bitstream
        pack_superframe(&quantized, output);

        // Update state
        self.prev_anchor_lsf = superframe.anchor().lsf;
        self.initialized = true;
    }

    /// Analyze a single frame of samples into FrameParams.
    fn analyze_frame(&mut self, samples: &[f32]) -> FrameParams {
        let n = samples.len().min(FRAME_SAMPLES);

        // Copy and apply pre-emphasis
        let mut buf = [0.0f32; FRAME_SAMPLES];
        buf[..n].copy_from_slice(&samples[..n]);
        pre_emphasis(&mut buf[..n], PRE_EMPHASIS_COEFF);

        // Window for LPC analysis
        let mut windowed = [0.0f32; FRAME_SAMPLES];
        windowed[..n].copy_from_slice(&buf[..n]);
        hamming_window(&mut windowed[..n]);

        // LPC analysis
        let r = autocorrelation(&windowed[..n], LPC_ORDER);
        let (lpc_coeffs, _pred_error) = levinson_durbin(&r, LPC_ORDER);

        // LPC → LSF (fall back to default if unstable)
        let lsf = lpc_to_lsf(&lpc_coeffs)
            .unwrap_or_else(crate::core_types::default_lsf);

        // Pitch detection (on pre-emphasized signal, not windowed)
        let clip_thresh = adaptive_clip_threshold(&buf[..n], CLIP_FRACTION);
        let pitch_result = estimate_pitch(&buf[..n], clip_thresh);

        // Bandpass voicing analysis (on original samples, not pre-emphasized)
        let bands = self.filter_bank.analyze(&samples[..n]);
        let voicing = bandpass_voicing_strength(
            &bands,
            pitch_result.period,
            n,
        );

        // Compute frame gain (RMS in dB)
        let rms = compute_rms(samples);
        let gain_db = if rms > 1e-10 {
            20.0 * log10f(rms)
        } else {
            -60.0
        };

        // Jitter: higher for unvoiced/noisy frames
        let jitter = if pitch_result.voiced {
            clampf(1.0 - pitch_result.correlation, 0.0, 1.0)
        } else {
            0.75
        };

        FrameParams {
            lsf,
            pitch: pitch_result.period,
            bandpass_voicing: voicing,
            gain: gain_db,
            jitter,
        }
    }

    /// Reset encoder state.
    pub fn reset(&mut self) {
        self.filter_bank.reset();
        self.prev_anchor_lsf = crate::core_types::default_lsf();
        self.initialized = false;
    }

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

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

/// Compute RMS energy of a sample buffer.
fn compute_rms(samples: &[f32]) -> f32 {
    if samples.is_empty() {
        return 0.0;
    }
    let sum_sq: f32 = samples.iter().map(|&s| s * s).sum();
    sqrtf(sum_sq / samples.len() as f32)
}

#[cfg(all(test, feature = "std"))]
mod tests {
    use super::*;
    use crate::core_types::SAMPLE_RATE;

    fn sine_wave(freq: f32, n: usize) -> Vec<f32> {
        (0..n)
            .map(|i| {
                0.5 * (2.0 * core::f32::consts::PI * freq * i as f32 / SAMPLE_RATE as f32).sin()
            })
            .collect()
    }

    fn to_superframe_buf(samples: &[f32]) -> [f32; SUPERFRAME_SAMPLES] {
        let mut buf = [0.0f32; SUPERFRAME_SAMPLES];
        let n = samples.len().min(SUPERFRAME_SAMPLES);
        buf[..n].copy_from_slice(&samples[..n]);
        buf
    }

    #[test]
    fn test_encode_produces_6_bytes() {
        let mut enc = Encoder::new();
        let input = to_superframe_buf(&sine_wave(200.0, SUPERFRAME_SAMPLES));
        let mut output = [0u8; SUPERFRAME_BYTES_600];

        enc.encode(&input, &mut output);

        // Output should be exactly 6 bytes and not all zeros
        // (a sine wave should produce non-trivial parameters)
        let nonzero = output.iter().any(|&b| b != 0);
        assert!(nonzero, "Encoded sine wave should produce non-zero bytes");
    }

    #[test]
    fn test_encode_silence() {
        let mut enc = Encoder::new();
        let input = [0.0f32; SUPERFRAME_SAMPLES];
        let mut output = [0u8; SUPERFRAME_BYTES_600];

        enc.encode(&input, &mut output);
        // Should not panic, output is valid (may or may not be all zeros)
    }

    #[test]
    fn test_encode_deterministic() {
        let samples = sine_wave(150.0, SUPERFRAME_SAMPLES);

        let mut enc1 = Encoder::new();
        let mut enc2 = Encoder::new();
        let input = to_superframe_buf(&samples);
        let mut out1 = [0u8; SUPERFRAME_BYTES_600];
        let mut out2 = [0u8; SUPERFRAME_BYTES_600];

        enc1.encode(&input, &mut out1);
        enc2.encode(&input, &mut out2);

        assert_eq!(out1, out2, "Same input must produce same output");
    }

    #[test]
    fn test_encode_different_inputs_differ() {
        let sine = to_superframe_buf(&sine_wave(200.0, SUPERFRAME_SAMPLES));

        // Noise-like signal
        let mut noise = [0.0f32; SUPERFRAME_SAMPLES];
        let mut seed: u32 = 42;
        for s in noise.iter_mut() {
            seed = seed.wrapping_mul(1103515245).wrapping_add(12345);
            *s = 0.5 * (seed as i32 as f32) / (i32::MAX as f32);
        }

        let mut enc = Encoder::new();
        let mut out_sine = [0u8; SUPERFRAME_BYTES_600];
        let mut out_noise = [0u8; SUPERFRAME_BYTES_600];

        enc.encode(&sine, &mut out_sine);
        enc.reset();
        enc.encode(&noise, &mut out_noise);

        assert_ne!(out_sine, out_noise, "Different signals should encode differently");
    }

    #[test]
    fn test_encoder_state_updates() {
        let mut enc = Encoder::new();
        assert!(!enc.is_initialized());

        let input = to_superframe_buf(&sine_wave(100.0, SUPERFRAME_SAMPLES));
        let mut output = [0u8; SUPERFRAME_BYTES_600];

        enc.encode(&input, &mut output);
        assert!(enc.is_initialized());

        // prev_anchor_lsf should have been updated from default
        // (may still be close to default for a sine, but should be set)
        let lsf = enc.prev_anchor_lsf();
        for i in 1..NUM_LSF {
            assert!(
                lsf[i] > lsf[i - 1],
                "Anchor LSFs should be ordered: [{}]={} <= [{}]={}",
                i - 1, lsf[i - 1], i, lsf[i]
            );
        }
    }

    #[test]
    fn test_encoder_reset() {
        let mut enc = Encoder::new();
        let input = to_superframe_buf(&sine_wave(200.0, SUPERFRAME_SAMPLES));
        let mut output = [0u8; SUPERFRAME_BYTES_600];

        enc.encode(&input, &mut output);
        assert!(enc.is_initialized());

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

    #[test]
    fn test_encode_multi_superframe() {
        // Encode 3 consecutive superframes, verify no panics
        let mut enc = Encoder::new();
        let samples = sine_wave(150.0, SUPERFRAME_SAMPLES * 3);
        let mut output = [0u8; SUPERFRAME_BYTES_600];

        for i in 0..3 {
            let start = i * SUPERFRAME_SAMPLES;
            let input = to_superframe_buf(&samples[start..start + SUPERFRAME_SAMPLES]);
            enc.encode(&input, &mut output);
        }
        assert!(enc.is_initialized());
    }

    #[test]
    fn test_compute_rms() {
        // DC signal of 0.5 → RMS = 0.5
        let dc = [0.5f32; 100];
        assert!((compute_rms(&dc) - 0.5).abs() < 1e-4);

        // Silence → RMS = 0
        let silence = [0.0f32; 100];
        assert!(compute_rms(&silence) < 1e-10);

        // Sine wave amplitude 1.0 → RMS ≈ 0.707
        let sine: Vec<f32> = (0..1000)
            .map(|i| (2.0 * core::f32::consts::PI * i as f32 / 100.0).sin())
            .collect();
        let rms = compute_rms(&sine);
        assert!(
            (rms - 0.707).abs() < 0.02,
            "Sine RMS should be ~0.707, got {}",
            rms
        );
    }

    #[test]
    fn test_loud_vs_quiet_gain() {
        // Louder input should produce a different encoding than quiet input
        let loud: Vec<f32> = sine_wave(200.0, SUPERFRAME_SAMPLES)
            .iter()
            .map(|&s| s * 2.0)
            .collect();
        let quiet: Vec<f32> = sine_wave(200.0, SUPERFRAME_SAMPLES)
            .iter()
            .map(|&s| s * 0.01)
            .collect();

        let mut enc = Encoder::new();
        let mut out_loud = [0u8; SUPERFRAME_BYTES_600];
        let mut out_quiet = [0u8; SUPERFRAME_BYTES_600];

        enc.encode(&to_superframe_buf(&loud), &mut out_loud);
        enc.reset();
        enc.encode(&to_superframe_buf(&quiet), &mut out_quiet);

        assert_ne!(out_loud, out_quiet, "Loud and quiet should encode differently");
    }
}