lc3-codec 0.2.0

use crate::{
    common::{complex::Scaler, config::FrameDuration},
    tables::temporal_noise_shaping_tables::{MAXLAG, TNS_NUMFILTERS_MAX},
};
use core::f64::consts::PI;
use heapless::Vec;
#[allow(unused_imports)]
use num_traits::real::Real;

use super::side_info::Bandwidth;

// checked against spec

/// Apply temportal noise shaping for decoded spectral lines
///
/// # Arguments
///
/// * `duration` - Frame Duration in milliseconds (N_MS)
/// * `bandwidth` - Bandwidth cutoff index (P_BW)
/// * `num_tns_filters` - Number of TNS filters
/// * `reflect_coef_order` - Order of quantized reflection coefficients
/// * `reflect_coef_ints` - Integer reflection coefficients
/// * `spec_lines` - Spectral lines of the frame to mutate (input and output)
pub fn apply_temporal_noise_shaping(
    duration: FrameDuration,
    bandwidth: Bandwidth,
    num_tns_filters: usize,
    reflect_coef_order: &[usize],
    reflect_coef_ints: &[usize],
    spec_lines: &mut [Scaler],
) {
    // NOTE: We should probably check the range of the values in rc_order and rc_i
    // because out of range values will cause a panic

    let mut bands = split_into_frequency_bands(duration, bandwidth, spec_lines);

    // quantized reflection coefficients will have no more than TNS_NUMFILTERS_MAX * MAXLAG entries
    let mut rc_quant: [Scaler; TNS_NUMFILTERS_MAX * MAXLAG] = [0.0; TNS_NUMFILTERS_MAX * MAXLAG];
    const QUANTIZER_STEPSIZE: Scaler = (PI / 17.0) as Scaler;

    // 0.489 ms - precompute 8 entries because using sin is very expensive in loops later on
    for (rc_q_item, rc_i_item) in rc_quant.iter_mut().zip(reflect_coef_ints) {
        if *rc_i_item != 0 {
            *rc_q_item = (QUANTIZER_STEPSIZE * (*rc_i_item as i32 - 8) as Scaler).sin()
        }
    }

    let mut lattice_state = [0.0; 8];

    // loop through each tns filter
    for (f, (band, order)) in bands
        .iter_mut()
        .zip(reflect_coef_order[..num_tns_filters].iter())
        .enumerate()
    {
        if *order > 0 {
            let offset = f * 8;

            // mutate each spectral line in the band associated with this filter
            for spectral_line in band.iter_mut() {
                let k = *order - 1;
                let mut t = *spectral_line - rc_quant[k + offset] * lattice_state[k];
                for k in (0..*order - 1).rev() {
                    let rc = rc_quant[k + offset];
                    t -= rc * lattice_state[k];
                    lattice_state[k + 1] = rc * t + lattice_state[k];
                }

                *spectral_line = t;
                lattice_state[0] = t;
            }
        }
    }
}

/// Split input spectral lines into slices according to frequency bands associated with the frame duration and bandwidth
///
/// # Arguments
///
/// * `duration` - Frame Duration in milliseconds
/// * `bandwidth` - Bandwidth cutoff index
/// * `spec_lines` - Spectral lines of the frame to split into bands (slices)
fn split_into_frequency_bands(
    duration: FrameDuration,
    bandwidth: Bandwidth,
    spec_lines: &mut [Scaler],
) -> Vec<&mut [Scaler], 2> {
    let mut bands: Vec<&mut [Scaler], 2> = Vec::new();
    match duration {
        FrameDuration::TenMs => match bandwidth {
            Bandwidth::NarrowBand => {
                bands.push(&mut spec_lines[12..80]).ok();
            }
            Bandwidth::WideBand => {
                bands.push(&mut spec_lines[12..160]).ok();
            }
            Bandwidth::SemiSuperWideBand => {
                bands.push(&mut spec_lines[12..240]).ok();
            }
            Bandwidth::SuperWideBand => {
                // 12..160 and 160..320
                let (left, right) = spec_lines.split_at_mut(160);
                bands.push(&mut left[12..]).ok();
                bands.push(&mut right[..160]).ok();
            }
            Bandwidth::FullBand => {
                // 12..200 and 200..400
                let (left, right) = spec_lines.split_at_mut(200);
                bands.push(&mut left[12..]).ok();
                bands.push(&mut right[..200]).ok();
            }
        },
        FrameDuration::SevenPointFiveMs => match bandwidth {
            Bandwidth::NarrowBand => {
                bands.push(&mut spec_lines[9..60]).ok();
            }
            Bandwidth::WideBand => {
                bands.push(&mut spec_lines[9..120]).ok();
            }
            Bandwidth::SemiSuperWideBand => {
                bands.push(&mut spec_lines[9..180]).ok();
            }
            Bandwidth::SuperWideBand => {
                // 9..120 and 120..240
                let (left, right) = spec_lines.split_at_mut(120);
                bands.push(&mut left[9..]).ok();
                bands.push(&mut right[..120]).ok();
            }
            Bandwidth::FullBand => {
                // 9..150 and 150..300
                let (left, right) = spec_lines.split_at_mut(150);
                bands.push(&mut left[9..]).ok();
                bands.push(&mut right[..150]).ok();
            }
        },
    };
    bands
}

#[cfg(test)]
mod tests {
    extern crate std;
    use super::*;
    use crate::common::config::FrameDuration;

    #[test]
    fn decode_test() {
        let duration = FrameDuration::TenMs;
        let bandwidth = Bandwidth::FullBand;
        let num_tns_filters = 2;
        let reflect_coef_order = [8, 0];
        let reflect_coef_ints = [6, 10, 7, 8, 7, 9, 7, 7];
        let mut spec_lines: [f32; 400] = [
            -568.56555, 1972.808, 3987.5906, 2461.2402, -263.29547, -2949.6724, -2217.0242, 141.18742, 1270.6868,
            -385.4035, 1667.538, 568.56555, -538.0386, -1606.4839, 538.0386, -263.29547, -49.60639, 751.7276,
            -476.98453, -141.18742, 1423.3219, -2033.862, -599.0926, 202.24144, -934.8897, -202.24144, 324.3495,
            1056.9977, -80.1334, 751.7276, 1667.538, -263.29547, 507.51154, 141.18742, 629.61957, -446.45752,
            507.51154, -385.4035, 202.24144, -1331.7408, 1453.8489, -385.4035, -49.60639, 141.18742, 0.0, 80.1334,
            -488.43216, 183.16206, -122.10804, 0.0, 0.0, 122.10804, -366.32413, -122.10804, -244.21608, -183.16206,
            0.0, 61.05402, 0.0, -427.37814, -244.21608, -671.59424, -427.37814, -366.32413, 0.0, 183.16206, 610.5402,
            122.10804, 549.4862, 183.16206, 427.37814, -122.10804, 122.10804, -122.10804, -61.05402, -61.05402, 0.0,
            0.0, 0.0, 0.0, 61.05402, 244.21608, -61.05402, -61.05402, -61.05402, 122.10804, 61.05402, 183.16206,
            183.16206, 61.05402, -122.10804, -122.10804, 61.05402, 305.2701, 122.10804, -122.10804, -305.2701,
            -244.21608, -122.10804, 0.0, 0.0, 0.0, -61.05402, 61.05402, -61.05402, 61.05402, 61.05402, 61.05402,
            -61.05402, -183.16206, -61.05402, 61.05402, 183.16206, 0.0, 0.0, 61.05402, 122.10804, -61.05402, -61.05402,
            61.05402, 122.10804, 122.10804, 61.05402, 122.10804, 0.0, -122.10804, -183.16206, 0.0, 183.16206,
            122.10804, 0.0, 0.0, 0.0, 0.0, 61.05402, 61.05402, 0.0, -61.05402, 0.0, 0.0, 0.0, -61.05402, -61.05402,
            61.05402, 122.10804, 0.0, 0.0, 0.0, 19.079382, 0.0, 0.0, 0.0, -61.05402, 122.10804, 0.0, -61.05402,
            -61.05402, 0.0, 122.10804, 0.0, 122.10804, 0.0, 0.0, -122.10804, -61.05402, 0.0, 61.05402, 0.0, -61.05402,
            0.0, 61.05402, 122.10804, 0.0, -61.05402, 0.0, 61.05402, 0.0, 61.05402, 0.0, 61.05402, -61.05402,
            -61.05402, 0.0, 0.0, 0.0, 0.0, -61.05402, -61.05402, 0.0, 0.0, 61.05402, 0.0, -61.05402, 0.0, 0.0, 0.0,
            0.0, 61.05402, 0.0, 61.05402, 0.0, 0.0, 0.0, 61.05402, 0.0, 61.05402, 0.0, 61.05402, 0.0, 0.0, 0.0,
            -61.05402, 61.05402, 122.10804, -61.05402, 0.0, 122.10804, 0.0, 61.05402, -61.05402, 0.0, 0.0, -61.05402,
            0.0, 0.0, 122.10804, -61.05402, -61.05402, 0.0, 0.0, 0.0, 0.0, -61.05402, 61.05402, -61.05402, 61.05402,
            61.05402, 61.05402, 0.0, 0.0, 61.05402, 0.0, 0.0, 0.0, 61.05402, 61.05402, 61.05402, 0.0, 0.0, 61.05402,
            0.0, 0.0, 61.05402, 122.10804, -122.10804, 0.0, -61.05402, -61.05402, 0.0, 0.0, 0.0, 0.0, -61.05402, 0.0,
            0.0, -61.05402, 61.05402, 61.05402, 61.05402, -122.10804, -61.05402, -61.05402, 0.0, 61.05402, -61.05402,
            0.0, 0.0, 61.05402, 0.0, 0.0, 0.0, -61.05402, 61.05402, -61.05402, 0.0, 0.0, -61.05402, 61.05402, 0.0, 0.0,
            0.0, 61.05402, 61.05402, 0.0, -61.05402, -61.05402, 0.0, -61.05402, 0.0, -61.05402, -61.05402, 0.0,
            61.05402, 0.0, 0.0, 0.0, 0.0, -61.05402, 0.0, 61.05402, 61.05402, 0.0, 0.0, 0.0, 0.0, -61.05402, -61.05402,
            -61.05402, 0.0, 0.0, 0.0, -61.05402, 0.0, 0.0, 61.05402, 0.0, -61.05402, 0.0, 61.05402, 0.0, 0.0, 0.0,
            -61.05402, 0.0, 0.0, 0.0, -61.05402, 0.0, 0.0, -61.05402, 0.0, 0.0, 0.0, 0.0, -61.05402, 61.05402, 0.0,
            0.0, 0.0, 0.0, 61.05402, 0.0, 0.0, 0.0, 0.0, -61.05402, -61.05402, 61.05402, 0.0, 0.0, 0.0, -19.079382,
            -19.079382, -19.079382, -19.079382, 19.079382, -19.079382, 19.079382, 0.0, 0.0, 0.0, -61.05402, -61.05402,
            0.0, 0.0, 61.05402, -61.05402, -61.05402, 0.0, 0.0, 61.05402, 0.0, 0.0, 0.0, 0.0, -61.05402, 0.0, 61.05402,
            0.0, 61.05402, 0.0, 0.0, 0.0, 0.0, 0.0, -61.05402, 0.0, 0.0, 61.05402, 61.05402,
        ];

        apply_temporal_noise_shaping(
            duration,
            bandwidth,
            num_tns_filters,
            &reflect_coef_order,
            &reflect_coef_ints,
            &mut spec_lines,
        );

        #[rustfmt::skip]
        let spec_lines_expected: [f32; 400] = [
            -568.56555, 1972.808, 3987.5906, 2461.2402, -263.29547, -2949.6724, -2217.0242, 141.18742, 1270.6868,
            -385.4035, 1667.538, 568.56555, -538.0386, -1924.9376, -330.6227, 358.01404, -103.25958, 640.72064,
            -380.6663, -565.7268, 1425.4789, -1620.2815, -2171.777, 65.41626, -630.4777, -241.61359, 70.17501,
            648.6986, 913.1724, 377.05695, 1248.8955, 447.55414, 122.224335, 387.0873, 670.9608, 202.11314, 470.52597,
            9.777176, 250.14848, -869.4124, 705.27844, 764.594, -196.6836, 162.10555, 32.383907, 212.68384, -200.68219,
            -385.05658, 74.749176, 244.4491, 75.39133, 71.7427, -370.58853, -284.1499, -221.86317, -378.81683,
            -78.224205, 120.7565, 49.39573, -431.70422, -639.2284, -863.50916, -762.39795, -557.4247, -210.2326,
            207.37712, 667.06, 268.28708, 266.76526, 180.27934, 266.1345, 18.510933, -71.35217, -18.678001, 61.719135,
            78.072945, 41.07525, 31.682747, 65.34431, 39.83117, 50.713173, 265.0006, 87.029144, -105.01486, -92.9501,
            109.07149, 213.51758, 237.69977, 244.04578, 169.18098, -60.653255, -202.20692, -5.709038, 418.32336,
            390.42206, -52.63867, -409.54736, -371.99884, -123.08209, 15.460239, -13.179712, 10.579021, 10.312677,
            49.179047, -104.04299, -108.42853, 45.232265, 106.2404, -18.412956, -252.0934, -189.36148, 85.68147,
            260.1948, 66.20917, -89.14935, 37.372047, 196.89316, 4.2800827, -186.45068, -10.664055, 251.7307,
            283.97842, 87.595795, 65.05641, 91.597, -64.72161, -260.67474, -120.33, 275.7829, 360.36444, 76.64067,
            -85.21214, -19.708641, 59.083878, 71.03325, 37.62242, 41.2656, 33.881153, 25.94694, -3.6607666, -15.336594,
            -46.333282, -68.63583, 51.637062, 181.11745, 75.242294, -25.19682, -26.834858, 23.40418, 38.82212,
            -11.055281, -14.305193, -22.998611, 129.22333, 87.56438, -91.969635, -113.919846, -21.006641, 165.25201,
            93.82325, 70.90962, 56.982315, 15.586489, -109.50122, -156.43475, -32.77449, 127.61873, 86.515625,
            -71.82623, -60.272907, 74.61493, 174.22997, 35.206467, -130.87024, -37.255043, 125.10912, 72.04756,
            28.948765, 1.5288572, 89.77242, 28.662163, -122.471275, -64.80139, 35.938797, 56.74036, 19.603413,
            -92.18097, -95.45599, 4.2908516, 4.6752024, 39.53695, 19.507534, -69.66742, -24.112461, -4.6465445,
            -19.89853, -2.2964888, 51.27441, 41.12713, 64.92945, 0.0, 0.0, 0.0, 61.05402, 0.0, 61.05402, 0.0, 61.05402,
            0.0, 0.0, 0.0, -61.05402, 61.05402, 122.10804, -61.05402, 0.0, 122.10804, 0.0, 61.05402, -61.05402, 0.0,
            0.0, -61.05402, 0.0, 0.0, 122.10804, -61.05402, -61.05402, 0.0, 0.0, 0.0, 0.0, -61.05402, 61.05402,
            -61.05402, 61.05402, 61.05402, 61.05402, 0.0, 0.0, 61.05402, 0.0, 0.0, 0.0, 61.05402, 61.05402, 61.05402,
            0.0, 0.0, 61.05402, 0.0, 0.0, 61.05402, 122.10804, -122.10804, 0.0, -61.05402, -61.05402, 0.0, 0.0, 0.0,
            0.0, -61.05402, 0.0, 0.0, -61.05402, 61.05402, 61.05402, 61.05402, -122.10804, -61.05402, -61.05402, 0.0,
            61.05402, -61.05402, 0.0, 0.0, 61.05402, 0.0, 0.0, 0.0, -61.05402, 61.05402, -61.05402, 0.0, 0.0,
            -61.05402, 61.05402, 0.0, 0.0, 0.0, 61.05402, 61.05402, 0.0, -61.05402, -61.05402, 0.0, -61.05402, 0.0,
            -61.05402, -61.05402, 0.0, 61.05402, 0.0, 0.0, 0.0, 0.0, -61.05402, 0.0, 61.05402, 61.05402, 0.0, 0.0, 0.0,
            0.0, -61.05402, -61.05402, -61.05402, 0.0, 0.0, 0.0, -61.05402, 0.0, 0.0, 61.05402, 0.0, -61.05402, 0.0,
            61.05402, 0.0, 0.0, 0.0, -61.05402, 0.0, 0.0, 0.0, -61.05402, 0.0, 0.0, -61.05402, 0.0, 0.0, 0.0, 0.0,
            -61.05402, 61.05402, 0.0, 0.0, 0.0, 0.0, 61.05402, 0.0, 0.0, 0.0, 0.0, -61.05402, -61.05402, 61.05402, 0.0,
            0.0, 0.0, -19.079382, -19.079382, -19.079382, -19.079382, 19.079382, -19.079382, 19.079382, 0.0, 0.0, 0.0,
            -61.05402, -61.05402, 0.0, 0.0, 61.05402, -61.05402, -61.05402, 0.0, 0.0, 61.05402, 0.0, 0.0, 0.0, 0.0,
            -61.05402, 0.0, 61.05402, 0.0, 61.05402, 0.0, 0.0, 0.0, 0.0, 0.0, -61.05402, 0.0, 0.0, 61.05402, 61.05402,
        ];

        assert_eq![spec_lines, spec_lines_expected];
    }
}