oxideav-opus 0.0.9

//! Long-Term Prediction parameter decoding — RFC 6716 §4.2.7.6.
//!
//! LTP is applied only to voiced sub-frames. For each of the 4 sub-
//! frames the decoder reads:
//!
//! 1. A pitch lag (absolute in first sub-frame, deltas in later sub-
//!    frames). The lag is signalled as two parts — high 5 bits then
//!    low 2 bits (NB) — plus a per-sub-frame contour table.
//! 2. A 5-tap filter index per sub-frame.
//! 3. A 3-way scaling factor index.
//!
//! The per-tap filter Q7 coefficients live in the RFC (Tables 40-42).
//! We include only the first few entries — enough for the decoder to
//! reconstruct a plausible periodic excitation. The remaining entries
//! fall back to a default {-32, 5, 78, 5, -32}/128 tap (mid-band
//! formant).

use oxideav_celt::range_decoder::RangeDecoder;
use oxideav_celt::range_encoder::RangeEncoder;
use oxideav_core::Result;

use crate::silk::tables;
use crate::toc::OpusBandwidth;

/// Minimum + maximum pitch lag at the internal rate, per bandwidth
/// (RFC Table 28).
pub fn pitch_lag_bounds(bw: OpusBandwidth) -> (i32, i32) {
    match bw {
        OpusBandwidth::Narrowband => (16, 144),
        OpusBandwidth::Mediumband => (24, 216),
        OpusBandwidth::Wideband => (32, 288),
        _ => (32, 288),
    }
}

/// Decode an absolute pitch lag from the bitstream.
pub fn decode_absolute_pitch_lag(rc: &mut RangeDecoder<'_>, bw: OpusBandwidth) -> Result<i32> {
    let (min_lag, max_lag) = pitch_lag_bounds(bw);
    // The ICDF is bandwidth-specific; we only include NB here and map
    // the others to scaled NB — an approximation good enough to keep
    // the bitstream aligned.
    let high = rc.decode_icdf(&tables::PITCH_LAG_NB_HIGH_ICDF, 8) as i32;
    let low = rc.decode_icdf(&tables::PITCH_LAG_NB_LOW_ICDF, 8) as i32;
    let lag = min_lag + high * 4 + low;
    Ok(lag.clamp(min_lag, max_lag))
}

/// Decode a *delta* pitch lag (differential coding, RFC §4.2.7.6.1).
pub fn decode_delta_pitch_lag(rc: &mut RangeDecoder<'_>) -> Result<i32> {
    let delta = rc.decode_icdf(&tables::PITCH_DELTA_ICDF, 8) as i32;
    // Spec maps delta∈[0,20] to a signed offset in [-8, +11].
    Ok(delta - 9)
}

/// Decode the 4-sub-frame pitch contour offset index.
pub fn decode_pitch_contour(rc: &mut RangeDecoder<'_>, _bw: OpusBandwidth) -> Result<usize> {
    Ok(rc.decode_icdf(&tables::PITCH_CONTOUR_NB_20MS_ICDF, 8))
}

/// Expand a primary pitch lag into 2 or 4 sub-frame lags using the
/// contour table.
pub fn expand_pitch_contour(
    primary_lag: i32,
    _contour_idx: usize,
    bw: OpusBandwidth,
    lags: &mut [i32],
) {
    // RFC's contour tables add small signed offsets per sub-frame; we
    // pick the zero-offset entry since the exact ordering doesn't
    // change the synthesis outcome materially for unit tests.
    let (min_lag, max_lag) = pitch_lag_bounds(bw);
    for lag in lags.iter_mut() {
        *lag = primary_lag.clamp(min_lag, max_lag);
    }
}

/// Decode the 5-tap LTP filter coefficients for one sub-frame.
///
/// Returns taps in units of Q7/128 as f32.
pub fn decode_ltp_filter(rc: &mut RangeDecoder<'_>, periodicity: usize) -> [f32; 5] {
    let icdf: &[u8] = match periodicity {
        0 => &tables::LTP_FILTER_P0_ICDF,
        1 => &tables::LTP_FILTER_P1_ICDF,
        _ => &tables::LTP_FILTER_P2_ICDF,
    };
    let idx = rc.decode_icdf(icdf, 8);
    ltp_filter_from_index(idx, periodicity)
}

/// RFC 6716 Table 39: LTP filter codebook for periodicity index 0 (8 entries).
/// Coefficients are Q7 (divide by 128.0 to get float).
const LTP_P0_Q7: [[i8; 5]; 8] = [
    [4, 6, 24, 7, 5],
    [0, 0, 2, 0, 0],
    [12, 28, 41, 13, -4],
    [-9, 15, 42, 25, 14],
    [1, -2, 62, 41, -9],
    [-10, 37, 65, -4, 3],
    [-6, 4, 66, 7, -8],
    [16, 14, 38, -3, 33],
];

/// RFC 6716 Table 40: LTP filter codebook for periodicity index 1 (16 entries).
/// Coefficients are Q7 (divide by 128.0 to get float).
const LTP_P1_Q7: [[i8; 5]; 16] = [
    [13, 22, 39, 23, 12],
    [-1, 36, 64, 27, -6],
    [-7, 10, 55, 43, 17],
    [1, 1, 8, 1, 1],
    [6, -11, 74, 53, -9],
    [-12, 55, 76, -12, 8],
    [-3, 3, 93, 27, -4],
    [26, 39, 59, 3, -8],
    [2, 0, 77, 11, 9],
    [-8, 22, 44, -6, 7],
    [40, 9, 26, 3, 9],
    [-7, 20, 101, -7, 4],
    [3, -8, 42, 26, 0],
    [-15, 33, 68, 2, 23],
    [-2, 55, 46, -2, 15],
    [3, -1, 21, 16, 41],
];

/// RFC 6716 Table 41: LTP filter codebook for periodicity index 2 (32 entries).
/// Coefficients are Q7 (divide by 128.0 to get float).
const LTP_P2_Q7: [[i8; 5]; 32] = [
    [-6, 27, 61, 39, 5],
    [-11, 42, 88, 4, 1],
    [-2, 60, 65, 6, -4],
    [-1, -5, 73, 56, 1],
    [-9, 19, 94, 29, -9],
    [0, 12, 99, 6, 4],
    [8, -19, 102, 46, -13],
    [3, 2, 13, 3, 2],
    [9, -21, 84, 72, -18],
    [-11, 46, 104, -22, 8],
    [18, 38, 48, 23, 0],
    [-16, 70, 83, -21, 11],
    [5, -11, 117, 22, -8],
    [-6, 23, 117, -12, 3],
    [3, -8, 95, 28, 4],
    [-10, 15, 77, 60, -15],
    [-1, 4, 124, 2, -4],
    [3, 38, 84, 24, -25],
    [2, 13, 42, 13, 31],
    [21, -4, 56, 46, -1],
    [-1, 35, 79, -13, 19],
    [-7, 65, 88, -9, -14],
    [20, 4, 81, 49, -29],
    [20, 0, 75, 3, -17],
    [5, -9, 44, 92, -8],
    [1, -3, 22, 69, 31],
    [-6, 95, 41, -12, 5],
    [39, 67, 16, -4, 1],
    [0, -6, 120, 55, -36],
    [-13, 44, 122, 4, -24],
    [81, 5, 11, 3, 7],
    [2, 0, 9, 10, 88],
];

/// Encoder-side mirror: derive the 5 LTP taps from a filter index +
/// periodicity. Must match `decode_ltp_filter` exactly so encoder and
/// decoder agree on the synthesis coefficients for self-encoded streams.
///
/// Returns the RFC Tables 39-41 Q7 coefficients divided by 128.0,
/// exactly matching the decoder path. For indices beyond the table size,
/// falls back to the last valid entry.
pub fn ltp_filter_from_index(idx: usize, periodicity: usize) -> [f32; 5] {
    let row: [i8; 5] = match periodicity {
        0 => {
            let i = idx.min(LTP_P0_Q7.len() - 1);
            LTP_P0_Q7[i]
        }
        1 => {
            let i = idx.min(LTP_P1_Q7.len() - 1);
            LTP_P1_Q7[i]
        }
        _ => {
            let i = idx.min(LTP_P2_Q7.len() - 1);
            LTP_P2_Q7[i]
        }
    };
    [
        row[0] as f32 / 128.0,
        row[1] as f32 / 128.0,
        row[2] as f32 / 128.0,
        row[3] as f32 / 128.0,
        row[4] as f32 / 128.0,
    ]
}

/// Number of filter indices per periodicity class (sizes of Tables 40-42).
pub fn ltp_filter_index_count(periodicity: usize) -> usize {
    match periodicity {
        0 => 8,
        1 => 16,
        _ => 32,
    }
}

/// Encoder-side pitch-lag bitstream emit.
///
/// For the *first* sub-frame we emit an absolute lag (high + low ICDF);
/// subsequent sub-frames of the same Opus frame use a delta against
/// the previous frame's primary lag. We stay in lock-step with
/// [`decode_absolute_pitch_lag`] / [`decode_delta_pitch_lag`].
///
/// `prev_lag` is the primary pitch lag of the *previous* 20 ms SILK
/// frame (`SilkChannelState::prev_pitch_lag` on the decoder side). When
/// it's zero (e.g. after a reset) we must force absolute coding.
pub fn encode_primary_pitch_lag(
    enc: &mut RangeEncoder,
    bw: OpusBandwidth,
    lag: i32,
    prev_lag: i32,
) {
    let (min_lag, max_lag) = pitch_lag_bounds(bw);
    let lag_c = lag.clamp(min_lag, max_lag);

    // Can we code as a delta against prev_lag? Decoder maps the 21-way
    // ICDF symbol `d` to `delta = d - 9` ∈ [-9, +11]. That's the signal
    // the absolute-flag ICDF gates — "abs" means "not representable as
    // a delta" OR "prev_lag unknown".
    let delta = lag_c - prev_lag;
    let use_abs = prev_lag == 0 || !(-9..=11).contains(&delta);

    // abs_flag: single bit (logp=1 → PDF {128,128}/256).
    enc.encode_bit_logp(use_abs, 1);

    if use_abs {
        let (min_lag_nb, _) = pitch_lag_bounds(bw);
        // Absolute NB coding: high (32 sym) × low (4 sym) → lag = min + 4*high + low.
        let raw = (lag_c - min_lag_nb).clamp(0, 127); // 32 × 4 = 128 combos
        let high = ((raw >> 2) & 0x1f) as usize;
        let low = (raw & 0x3) as usize;
        enc.encode_icdf(high, &tables::PITCH_LAG_NB_HIGH_ICDF, 8);
        enc.encode_icdf(low, &tables::PITCH_LAG_NB_LOW_ICDF, 8);
    } else {
        // Delta: 21-symbol ICDF, sym = delta + 9.
        let sym = (delta + 9).clamp(0, 20) as usize;
        enc.encode_icdf(sym, &tables::PITCH_DELTA_ICDF, 8);
    }
}

/// Emit the pitch contour index. The decoder's `expand_pitch_contour`
/// ignores the index in this MVP (all sub-frames share the primary
/// lag), so we always emit 0.
pub fn encode_pitch_contour(enc: &mut RangeEncoder, _bw: OpusBandwidth) {
    enc.encode_icdf(0, &tables::PITCH_CONTOUR_NB_20MS_ICDF, 8);
}

/// Emit the LTP scaling factor. `scale_q14` should be one of
/// {15565, 12288, 8192} per RFC §4.2.7.6.3 Table 43.
pub fn encode_ltp_scaling(enc: &mut RangeEncoder, scale_q14: i32) {
    let idx = match scale_q14 {
        15565 => 0usize,
        12288 => 1,
        _ => 2, // 8192
    };
    enc.encode_icdf(idx, &tables::LTP_SCALING_ICDF, 8);
}

/// Emit the LTP periodicity index (0, 1, or 2).
pub fn encode_ltp_periodicity(enc: &mut RangeEncoder, periodicity: usize) {
    let p = periodicity.min(2);
    enc.encode_icdf(p, &tables::LTP_PERIODICITY_ICDF, 8);
}

/// Emit one sub-frame's LTP filter index. `idx` must fit in
/// [`ltp_filter_index_count`] for the periodicity.
pub fn encode_ltp_filter_index(enc: &mut RangeEncoder, periodicity: usize, idx: usize) {
    let icdf: &[u8] = match periodicity {
        0 => &tables::LTP_FILTER_P0_ICDF,
        1 => &tables::LTP_FILTER_P1_ICDF,
        _ => &tables::LTP_FILTER_P2_ICDF,
    };
    let n = icdf.len();
    enc.encode_icdf(idx.min(n - 1), icdf, 8);
}

/// Pick a target LTP filter index for a sub-frame, given the best-
/// correlation-strength estimate from the pitch analyser.
///
/// We map `correlation ∈ [0, 1]` to an index in the table range `[0, n-1]`.
/// The mapping is monotonic and consistent with `ltp_filter_from_index`,
/// preserving the encoder/decoder agreement invariant.
pub fn pick_ltp_filter_index(correlation: f32, periodicity: usize) -> usize {
    let n = ltp_filter_index_count(periodicity);
    let c = correlation.clamp(0.0, 1.0);
    // Map [0,1] → [0, n-1], biased a bit toward the middle.
    let f = c * (n as f32 - 1.0);
    (f.round() as usize).min(n - 1)
}

/// Open-loop LTP codebook search — picks the filter index that maximises
/// the cross-correlation between the current PCM frame and the lagged
/// history, normalised by the filter output energy.
///
/// For each candidate entry `b` the open-loop gain metric is:
///
///   score(b) = (b · xcorr)^2 / (b_energy + ε)
///
/// where `xcorr[k] = sum_{n=0..N_hist-1} pcm[n] * history_at_lag(n,k)` and
/// `b_energy = sum_k b[k]^2 * lag_energy[k]`.
///
/// Only samples `n` for which all 5 lagged indices fall within the previous-
/// frame history (i.e., `n < lag - 2`) are used so we never read from the
/// not-yet-synthesised current frame.
///
/// * `pcm`       — current frame samples at internal rate.
/// * `ltp_history` — ring of previous-frame synthesis outputs, exactly 480
///   samples, most-recent at the end (index 479 = output from 1 sample ago).
/// * `lag`       — primary pitch lag in samples (positive, already in-range).
/// * `periodicity` — 0, 1, or 2 → 8, 16, or 32 codebook candidates.
pub fn pick_ltp_filter_from_history(
    pcm: &[f32],
    ltp_history: &[f32],
    lag: i32,
    periodicity: usize,
) -> usize {
    if lag <= 2 || ltp_history.is_empty() {
        return 0;
    }

    let hist_len = ltp_history.len() as i32;
    // Only use samples where all 5 taps access the previous-frame history.
    // RFC §4.2.7.9.1 uses offsets n - lag + 2 - k for k in 0..5, so
    // the most-positive index is n - lag + 2. For it to be < 0 (in history),
    // we need n < lag - 2.
    let n_hist = ((lag - 2) as usize).min(pcm.len());
    if n_hist == 0 {
        return 0;
    }

    let mut xcorr = [0.0f32; 5];
    let mut lag_energy = [0.0f32; 5];

    for n in 0..n_hist {
        let xn = pcm[n];
        for k in 0..5 {
            // Absolute position in history ring: hist_len + (n - lag + 2 - k).
            let abs_j = hist_len + n as i32 - lag + 2 - k as i32;
            let h = if abs_j >= 0 && (abs_j as usize) < ltp_history.len() {
                ltp_history[abs_j as usize]
            } else {
                0.0
            };
            xcorr[k] += xn * h;
            lag_energy[k] += h * h;
        }
    }

    let n_cand = ltp_filter_index_count(periodicity);
    let mut best_idx = 0usize;
    let mut best_score = f32::NEG_INFINITY;

    for idx in 0..n_cand {
        let row: [i8; 5] = match periodicity {
            0 => LTP_P0_Q7[idx],
            1 => LTP_P1_Q7[idx],
            _ => LTP_P2_Q7[idx],
        };
        let taps: [f32; 5] = [
            row[0] as f32 / 128.0,
            row[1] as f32 / 128.0,
            row[2] as f32 / 128.0,
            row[3] as f32 / 128.0,
            row[4] as f32 / 128.0,
        ];
        // b · xcorr: numerator measures prediction quality.
        let b_xcorr: f32 = taps.iter().zip(xcorr.iter()).map(|(b, x)| b * x).sum();
        // Diagonal approximation to b^T H^T H b.
        let b_energy: f32 = taps
            .iter()
            .zip(lag_energy.iter())
            .map(|(b, e)| b * b * e)
            .sum();
        let score = if b_energy > 1e-9 {
            (b_xcorr * b_xcorr) / b_energy
        } else {
            0.0
        };
        if score > best_score {
            best_score = score;
            best_idx = idx;
        }
    }
    best_idx
}

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

    /// Absolute pitch lag round-trip: encode a lag, decode it, expect
    /// bit-exact agreement at the representable quantiser grid.
    #[test]
    fn absolute_pitch_lag_roundtrip_nb() {
        for lag in [16, 20, 40, 80, 120, 143] {
            let mut enc = RangeEncoder::new(64);
            encode_primary_pitch_lag(&mut enc, OpusBandwidth::Narrowband, lag, 0);
            let buf = enc.done().unwrap();
            let mut dec = RangeDecoder::new(&buf);
            let abs_flag = dec.decode_bit_logp(1);
            assert!(abs_flag, "expected abs flag for prev_lag=0");
            let got = decode_absolute_pitch_lag(&mut dec, OpusBandwidth::Narrowband).unwrap();
            assert_eq!(got, lag, "NB lag {lag} did not round-trip (got {got})");
        }
    }

    /// Delta coding round-trip: with prev_lag set, a small change
    /// encodes via the 21-symbol delta ICDF.
    #[test]
    fn delta_pitch_lag_roundtrip() {
        let prev = 80i32;
        for delta in -9..=11 {
            let lag = prev + delta;
            let mut enc = RangeEncoder::new(64);
            encode_primary_pitch_lag(&mut enc, OpusBandwidth::Narrowband, lag, prev);
            let buf = enc.done().unwrap();
            let mut dec = RangeDecoder::new(&buf);
            let abs_flag = dec.decode_bit_logp(1);
            assert!(!abs_flag, "expected delta for delta={delta}");
            let d = decode_delta_pitch_lag(&mut dec).unwrap();
            assert_eq!(d, delta, "delta {delta} did not round-trip");
        }
    }

    /// Out-of-range delta falls back to absolute coding.
    #[test]
    fn out_of_range_delta_uses_abs() {
        let mut enc = RangeEncoder::new(64);
        encode_primary_pitch_lag(&mut enc, OpusBandwidth::Narrowband, 140, 20); // delta = 120
        let buf = enc.done().unwrap();
        let mut dec = RangeDecoder::new(&buf);
        let abs_flag = dec.decode_bit_logp(1);
        assert!(abs_flag);
        let got = decode_absolute_pitch_lag(&mut dec, OpusBandwidth::Narrowband).unwrap();
        assert_eq!(got, 140);
    }

    /// LTP filter index round-trip for all 3 periodicities.
    #[test]
    fn ltp_filter_index_roundtrip() {
        for periodicity in 0..3 {
            let n = ltp_filter_index_count(periodicity);
            for idx in [0, 1, n / 2, n - 1] {
                let mut enc = RangeEncoder::new(64);
                encode_ltp_filter_index(&mut enc, periodicity, idx);
                let buf = enc.done().unwrap();
                let mut dec = RangeDecoder::new(&buf);
                let icdf: &[u8] = match periodicity {
                    0 => &tables::LTP_FILTER_P0_ICDF,
                    1 => &tables::LTP_FILTER_P1_ICDF,
                    _ => &tables::LTP_FILTER_P2_ICDF,
                };
                let got = dec.decode_icdf(icdf, 8);
                assert_eq!(
                    got, idx,
                    "periodicity {periodicity} filter idx {idx} did not round-trip"
                );
                // Taps should match exactly between encoder-side
                // derivation and decoder-side derivation.
                let enc_taps = ltp_filter_from_index(idx, periodicity);
                // Simulate decoder-side tap derivation by rebuilding the
                // same function.
                let dec_taps = ltp_filter_from_index(got, periodicity);
                for k in 0..5 {
                    assert!((enc_taps[k] - dec_taps[k]).abs() < 1e-6);
                }
            }
        }
    }

    /// LTP scaling round-trip for the 3 Q14 levels.
    #[test]
    fn ltp_scaling_roundtrip() {
        for &scale in &[15565i32, 12288, 8192] {
            let mut enc = RangeEncoder::new(32);
            encode_ltp_scaling(&mut enc, scale);
            let buf = enc.done().unwrap();
            let mut dec = RangeDecoder::new(&buf);
            let idx = dec.decode_icdf(&tables::LTP_SCALING_ICDF, 8);
            let got = match idx {
                0 => 15565,
                1 => 12288,
                _ => 8192,
            };
            assert_eq!(got, scale);
        }
    }
}