melpe-rs 0.1.2

/// Bandpass voicing and mixed excitation for MELPe
///
/// MELPe splits 0-4kHz into 5 bands and determines voicing per band.
/// Mixed excitation blends periodic (pulse) and aperiodic (noise) sources
/// per band, then sums them.

use crate::core_types::{NUM_BANDS, SAMPLE_RATE, FRAME_SAMPLES};
use crate::math::{sinf, cosf, sinhf, sqrtf, roundf, clampf};

/// Band edge frequencies in Hz for the 5 MELPe bands
/// Band 0: 0-500, Band 1: 500-1000, Band 2: 1000-2000, Band 3: 2000-3000, Band 4: 3000-4000
pub const BAND_EDGES: [(f32, f32); NUM_BANDS] = [
    (0.0, 500.0),
    (500.0, 1000.0),
    (1000.0, 2000.0),
    (2000.0, 3000.0),
    (3000.0, 4000.0),
];

// ── Second-order bandpass filter (biquad) ──

#[derive(Debug, Clone, Copy)]
pub struct Biquad {
    b0: f32, b1: f32, b2: f32,
    a1: f32, a2: f32,
    // State
    x1: f32, x2: f32,
    y1: f32, y2: f32,
}

impl Biquad {
    /// Design a bandpass biquad for the given center frequency and bandwidth.
    pub fn bandpass(center_hz: f32, bandwidth_hz: f32, sample_rate: f32) -> Self {
        let w0 = 2.0 * core::f32::consts::PI * center_hz / sample_rate;
        let bw = 2.0 * core::f32::consts::PI * bandwidth_hz / sample_rate;
        let alpha = sinf(w0) * sinhf(bw / 2.0);

        let b0 = alpha;
        let b1 = 0.0;
        let b2 = -alpha;
        let a0 = 1.0 + alpha;
        let a1 = -2.0 * cosf(w0);
        let a2 = 1.0 - alpha;

        Self {
            b0: b0 / a0, b1: b1 / a0, b2: b2 / a0,
            a1: a1 / a0, a2: a2 / a0,
            x1: 0.0, x2: 0.0, y1: 0.0, y2: 0.0,
        }
    }

    /// Process one sample through the filter.
    pub fn process(&mut self, x: f32) -> f32 {
        let y = self.b0 * x + self.b1 * self.x1 + self.b2 * self.x2
              - self.a1 * self.y1 - self.a2 * self.y2;
        self.x2 = self.x1;
        self.x1 = x;
        self.y2 = self.y1;
        self.y1 = y;
        y
    }

    /// Filter an entire buffer, writing output to `out`.
    pub fn filter_buf(&mut self, input: &[f32], out: &mut [f32]) {
        let n = input.len().min(out.len());
        for i in 0..n {
            out[i] = self.process(input[i]);
        }
    }

    pub fn reset(&mut self) {
        self.x1 = 0.0; self.x2 = 0.0;
        self.y1 = 0.0; self.y2 = 0.0;
    }
}

// ── Bandpass voicing analysis ──

/// Filter bank: 5 bandpass filters for MELPe bands
pub struct FilterBank {
    filters: [Biquad; NUM_BANDS],
}

impl FilterBank {
    pub fn new() -> Self {
        let sr = SAMPLE_RATE as f32;
        let mut filters = [Biquad::bandpass(1000.0, 500.0, sr); NUM_BANDS]; // placeholder

        for i in 0..NUM_BANDS {
            let (lo, hi) = BAND_EDGES[i];
            let center = (lo + hi) / 2.0;
            let bw = hi - lo;
            // Avoid 0 Hz center for band 0
            let center = if center < 1.0 { bw / 2.0 } else { center };
            filters[i] = Biquad::bandpass(center, bw, sr);
        }

        Self { filters }
    }

    /// Split signal into 5 bands. Returns per-band output buffers.
    pub fn analyze(&mut self, signal: &[f32]) -> [[f32; FRAME_SAMPLES]; NUM_BANDS] {
        let mut bands = [[0.0f32; FRAME_SAMPLES]; NUM_BANDS];
        let n = signal.len().min(FRAME_SAMPLES);

        for b in 0..NUM_BANDS {
            self.filters[b].filter_buf(&signal[..n], &mut bands[b][..n]);
        }

        bands
    }

    pub fn reset(&mut self) {
        for f in &mut self.filters {
            f.reset();
        }
    }
}

/// Compute band energies from per-band signals.
pub fn band_energies(bands: &[[f32; FRAME_SAMPLES]; NUM_BANDS], len: usize) -> [f32; NUM_BANDS] {
    let mut energies = [0.0f32; NUM_BANDS];
    let n = len.min(FRAME_SAMPLES);
    for b in 0..NUM_BANDS {
        let mut sum = 0.0f32;
        for i in 0..n {
            sum += bands[b][i] * bands[b][i];
        }
        energies[b] = sum / n as f32;
    }
    energies
}

/// Determine per-band voicing strengths.
///
/// Uses the correlation between the band signal and the periodic
/// component at the given pitch. Returns voicing strength 0.0..1.0 per band.
pub fn bandpass_voicing_strength(
    bands: &[[f32; FRAME_SAMPLES]; NUM_BANDS],
    pitch_period: f32,
    len: usize,
) -> [f32; NUM_BANDS] {
    let mut voicing = [0.0f32; NUM_BANDS];
    let n = len.min(FRAME_SAMPLES);
    let period = roundf(pitch_period) as usize;

    if period < 2 || period >= n {
        return voicing; // all unvoiced
    }

    for b in 0..NUM_BANDS {
        // Normalized correlation at the pitch lag within this band
        let mut num = 0.0f32;
        let mut e0 = 0.0f32;
        let mut e1 = 0.0f32;

        for i in period..n {
            num += bands[b][i] * bands[b][i - period];
            e0 += bands[b][i] * bands[b][i];
            e1 += bands[b][i - period] * bands[b][i - period];
        }

        let denom = sqrtf(e0 * e1);
        if denom > 1e-10 {
            voicing[b] = clampf(num / denom, 0.0, 1.0);
        }
    }

    voicing
}

// ── Mixed excitation generation ──

/// Simple LCG pseudo-random noise generator (no_std friendly)
pub struct NoiseGen {
    state: u32,
}

impl NoiseGen {
    pub const fn new(seed: u32) -> Self {
        Self { state: seed }
    }

    /// Generate one noise sample in [-1.0, 1.0]
    pub fn next(&mut self) -> f32 {
        self.state = self.state.wrapping_mul(1103515245).wrapping_add(12345);
        // Map to [-1, 1]
        (self.state as i32 as f32) / (i32::MAX as f32)
    }
}

/// Generate a mixed excitation frame.
///
/// Blends periodic pulse train and noise per band based on voicing strengths,
/// then sums across bands.
///
/// - `pitch_period`: pitch in samples
/// - `voicing`: per-band voicing strength (0=noise, 1=pulse)
/// - `jitter`: aperiodicity applied to pulse train (0..1)
/// - `output`: buffer to write the excitation signal
/// - `pulse_phase`: mutable phase accumulator for continuity across frames
/// - `noise`: noise generator
pub fn mixed_excitation(
    pitch_period: f32,
    voicing: &[f32; NUM_BANDS],
    jitter: f32,
    output: &mut [f32],
    pulse_phase: &mut f32,
    noise: &mut NoiseGen,
) {
    let n = output.len();

    // Generate raw pulse train and noise
    let mut pulse = [0.0f32; FRAME_SAMPLES];
    let mut noise_buf = [0.0f32; FRAME_SAMPLES];
    let len = n.min(FRAME_SAMPLES);

    for i in 0..len {
        // Noise component
        noise_buf[i] = noise.next();

        // Pulse train with optional jitter
        *pulse_phase += 1.0;
        let jittered_period = pitch_period * (1.0 + jitter * (noise.next() * 0.1));
        if *pulse_phase >= jittered_period {
            *pulse_phase -= jittered_period;
            pulse[i] = sqrtf(pitch_period); // scale pulse amplitude
        }
    }

    // For each band, blend pulse and noise according to voicing strength
    // Simple approach: weight the overall mix by average voicing
    // Full implementation would filter pulse & noise per band then sum
    // Here we do a weighted blend per-band using band center as a proxy

    // Compute overall blend weights at each sample by interpolating band voicing
    for i in 0..len {
        let mut sample = 0.0f32;
        for b in 0..NUM_BANDS {
            let v = voicing[b];
            let band_excitation = v * pulse[i] + (1.0 - v) * noise_buf[i];
            sample += band_excitation;
        }
        output[i] = sample / NUM_BANDS as f32;
    }
}

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

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

    #[test]
    fn test_biquad_bandpass_passes_center() {
        // 1kHz bandpass should pass 1kHz signal
        let mut bp = Biquad::bandpass(1000.0, 500.0, SAMPLE_RATE as f32);
        let input = sine_wave(1000.0, 256);
        let mut output = vec![0.0f32; 256];
        bp.filter_buf(&input, &mut output);

        // After transient (first ~50 samples), output should have significant energy
        let energy: f32 = output[50..].iter().map(|x| x * x).sum::<f32>() / 206.0;
        assert!(energy > 0.01, "Bandpass should pass center freq, energy={}", energy);
    }

    #[test]
    fn test_biquad_bandpass_rejects_far() {
        // 1kHz bandpass should attenuate 3.5kHz signal
        let mut bp = Biquad::bandpass(1000.0, 500.0, SAMPLE_RATE as f32);
        let input = sine_wave(3500.0, 256);
        let mut output = vec![0.0f32; 256];
        bp.filter_buf(&input, &mut output);

        let energy_out: f32 = output[50..].iter().map(|x| x * x).sum::<f32>() / 206.0;
        let energy_in: f32 = input[50..].iter().map(|x| x * x).sum::<f32>() / 206.0;

        assert!(
            energy_out < energy_in * 0.2,
            "Bandpass should reject far freq: in={}, out={}",
            energy_in, energy_out
        );
    }

    #[test]
    fn test_filter_bank_splits() {
        let mut fb = FilterBank::new();
        // 300 Hz signal should mostly appear in band 0 (0-500 Hz)
        let signal = sine_wave(300.0, FRAME_SAMPLES);
        let bands = fb.analyze(&signal);

        let energies = band_energies(&bands, FRAME_SAMPLES);
        // Band 0 should have the most energy
        assert!(
            energies[0] > energies[2],
            "Band 0 should dominate for 300Hz: {:?}",
            energies
        );
        assert!(
            energies[0] > energies[3],
            "Band 0 should dominate for 300Hz: {:?}",
            energies
        );
    }

    #[test]
    fn test_voicing_strength_voiced() {
        let mut fb = FilterBank::new();
        // Periodic signal → high voicing
        let signal = sine_wave(150.0, FRAME_SAMPLES);
        let bands = fb.analyze(&signal);
        let pitch_period = SAMPLE_RATE as f32 / 150.0; // ~53.3 samples

        let voicing = bandpass_voicing_strength(&bands, pitch_period, FRAME_SAMPLES);

        // At least one band (the one containing 150 Hz) should show strong voicing
        let max_v = voicing.iter().cloned().fold(0.0f32, f32::max);
        assert!(max_v > 0.5, "Periodic signal should show voicing: {:?}", voicing);
    }

    #[test]
    fn test_voicing_strength_unvoiced() {
        let mut fb = FilterBank::new();
        // Noise → low voicing
        let mut noise = [0.0f32; FRAME_SAMPLES];
        let mut seed: u32 = 42;
        for s in noise.iter_mut() {
            seed = seed.wrapping_mul(1103515245).wrapping_add(12345);
            *s = (seed as i32 as f32) / (i32::MAX as f32);
        }
        let bands = fb.analyze(&noise);
        let voicing = bandpass_voicing_strength(&bands, 80.0, FRAME_SAMPLES);

        let avg: f32 = voicing.iter().sum::<f32>() / NUM_BANDS as f32;
        assert!(avg < 0.5, "Noise should have low avg voicing: {}", avg);
    }

    #[test]
    fn test_noise_gen_range() {
        let mut ng = NoiseGen::new(12345);
        for _ in 0..1000 {
            let s = ng.next();
            assert!(s >= -1.0 && s <= 1.0, "Noise out of range: {}", s);
        }
    }

    #[test]
    fn test_noise_gen_variance() {
        let mut ng = NoiseGen::new(99);
        let mut sum = 0.0f32;
        let mut sum_sq = 0.0f32;
        let n = 10000;
        for _ in 0..n {
            let s = ng.next();
            sum += s;
            sum_sq += s * s;
        }
        let mean = sum / n as f32;
        let var = sum_sq / n as f32 - mean * mean;
        // Should be roughly uniform: mean~0, variance~1/3
        assert!(mean.abs() < 0.1, "Noise mean should be ~0, got {}", mean);
        assert!(var > 0.1 && var < 0.6, "Noise variance unexpected: {}", var);
    }

    #[test]
    fn test_mixed_excitation_voiced() {
        let mut phase = 0.0f32;
        let mut ng = NoiseGen::new(7);
        let mut output = [0.0f32; FRAME_SAMPLES];
        let voicing = [1.0; NUM_BANDS]; // fully voiced

        mixed_excitation(80.0, &voicing, 0.0, &mut output, &mut phase, &mut ng);

        // Should contain pulses — check for spiky content
        let peak = output.iter().fold(0.0f32, |m, &s| m.max(s.abs()));
        assert!(peak > 0.1, "Voiced excitation should have pulses, peak={}", peak);
    }

    #[test]
    fn test_mixed_excitation_unvoiced() {
        let mut phase = 0.0f32;
        let mut ng = NoiseGen::new(7);
        let mut output = [0.0f32; FRAME_SAMPLES];
        let voicing = [0.0; NUM_BANDS]; // fully unvoiced

        mixed_excitation(80.0, &voicing, 0.0, &mut output, &mut phase, &mut ng);

        // Should be noise-like — check variance is reasonable
        let energy: f32 = output.iter().map(|x| x * x).sum::<f32>() / FRAME_SAMPLES as f32;
        assert!(energy > 0.01, "Unvoiced should have noise energy, got {}", energy);
    }
}