melpe-rs 0.1.2

/// Quantization for 600 bps MELPe superframes
///
/// Scalar quantization of LSFs, log-domain pitch and gain,
/// binary voicing collapse, and LSF interpolation for the
/// 3-frame superframe structure.

use crate::core_types::{
    NUM_LSF, NUM_BANDS, SUPERFRAME_FRAMES,
    LSF_BITS_600, PITCH_BITS_600, GAIN_BITS_600,
    PITCH_MIN, PITCH_MAX,
    FrameParams, SuperFrame, SuperFrameQuantized,
};
use crate::math::{logf, expf, roundf, clampf};

// ── LSF quantization ──

/// LSF range: each LSF lives in (0, PI). We define per-coefficient
/// min/max boundaries to improve quantization efficiency.
/// Ranges cover typical speech LSFs plus margin for the default (evenly-spaced) set.
const LSF_MIN: [f32; NUM_LSF] = [
    0.05, 0.15, 0.30, 0.45, 0.60, 0.80, 1.00, 1.25, 1.55, 1.85,
];
const LSF_MAX: [f32; NUM_LSF] = [
    0.55, 0.85, 1.15, 1.45, 1.75, 2.05, 2.35, 2.65, 2.90, 3.10,
];

/// Minimum LSF spacing in radians to guarantee filter stability
const LSF_MIN_SPACING: f32 = 0.03;

/// Scalar quantize a single LSF value to `n_bits` levels within [min, max].
fn quantize_uniform(value: f32, min: f32, max: f32, n_bits: usize) -> u8 {
    let levels = (1u32 << n_bits) as f32;
    let clamped = clampf(value, min, max);
    let normalized = (clamped - min) / (max - min); // 0.0..1.0
    let index = roundf(normalized * (levels - 1.0)) as u8;
    index.min((levels as u8) - 1)
}

/// Dequantize a uniform scalar index back to a float value.
fn dequantize_uniform(index: u8, min: f32, max: f32, n_bits: usize) -> f32 {
    let levels = (1u32 << n_bits) as f32;
    let normalized = index as f32 / (levels - 1.0);
    min + normalized * (max - min)
}

/// Quantize an array of 10 LSFs using per-coefficient bit allocations.
pub fn quantize_lsf(lsf: &[f32; NUM_LSF]) -> [u8; NUM_LSF] {
    let mut indices = [0u8; NUM_LSF];
    for i in 0..NUM_LSF {
        indices[i] = quantize_uniform(lsf[i], LSF_MIN[i], LSF_MAX[i], LSF_BITS_600[i]);
    }
    indices
}

/// Dequantize LSF indices back to LSF values, then enforce ordering.
pub fn dequantize_lsf(indices: &[u8; NUM_LSF]) -> [f32; NUM_LSF] {
    let mut lsf = [0.0f32; NUM_LSF];
    for i in 0..NUM_LSF {
        lsf[i] = dequantize_uniform(indices[i], LSF_MIN[i], LSF_MAX[i], LSF_BITS_600[i]);
    }
    enforce_lsf_ordering(&mut lsf);
    lsf
}

/// Enforce strictly increasing LSF ordering with minimum spacing.
/// Pushes LSFs apart if they're too close after quantization.
pub fn enforce_lsf_ordering(lsf: &mut [f32; NUM_LSF]) {
    // Forward pass: ensure each LSF is at least MIN_SPACING above the previous
    for i in 1..NUM_LSF {
        let floor = lsf[i - 1] + LSF_MIN_SPACING;
        if lsf[i] < floor {
            lsf[i] = floor;
        }
    }
    // Clamp the last one below PI
    let pi = core::f32::consts::PI;
    if lsf[NUM_LSF - 1] > pi - 0.02 {
        lsf[NUM_LSF - 1] = pi - 0.02;
        // Backward pass if we hit the ceiling
        for i in (0..NUM_LSF - 1).rev() {
            let ceiling = lsf[i + 1] - LSF_MIN_SPACING;
            if lsf[i] > ceiling {
                lsf[i] = ceiling;
            }
        }
    }
}

// ── LSF interpolation for 600 bps ──

/// Interpolation weights for deriving frames 0 and 1 from
/// the previous superframe's anchor and the current anchor.
/// Frame 0 weight on current anchor (rest is previous anchor):
const INTERP_WEIGHT: [f32; SUPERFRAME_FRAMES] = [
    0.333, // frame 0: 1/3 current + 2/3 previous
    0.667, // frame 1: 2/3 current + 1/3 previous
    1.000, // frame 2: anchor itself (transmitted)
];

/// Interpolate LSFs for all 3 frames given the current anchor and previous anchor.
pub fn interpolate_lsf(
    prev_anchor: &[f32; NUM_LSF],
    curr_anchor: &[f32; NUM_LSF],
) -> [[f32; NUM_LSF]; SUPERFRAME_FRAMES] {
    let mut result = [[0.0f32; NUM_LSF]; SUPERFRAME_FRAMES];

    for f in 0..SUPERFRAME_FRAMES {
        let w = INTERP_WEIGHT[f];
        for i in 0..NUM_LSF {
            result[f][i] = (1.0 - w) * prev_anchor[i] + w * curr_anchor[i];
        }
        enforce_lsf_ordering(&mut result[f]);
    }

    result
}

// ── Pitch quantization ──

/// Log-domain pitch quantization.
/// Maps [PITCH_MIN, PITCH_MAX] to [0, 2^PITCH_BITS - 1] in log scale.
pub fn quantize_pitch(period: f32) -> u8 {
    let min_log = logf(PITCH_MIN as f32);
    let max_log = logf(PITCH_MAX as f32);
    let levels = (1u32 << PITCH_BITS_600) as f32;

    let clamped = clampf(period, PITCH_MIN as f32, PITCH_MAX as f32);
    let log_val = logf(clamped);
    let normalized = (log_val - min_log) / (max_log - min_log);
    let index = roundf(normalized * (levels - 1.0)) as u8;
    index.min((levels as u8) - 1)
}

/// Dequantize pitch index back to period in samples.
pub fn dequantize_pitch(index: u8) -> f32 {
    let min_log = logf(PITCH_MIN as f32);
    let max_log = logf(PITCH_MAX as f32);
    let levels = (1u32 << PITCH_BITS_600) as f32;

    let normalized = index as f32 / (levels - 1.0);
    expf(min_log + normalized * (max_log - min_log))
}

// ── Gain quantization ──

/// Gain range in dB for quantization
const GAIN_MIN_DB: f32 = -60.0;
const GAIN_MAX_DB: f32 = 0.0;

/// Quantize gain (dB) to GAIN_BITS_600-bit index.
pub fn quantize_gain(gain_db: f32) -> u8 {
    quantize_uniform(gain_db, GAIN_MIN_DB, GAIN_MAX_DB, GAIN_BITS_600)
}

/// Dequantize gain index back to dB.
pub fn dequantize_gain(index: u8) -> f32 {
    dequantize_uniform(index, GAIN_MIN_DB, GAIN_MAX_DB, GAIN_BITS_600)
}

// ── Voicing quantization ──

/// Collapse 5-band voicing to binary V/UV decision.
/// Returns true (voiced) if mean bandpass voicing >= threshold.
pub fn collapse_voicing(bandpass_voicing: &[f32; NUM_BANDS], threshold: f32) -> bool {
    let mean: f32 = bandpass_voicing.iter().sum::<f32>() / NUM_BANDS as f32;
    mean >= threshold
}

/// Pack 3 per-frame V/UV decisions + 1 transition bit into a u8.
/// Bit layout: [V0, V1, V2, T] where T = transition detected.
/// Transition bit is set if voicing changes across the superframe.
pub fn pack_voicing(voiced: &[bool; SUPERFRAME_FRAMES]) -> u8 {
    let mut bits: u8 = 0;
    for i in 0..SUPERFRAME_FRAMES {
        if voiced[i] {
            bits |= 1 << i;
        }
    }
    // Transition bit: set if not all frames agree
    let all_same = voiced[0] == voiced[1] && voiced[1] == voiced[2];
    if !all_same {
        bits |= 1 << 3;
    }
    bits
}

/// Unpack voicing bits back to per-frame V/UV decisions.
pub fn unpack_voicing(bits: u8) -> ([bool; SUPERFRAME_FRAMES], bool) {
    let voiced = [
        (bits & (1 << 0)) != 0,
        (bits & (1 << 1)) != 0,
        (bits & (1 << 2)) != 0,
    ];
    let transition = (bits & (1 << 3)) != 0;
    (voiced, transition)
}

// ── Full superframe quantization ──

/// Voicing threshold for 600 bps binary collapse
const VOICING_THRESHOLD_600: f32 = 0.35;

/// Quantize an entire SuperFrame into a SuperFrameQuantized.
pub fn quantize_superframe(sf: &SuperFrame) -> SuperFrameQuantized {
    let anchor = sf.anchor();

    let lsf_indices = quantize_lsf(&anchor.lsf);
    let pitch_index = quantize_pitch(anchor.pitch);

    // Use anchor gain
    let gain_index = quantize_gain(anchor.gain);

    // Binary voicing per frame
    let voiced = sf.binary_voicing(VOICING_THRESHOLD_600);
    let voicing_bits = pack_voicing(&voiced);

    // Jitter: simple 1-bit threshold
    let jitter_bit = if anchor.jitter > 0.5 { 1 } else { 0 };

    SuperFrameQuantized {
        lsf_indices,
        pitch_index,
        gain_index,
        voicing_bits,
        jitter_bit,
    }
}

/// Dequantize a SuperFrameQuantized back into 3 FrameParams.
/// Requires the previous superframe's anchor LSFs for interpolation.
pub fn dequantize_superframe(
    sq: &SuperFrameQuantized,
    prev_anchor_lsf: &[f32; NUM_LSF],
) -> [FrameParams; SUPERFRAME_FRAMES] {
    // Dequantize anchor parameters
    let anchor_lsf = dequantize_lsf(&sq.lsf_indices);
    let anchor_pitch = dequantize_pitch(sq.pitch_index);
    let anchor_gain = dequantize_gain(sq.gain_index);
    let (voiced, _transition) = unpack_voicing(sq.voicing_bits);
    let jitter = if sq.jitter_bit > 0 { 0.75 } else { 0.0 };

    // Interpolate LSFs across the 3 frames
    let interp_lsf = interpolate_lsf(prev_anchor_lsf, &anchor_lsf);

    let mut frames = [
        FrameParams::default(),
        FrameParams::default(),
        FrameParams::default(),
    ];

    for i in 0..SUPERFRAME_FRAMES {
        frames[i].lsf = interp_lsf[i];
        frames[i].pitch = anchor_pitch; // same pitch for all 3 at 600 bps
        frames[i].gain = anchor_gain;   // could interpolate, but 600 bps keeps it simple
        frames[i].jitter = jitter;

        // Set bandpass voicing: all bands same as binary decision
        let v = if voiced[i] { 1.0 } else { 0.0 };
        frames[i].bandpass_voicing = [v; NUM_BANDS];
    }

    frames
}

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

    // ── Uniform quantizer tests ──

    #[test]
    fn test_uniform_quant_midpoint() {
        // Mid-range value with 3 bits (8 levels)
        let idx = quantize_uniform(0.5, 0.0, 1.0, 3);
        assert_eq!(idx, 4); // 0.5 * 7 = 3.5 → rounds to 4
        let val = dequantize_uniform(idx, 0.0, 1.0, 3);
        assert!((val - 0.571).abs() < 0.01); // 4/7 ≈ 0.571
    }

    #[test]
    fn test_uniform_quant_boundaries() {
        // At minimum
        let idx = quantize_uniform(0.0, 0.0, 1.0, 3);
        assert_eq!(idx, 0);
        let val = dequantize_uniform(0, 0.0, 1.0, 3);
        assert!((val - 0.0).abs() < 1e-6);

        // At maximum
        let idx = quantize_uniform(1.0, 0.0, 1.0, 3);
        assert_eq!(idx, 7);
        let val = dequantize_uniform(7, 0.0, 1.0, 3);
        assert!((val - 1.0).abs() < 1e-6);
    }

    #[test]
    fn test_uniform_quant_clamps() {
        // Below min → clamps to 0
        let idx = quantize_uniform(-5.0, 0.0, 1.0, 3);
        assert_eq!(idx, 0);
        // Above max → clamps to max index
        let idx = quantize_uniform(99.0, 0.0, 1.0, 3);
        assert_eq!(idx, 7);
    }

    // ── LSF quantization tests ──

    #[test]
    fn test_lsf_quantize_roundtrip() {
        let lsf = default_lsf();
        let indices = quantize_lsf(&lsf);
        let recovered = dequantize_lsf(&indices);

        // Should be close (within quantization step)
        for i in 0..NUM_LSF {
            let range = LSF_MAX[i] - LSF_MIN[i];
            let step = range / ((1u32 << LSF_BITS_600[i]) - 1) as f32;
            assert!(
                (lsf[i] - recovered[i]).abs() < step + LSF_MIN_SPACING + 0.01,
                "LSF[{}] roundtrip: {} → {} (step={})",
                i, lsf[i], recovered[i], step
            );
        }
    }

    #[test]
    fn test_lsf_ordering_preserved() {
        let lsf = default_lsf();
        let indices = quantize_lsf(&lsf);
        let recovered = dequantize_lsf(&indices);

        // Must be strictly increasing after dequantization
        for i in 1..NUM_LSF {
            assert!(
                recovered[i] > recovered[i - 1],
                "LSF ordering broken at [{}]: {} <= {}",
                i, recovered[i], recovered[i - 1]
            );
        }
    }

    #[test]
    fn test_enforce_lsf_ordering_fixes_collapse() {
        // Deliberately break ordering
        let mut lsf = [0.3; NUM_LSF]; // all the same
        enforce_lsf_ordering(&mut lsf);

        for i in 1..NUM_LSF {
            assert!(
                lsf[i] > lsf[i - 1],
                "enforce_lsf_ordering failed at [{}]",
                i
            );
        }
        assert!(lsf[NUM_LSF - 1] < core::f32::consts::PI);
    }

    // ── LSF interpolation tests ──

    #[test]
    fn test_interpolation_anchor_exact() {
        let prev = default_lsf();
        let curr = default_lsf();
        let result = interpolate_lsf(&prev, &curr);

        // When prev == curr, all frames should match
        for f in 0..SUPERFRAME_FRAMES {
            for i in 0..NUM_LSF {
                assert!(
                    (result[f][i] - curr[i]).abs() < LSF_MIN_SPACING + 0.01,
                    "Frame {} LSF[{}] mismatch",
                    f, i
                );
            }
        }
    }

    #[test]
    fn test_interpolation_weights() {
        let mut prev = [0.0f32; NUM_LSF];
        let mut curr = [0.0f32; NUM_LSF];
        // Set up widely spaced values to see interpolation clearly
        for i in 0..NUM_LSF {
            prev[i] = 0.1 * (i as f32 + 1.0);
            curr[i] = 0.2 * (i as f32 + 1.0);
        }

        let result = interpolate_lsf(&prev, &curr);

        // Frame 2 (anchor) should be very close to curr
        for i in 0..NUM_LSF {
            assert!(
                (result[2][i] - curr[i]).abs() < LSF_MIN_SPACING + 0.01,
                "Anchor frame should match curr at [{}]",
                i
            );
        }

        // Frame 0 should be closer to prev than frame 1
        let dist0: f32 = (0..NUM_LSF).map(|i| (result[0][i] - prev[i]).abs()).sum();
        let dist1: f32 = (0..NUM_LSF).map(|i| (result[1][i] - prev[i]).abs()).sum();
        assert!(
            dist0 < dist1,
            "Frame 0 should be closer to prev: d0={}, d1={}",
            dist0, dist1
        );
    }

    #[test]
    fn test_interpolation_ordering_maintained() {
        let prev = default_lsf();
        let mut curr = default_lsf();
        // Shift current anchor up a bit
        for i in 0..NUM_LSF {
            curr[i] += 0.05;
        }
        let result = interpolate_lsf(&prev, &curr);

        for f in 0..SUPERFRAME_FRAMES {
            for i in 1..NUM_LSF {
                assert!(
                    result[f][i] > result[f][i - 1],
                    "Interpolated frame {} LSF ordering broken at [{}]",
                    f, i
                );
            }
        }
    }

    // ── Pitch quantization tests ──

    #[test]
    fn test_pitch_roundtrip_min() {
        let idx = quantize_pitch(PITCH_MIN as f32);
        assert_eq!(idx, 0);
        let recovered = dequantize_pitch(0);
        assert!((recovered - PITCH_MIN as f32).abs() < 1.0);
    }

    #[test]
    fn test_pitch_roundtrip_max() {
        let idx = quantize_pitch(PITCH_MAX as f32);
        assert_eq!(idx, (1 << PITCH_BITS_600) - 1);
        let recovered = dequantize_pitch(idx);
        assert!((recovered - PITCH_MAX as f32).abs() < 1.0);
    }

    #[test]
    fn test_pitch_roundtrip_mid() {
        let mid = 80.0f32; // typical male pitch period
        let idx = quantize_pitch(mid);
        let recovered = dequantize_pitch(idx);
        // Log quantization: relative error matters more than absolute
        let rel_error = (recovered - mid).abs() / mid;
        assert!(
            rel_error < 0.05,
            "Pitch roundtrip: {} → idx {} → {} (rel_err={})",
            mid, idx, recovered, rel_error
        );
    }

    #[test]
    fn test_pitch_clamps() {
        let idx_lo = quantize_pitch(5.0); // below PITCH_MIN
        assert_eq!(idx_lo, 0);
        let idx_hi = quantize_pitch(500.0); // above PITCH_MAX
        assert_eq!(idx_hi, (1 << PITCH_BITS_600) - 1);
    }

    // ── Gain quantization tests ──

    #[test]
    fn test_gain_roundtrip() {
        let gain = -30.0f32;
        let idx = quantize_gain(gain);
        let recovered = dequantize_gain(idx);
        let step = (GAIN_MAX_DB - GAIN_MIN_DB) / ((1u32 << GAIN_BITS_600) - 1) as f32;
        assert!(
            (recovered - gain).abs() < step + 0.1,
            "Gain roundtrip: {} → {} (step={})",
            gain, recovered, step
        );
    }

    #[test]
    fn test_gain_boundaries() {
        assert_eq!(quantize_gain(-60.0), 0);
        assert_eq!(quantize_gain(0.0), (1 << GAIN_BITS_600) - 1);
    }

    // ── Voicing tests ──

    #[test]
    fn test_collapse_voicing_voiced() {
        let bpv = [0.8, 0.9, 0.7, 0.6, 0.85];
        assert!(collapse_voicing(&bpv, 0.35));
    }

    #[test]
    fn test_collapse_voicing_unvoiced() {
        let bpv = [0.1, 0.2, 0.0, 0.1, 0.05];
        assert!(!collapse_voicing(&bpv, 0.35));
    }

    #[test]
    fn test_voicing_pack_unpack_roundtrip() {
        let cases: [[bool; 3]; 8] = [
            [false, false, false],
            [true, false, false],
            [false, true, false],
            [false, false, true],
            [true, true, false],
            [true, false, true],
            [false, true, true],
            [true, true, true],
        ];

        for voiced in &cases {
            let packed = pack_voicing(voiced);
            let (recovered, transition) = unpack_voicing(packed);
            assert_eq!(
                &recovered, voiced,
                "Voicing roundtrip failed for {:?}",
                voiced
            );

            let expected_transition = !(voiced[0] == voiced[1] && voiced[1] == voiced[2]);
            assert_eq!(transition, expected_transition);
        }
    }

    // ── Full superframe quantization tests ──

    #[test]
    fn test_superframe_quantize_dequantize() {
        let sf = SuperFrame::default();
        let sq = quantize_superframe(&sf);
        let prev_lsf = default_lsf();
        let frames = dequantize_superframe(&sq, &prev_lsf);

        // All frames should have valid LSF ordering
        for f in 0..SUPERFRAME_FRAMES {
            for i in 1..NUM_LSF {
                assert!(
                    frames[f].lsf[i] > frames[f].lsf[i - 1],
                    "Dequantized frame {} LSF ordering broken",
                    f
                );
            }
            // Gain should be finite
            assert!(frames[f].gain.is_finite());
            // Pitch in valid range
            assert!(frames[f].pitch >= PITCH_MIN as f32 - 1.0);
            assert!(frames[f].pitch <= PITCH_MAX as f32 + 1.0);
        }
    }

    #[test]
    fn test_superframe_voiced_propagates() {
        let mut sf = SuperFrame::default();
        // Make frame 1 voiced
        sf.frames[1].bandpass_voicing = [0.9; NUM_BANDS];

        let sq = quantize_superframe(&sf);
        let prev_lsf = default_lsf();
        let frames = dequantize_superframe(&sq, &prev_lsf);

        // Frame 1 should be voiced (all bands = 1.0)
        assert_eq!(frames[1].bandpass_voicing, [1.0; NUM_BANDS]);
        // Frame 0 should be unvoiced
        assert_eq!(frames[0].bandpass_voicing, [0.0; NUM_BANDS]);
    }
}