mcelp 1.0.1

//! Long-term (pitch) postfilter.
//!
//! Applied to the LPC residual before it is re-synthesised.  The reference
//! spreads it over five routines; the shape is:
//!
//! 1. rescale the residual history so the search has full precision;
//! 2. refine the transmitted pitch lag by +-1 sample, then over seven
//!    fractional phases in both directions, always maximising
//!    `correlation^2 / energy`;
//! 3. build the delayed signal a second time with a longer 16-tap fractional
//!    filter and keep whichever version correlates better;
//! 4. mix the delayed signal into the residual with a gain derived from the
//!    winning correlation.

use crate::SUBFRAME;
use crate::fixed::{acc, exp, hi, low, mul, sat, shift};

/// Residual history kept before the subframe being postfiltered.
pub const RES_HISTORY: usize = 152;
/// Samples of residual the search looks at: history plus current subframe.
pub const WINDOW: usize = RES_HISTORY + SUBFRAME;
/// Fractional phases tried on each side of the integer lag.
const PHASES: usize = 7;
/// Samples produced per interpolated phase (one more than a subframe so both
/// the leading and trailing window can be scored).
const INTERP_LEN: usize = SUBFRAME + 1;
/// Taps of the short interpolation filter used during the search.
const SEARCH_TAPS: usize = 4;
/// Taps of the long interpolation filter used once the phase is settled.
const REFINE_TAPS: usize = 16;

/// A normalised `correlation^2 / energy` score, kept as mantissa/exponent pairs
/// so candidates at different scales can be compared by cross multiplication.
#[derive(Clone, Copy, Default)]
struct Score {
    /// Numerator, `correlation^2`.
    numerator: i64,
    /// Denominator, the delayed signal's energy.
    denominator: i64,
    /// Correlation mantissa of the winner.
    correlation: i16,
    /// Fractional phase, 0 when the integer lag was best.
    phase: i16,
    /// Which of the two 80-sample windows the phase selects.
    direction: i16,
}

impl Score {
    /// Keep a fractional candidate when its correlation/energy ratio wins.
    fn consider(&mut self, candidate: Self, raw_energy: i16, peak_exp: i16) {
        if better_ratio(
            candidate.numerator,
            self.denominator,
            self.numerator,
            raw_energy,
            peak_exp,
        ) {
            *self = candidate;
        }
    }
}

/// Fractional-delay candidates and the two windows each phase exposes.
struct Interpolations {
    samples: [i16; PHASES * INTERP_LEN],
    energies: [[i16; PHASES]; 2],
    peak: i64,
}

impl Interpolations {
    fn new(integer_energy: i64) -> Self {
        Self {
            samples: [0; PHASES * INTERP_LEN],
            energies: [[0; PHASES]; 2],
            peak: integer_energy,
        }
    }

    /// Samples generated for one fractional phase.
    fn phase(&self, phase: usize) -> &[i16] {
        let start = phase * INTERP_LEN;
        &self.samples[start..start + INTERP_LEN]
    }

    /// Mutable samples generated for one fractional phase.
    fn phase_mut(&mut self, phase: usize) -> &mut [i16] {
        let start = phase * INTERP_LEN;
        &mut self.samples[start..start + INTERP_LEN]
    }

    /// One subframe window exposed by a phase.
    fn window(&self, phase: usize, direction: usize) -> &[i16] {
        &self.phase(phase)[direction..direction + SUBFRAME]
    }

    /// Generate and measure one fractional phase.
    fn build_phase(&mut self, res: &[i16; WINDOW], base: usize, phase: usize) {
        let (energies, edge) = {
            let block = self.phase_mut(phase);
            interpolate_phase(res, base, phase, block);
            interpolation_energies(block)
        };
        for (direction, energy) in energies.into_iter().enumerate() {
            self.energies[direction][phase] = energy;
        }
        self.peak = self.peak.max(edge);
    }
}

/// Score of the integer-delay candidate before fractional phases are tried.
fn integer_score(correlation: i64, energy: i64, corr_shift: i16, peak_exp: i16) -> Score {
    let correlation = hi(shift(correlation, corr_shift as i32));
    Score {
        numerator: acc((correlation as i64) * (correlation as i64) * 2),
        denominator: shift(energy, peak_exp as i32),
        correlation,
        phase: PHASES as i16,
        direction: 1,
    }
}

/// Score one of the two windows exposed by a fractional-delay phase.
fn fractional_score(
    current: &[i16],
    delayed: &[i16],
    energy: i16,
    phase: usize,
    direction: usize,
    corr_shift: i16,
    peak_exp: i16,
) -> Score {
    let scaled = shift(correlate(current, delayed), corr_shift as i32).max(0);
    let correlation = hi(scaled);
    Score {
        numerator: acc((correlation as i64) * (correlation as i64) * 2),
        denominator: shift((energy as i64) << 16, peak_exp as i32),
        correlation,
        phase: (PHASES - 1 - phase) as i16,
        direction: direction as i16,
    }
}

/// Outcome of the search.
pub struct Search {
    /// Integer lag the postfilter should use.
    pub lag: usize,
    /// Fractional phase, `0` meaning "no fractional filtering".
    pub phase: i16,
    /// Window selector that goes with [`Self::phase`].
    pub direction: i16,
    /// Correlation mantissa and exponent of the winner.
    pub correlation: i16,
    /// Energy mantissa of the winner.
    pub energy: i16,
    /// Exponents belonging to the two mantissas.
    pub corr_exp: i16,
    /// Exponent of the energy.
    pub energy_exp: i16,
    /// The interpolated signal that won; all zero when the winning phase was
    /// the integer one, which never reads it.
    pub interpolated: [i16; SUBFRAME],
}

/// Energy of `n` samples ending at `end` (exclusive) of `sig`.
fn energy(sig: &[i16]) -> i64 {
    let mut e = 0i64;
    for &v in sig {
        e = acc(e + (v as i64) * (v as i64) * 2);
    }
    e
}

/// Cross correlation of two equal-length blocks.
fn correlate(a: &[i16], b: &[i16]) -> i64 {
    let mut c = 0i64;
    for (&x, &y) in a.iter().zip(b) {
        c = acc(c + mul(x, y));
    }
    c
}

/// Refine the transmitted integer lag by one sample in either direction.
///
/// Returns the chosen lag, its correlation and its delayed-signal energy.
fn integer_candidate(res: &[i16; WINDOW], current: &[i16], lag0: i16) -> Option<(usize, i64, i64)> {
    let here = RES_HISTORY;
    let mut best_corr = -32768i64 << 16;
    let mut best_k = 0usize;
    for k in 0..3 {
        let delay = (lag0 as usize) - 1 + k;
        let correlation = correlate(current, &res[here - delay..here - delay + SUBFRAME]);
        if correlation >= best_corr {
            best_corr = correlation;
            best_k = k;
        }
        best_corr = sat(best_corr);
    }
    if best_corr <= 0 {
        return None;
    }

    let lag = (lag0 as usize) - 1 + best_k;
    let delayed = &res[here - lag..here - lag + SUBFRAME];
    let delayed_energy = energy(delayed);
    if acc(delayed_energy - (1i64 << 16)) < 0 {
        return None;
    }
    Some((lag, best_corr, sat(delayed_energy)))
}

/// Generate one short-filter fractional-delay phase.
fn interpolate_phase(res: &[i16; WINDOW], base: usize, phase: usize, out: &mut [i16]) {
    let coefficients =
        &crate::tables::SEARCH_FILTER[SEARCH_TAPS * phase..SEARCH_TAPS * phase + SEARCH_TAPS];
    for (n, sample) in out.iter_mut().enumerate() {
        let mut value = 0i64;
        for (tap, &coefficient) in coefficients.iter().enumerate() {
            value = acc(value + mul(coefficient, res[base + n - tap]));
        }
        *sample = hi(acc(value + 0x8000));
    }
}

/// Measure the two subframe windows exposed by one interpolated phase.
fn interpolation_energies(block: &[i16]) -> ([i16; 2], i64) {
    let energies = [
        hi(sat(energy(&block[..SUBFRAME]))),
        hi(sat(energy(&block[1..]))),
    ];
    // Only the stored 16-bit mantissa takes part in the running maximum.
    let edge = if (block[0] as i64).abs() > (block[SUBFRAME] as i64).abs() {
        energies[0]
    } else {
        energies[1]
    };
    (energies, (edge as i64) << 16)
}

/// Build every short-filter fractional phase and measure its two windows.
fn interpolation_candidates(
    res: &[i16; WINDOW],
    lag: usize,
    integer_energy: i64,
) -> Interpolations {
    let mut out = Interpolations::new(integer_energy);
    let base = RES_HISTORY - lag + 1;
    for phase in 0..PHASES {
        out.build_phase(res, base, phase);
    }
    out
}

/// Compare every fractional phase against the winning integer candidate.
fn consider_fractional_candidate(
    best: &mut Score,
    current: &[i16],
    candidates: &Interpolations,
    phase: usize,
    direction: usize,
    scale: (i16, i16),
) {
    let (corr_shift, peak_exp) = scale;
    let delayed = candidates.window(phase, direction);
    let denominator = candidates.energies[direction][phase];
    let candidate = fractional_score(
        current,
        delayed,
        denominator,
        phase,
        direction,
        corr_shift,
        peak_exp,
    );
    best.consider(candidate, denominator, peak_exp);
}

/// Compare every fractional phase against the winning integer candidate.
fn best_fractional_candidate(
    current: &[i16],
    integer_correlation: i64,
    integer_energy: i64,
    candidates: &Interpolations,
    corr_shift: i16,
    peak_exp: i16,
) -> Score {
    let mut best = integer_score(integer_correlation, integer_energy, corr_shift, peak_exp);

    for phase in 0..PHASES {
        for direction in 0..candidates.energies.len() {
            consider_fractional_candidate(
                &mut best,
                current,
                candidates,
                phase,
                direction,
                (corr_shift, peak_exp),
            );
        }
    }
    best
}

/// Convert the internal winning score into the public search result.
fn search_result(
    lag: usize,
    best: Score,
    candidates: &Interpolations,
    corr_shift: i16,
    peak_exp: i16,
) -> Search {
    let phase = (PHASES as i16) - best.phase;
    let mut interpolated = [0i16; SUBFRAME];
    if phase != 0 {
        let candidate = (phase - 1) as usize;
        interpolated.copy_from_slice(candidates.window(candidate, best.direction as usize));
    }

    Search {
        lag: lag + 1 - best.direction as usize,
        phase,
        direction: best.direction,
        correlation: best.correlation,
        energy: hi(best.denominator),
        corr_exp: corr_shift,
        energy_exp: peak_exp,
        interpolated,
    }
}

/// Normalisation parameters for the current residual block, when it has enough
/// energy for a meaningful pitch search.
fn current_energy_scale(current: &[i16]) -> Option<(i16, i16)> {
    let energy = energy(current);
    if acc(energy - (1i64 << 16)) < 0 {
        return None;
    }
    let exponent = exp(energy) as i16;
    Some((exponent, hi(shift(energy, exponent as i32))))
}

/// Final significance gate for the winning pitch candidate.
fn winner_is_significant(
    best: &Score,
    current_mantissa: i16,
    current_exponent: i16,
    corr_shift: i16,
    peak_exp: i16,
) -> bool {
    if best.correlation == 0 || acc(best.denominator - (1i64 << 16)) <= 0 {
        return false;
    }
    let scaled = shift(
        shift(
            acc(crate::fixed::mul32x16(best.denominator, current_mantissa)),
            -1,
        ),
        (2 * corr_shift - peak_exp - current_exponent) as i32,
    );
    acc(scaled - best.numerator) <= 0
}

/// Search for the best pitch delay to run the postfilter at.
///
/// `res` is the rescaled residual window; the subframe being filtered starts at
/// `RES_HISTORY`.  Returns `None` when the residual is too weak or no candidate
/// correlates positively, which switches the postfilter off.
pub fn search(res: &[i16; WINDOW], lag0: i16) -> Option<Search> {
    let here = RES_HISTORY;
    let current = &res[here..here + SUBFRAME];
    let (e0_exp, e0_mant) = current_energy_scale(current)?;

    let (lag, integer_correlation, integer_energy) = integer_candidate(res, current, lag0)?;
    let candidates = interpolation_candidates(res, lag, integer_energy);
    if acc(candidates.peak - (1i64 << 16)) < 0 {
        return None;
    }

    // Everything is now compared at a common scale.
    let peak_exp = exp(candidates.peak) as i16;
    let corr_shift = peak_exp.min(e0_exp);
    let best = best_fractional_candidate(
        current,
        integer_correlation,
        integer_energy,
        &candidates,
        corr_shift,
        peak_exp,
    );

    if !winner_is_significant(&best, e0_mant, e0_exp, corr_shift, peak_exp) {
        return None;
    }

    Some(search_result(lag, best, &candidates, corr_shift, peak_exp))
}

/// Is `new_num / new_den` a better ratio than `best_num / best_den`?
///
/// Compared by cross multiplication so no division is needed.  `best_den` is
/// held pre-scaled by `peak_exp` while the challenger's denominator is still a
/// raw mantissa, which is what the shift undoes.
fn better_ratio(new_num: i64, best_den: i64, best_num: i64, new_den: i16, peak_exp: i16) -> bool {
    let lhs = shift(mul32(best_den, new_num), -(peak_exp as i32));
    let rhs = crate::fixed::mul32x16(best_num, new_den);
    acc(lhs - rhs) > 0
}

/// 32x32 fractional product, spelled the way the reference builds it out of one
/// one unsigned multiply and three multiply-accumulates.
fn mul32(x: i64, y: i64) -> i64 {
    let xh = hi(x) as i64;
    let xl = low(x) as u16 as i64;
    let yh = hi(y) as i64;
    let yl = low(y) as u16 as i64;
    let mut a = acc(xl * yl);
    a = shift(a, -16);
    a = acc(a + yl * xh * 2);
    a = acc(a + xl * yh * 2);
    a = shift(a, -16);
    acc(a + xh * yh * 2)
}

/// Rebuild the delayed signal with the long filter and report its correlation
/// and energy.
pub fn refine(res: &[i16; WINDOW], lag: usize, phase: i16, out: &mut [i16; SUBFRAME]) -> Stats {
    let here = RES_HISTORY;
    let at = REFINE_TAPS * ((phase - 1) as usize);
    let h = &crate::tables::REFINE_FILTER[at..at + REFINE_TAPS];
    let base = here + 8 - lag;
    for n in 0..SUBFRAME {
        let mut a = 0i64;
        for (k, &c) in h.iter().enumerate() {
            a = acc(a + mul(c, res[base + n - k]));
        }
        out[n] = hi(sat(acc(a + 0x8000)));
    }

    let c = correlate(out, &res[here..here + SUBFRAME]);
    let mut stats = Stats::default();
    if c >= 0 {
        let e = exp(c) as i16;
        stats.corr_exp = e;
        stats.correlation = hi(shift(c, e as i32));
    }
    let en = energy(out);
    let e = exp(en) as i16;
    stats.energy_exp = e;
    stats.energy = hi(shift(en, e as i32));
    stats
}

/// A correlation/energy pair with the exponents needed to compare it.
#[derive(Clone, Copy, Default)]
pub struct Stats {
    /// Correlation mantissa.
    pub correlation: i16,
    /// Energy mantissa.
    pub energy: i16,
    /// Exponent of [`Self::correlation`].
    pub corr_exp: i16,
    /// Exponent of [`Self::energy`].
    pub energy_exp: i16,
}

/// Decide whether the refined delayed signal beats the one the search picked
///.  Returns true when the refinement wins.
pub fn refinement_wins(new: &Stats, old: &Stats) -> bool {
    if new.energy == 0 {
        return false;
    }
    let num_new = acc((new.correlation as i64) * (new.correlation as i64) * 2);
    let num_old = acc((old.correlation as i64) * (old.correlation as i64) * 2);
    let challenger = crate::fixed::mul32x16(num_new, old.energy);
    let incumbent = crate::fixed::mul32x16(num_old, new.energy);

    // Line the two products up before comparing them.
    let diff = 2 * (new.corr_exp as i64) - 2 * (old.corr_exp as i64) + (old.energy_exp as i64)
        - (new.energy_exp as i64);
    let (challenger, incumbent) = if diff > 0 {
        (shift(challenger, -(diff as i32)), incumbent)
    } else {
        (challenger, shift(incumbent, diff as i32))
    };
    acc(challenger - incumbent) > 0
}

/// Mix the delayed signal into the residual.
///
/// `gain` is the weight given to the original residual; the remainder goes to
/// the pitch-delayed copy.
pub fn mix(gain: i16, residual: &[i16], delayed: &[i16], out: &mut [i16; SUBFRAME]) {
    let complement = hi(acc((32768i64 << 16) - ((gain as i64) << 16)));
    for n in 0..SUBFRAME {
        let a = acc(mul(gain, residual[n]) + mul(complement, delayed[n]) + 0x8000);
        out[n] = hi(a);
    }
}

/// Weight given to the *original* residual when mixing.
///
/// Two thirds unless the delayed signal's energy dominates its correlation with
/// the residual, in which case the weight follows `E / (E + correlation)`.
pub fn mix_gain(stats: &Stats) -> i16 {
    /// Default weight, 2/3 in Q15.
    const DEFAULT: i16 = 21845;

    let shift_amount = (stats.corr_exp - stats.energy_exp) as i32;
    let scaled_energy = shift((stats.energy as i64) << 16, shift_amount);
    if acc(((stats.correlation as i64) << 16) - scaled_energy) > 0 {
        return DEFAULT;
    }
    let denominator = acc(scaled_energy + (stats.correlation as i64) * 16384 * 2);
    hi(shift(
        crate::fixed::divide(denominator, stats.energy),
        shift_amount,
    ))
}

/// Scale a residual window up so the lag search uses the full word.
fn normalise_residual(residual: &[i16; WINDOW]) -> ([i16; WINDOW], i32) {
    let mut peak = 0i64;
    for &value in residual.iter() {
        let magnitude = ((value as i64) << 16).abs();
        if magnitude > peak {
            peak = magnitude;
        }
    }
    let peak = sat(peak);
    let headroom = if peak == 0 { 12 } else { exp(peak) - 3 };

    let mut scaled = [0i16; WINDOW];
    for (output, &value) in scaled.iter_mut().zip(residual.iter()) {
        *output = hi(shift((value as i64) << 16, headroom));
    }
    (scaled, headroom)
}

/// Convert a search result into the statistics used for refinement and mixing.
fn search_stats(found: &Search) -> Stats {
    Stats {
        correlation: found.correlation,
        energy: found.energy,
        corr_exp: found.corr_exp,
        energy_exp: found.energy_exp,
    }
}

/// Rebuild the winning delayed signal, optionally with the longer filter.
fn rebuild_delayed(
    residual: &[i16; WINDOW],
    scaled: &[i16; WINDOW],
    found: &Search,
    headroom: i32,
    stats: &mut Stats,
) -> [i16; SUBFRAME] {
    let here = RES_HISTORY;
    if found.phase == 0 {
        let mut delayed = [0i16; SUBFRAME];
        delayed.copy_from_slice(&residual[here - found.lag..here - found.lag + SUBFRAME]);
        return delayed;
    }

    let mut delayed = found.interpolated;
    let mut refined = [0i16; SUBFRAME];
    let refined_stats = refine(scaled, found.lag, found.phase, &mut refined);
    if refinement_wins(&refined_stats, stats) {
        *stats = refined_stats;
        delayed = refined;
    }
    // Undo the search scaling before mixing.
    for value in delayed.iter_mut() {
        *value = hi(sat(shift((*value as i64) << 16, -headroom)));
    }
    delayed
}

/// Mix the winning delayed signal into the current residual subframe.
fn emit_filtered(
    residual: &[i16; WINDOW],
    delayed: &[i16; SUBFRAME],
    stats: &Stats,
    out: &mut [i16; SUBFRAME],
) {
    let gain = mix_gain(stats);
    mix(gain, &residual[RES_HISTORY..], delayed, out);
}

/// Run the long-term postfilter over one subframe.
///
/// `residual` is the unscaled residual window; the current subframe occupies
/// its last [`SUBFRAME`] samples.  Returns the pitch lag the filter settled on,
/// which the excitation builder of the next half-frame reads as a voicing hint.
pub fn filter(residual: &[i16; WINDOW], lag0: i16, out: &mut [i16; SUBFRAME]) -> i16 {
    let (scaled, headroom) = normalise_residual(residual);
    let found = match search(&scaled, lag0) {
        None => {
            out.copy_from_slice(&residual[RES_HISTORY..]);
            return 0;
        }
        Some(found) => found,
    };

    let mut stats = search_stats(&found);
    let delayed = rebuild_delayed(residual, &scaled, &found, headroom, &mut stats);
    emit_filtered(residual, &delayed, &stats, out);
    found.lag as i16
}