mcelp 1.0.1

//! The encoder's noise-analysis front end.
//!
//! This is the glue between the spectral analysis and the voice-activity
//! detector: the magnitude spectrum is brought back to unnormalised units,
//! folded into critical bands, and compared against the two long-term noise
//! references to produce the handful of features the detector votes on.
//!
//! [`Suppressor`] then chains that decision into the noise floors and the
//! per-band weights, and [`Frontend`] is the whole of it behind one call.

use crate::analysis::BINS;
use crate::bands::{self, BANDS};
use crate::fixed::{acc, hi, sat, shift};
use crate::vad::{self, Features, Frame};

/// Frames of silence the analysis sits out at start-up before it trusts its
/// own long-term estimates.
pub const WARMUP: i16 = 11;

/// Renormalise the magnitude spectrum out of the analysis headroom.
pub fn unnormalise(magnitude: &[i16; BINS], headroom: i16) -> [i16; BINS] {
    let mut out = [0i16; BINS];
    for (o, &m) in out.iter_mut().zip(magnitude.iter()) {
        *o = hi(sat(shift(acc((m as i64) << 16), -(headroom as i32))));
    }
    out
}

/// Build the detector's features from one frame's spectrum.
///
/// `fast` and `slow` are the two long-term band-energy estimates; the frame is
/// described by how far the current band levels sit above each of them, both
/// band by band and as a whole.
pub fn features(
    spectrum: &[i16; BINS],
    fast: &[i64; BANDS],
    slow: &[i64; BANDS],
    voicing: i16,
) -> Frame {
    let levels = bands::levels(spectrum);
    let summary = bands::summary(&levels);

    let mut frame = Frame {
        summary,
        voicing,
        ..Frame::default()
    };
    for (i, energies) in [fast, slow].into_iter().enumerate() {
        let tracked = bands::tracked_levels(energies);
        let excess = bands::excess(&levels, &tracked);
        frame.reference[i] = Features {
            tracked: bands::summary(&tracked),
            excess_sum: vad::excess_sum(&excess),
            difference: bands::summary_difference(summary, bands::summary(&tracked)),
            variance: bands::snr_statistics(&excess).1,
        };
    }
    frame
}

use crate::noise::NoiseEstimate;
use crate::shaping::{self, HOP, Levels, Overlap, SPAN, Scale, Smoother};
use crate::vad::{Decision, Vad};
use crate::weights::{self, Depths, Mix, Shaping};

/// Everything the noise-suppression front end carries between frames.
pub struct Suppressor {
    /// Frames still to sit out before the analysis trusts itself.
    warmup: i16,
    detector: Vad,
    fast: NoiseEstimate,
    slow: NoiseEstimate,
    levels: Levels,
    depths: Depths,
    shaping: Shaping,
    scale: Scale,
    smoother: Smoother,
    overlap: Overlap,
}

/// Per-band values needed to turn the current spectrum into a noise floor.
struct FloorShape {
    scale: [i16; BANDS],
    weight: [i16; BANDS],
    overall_level: i16,
    summary_level: i16,
}

/// Current per-band levels and their two frame-level summaries.
struct FloorLevels {
    band: [i16; BANDS],
    overall: i16,
    summary: i16,
}

/// Mutable per-band work arrays carried through floor refinement.
struct FloorBands {
    depth: [i64; BANDS],
    weight: [i16; BANDS],
    mix: [i64; BANDS],
}

impl FloorBands {
    fn new(levels: &FloorLevels, decision: Decision) -> Self {
        let (depth, weight) = weights::depths(levels.summary);
        let (mix, _) = Mix::run(
            &levels.band,
            levels.summary,
            levels.overall,
            decision.score,
            decision.active,
        );
        Self { depth, weight, mix }
    }
}

impl Default for Suppressor {
    fn default() -> Self {
        Suppressor {
            warmup: WARMUP,
            detector: Vad::default(),
            fast: NoiseEstimate::fast(),
            slow: NoiseEstimate::slow(),
            levels: Levels::default(),
            depths: Depths::default(),
            shaping: Shaping::default(),
            scale: Scale::default(),
            smoother: Smoother::default(),
            overlap: Overlap::default(),
        }
    }
}

impl Suppressor {
    /// Decide whether the current half-frame contains speech.
    fn decide(&mut self, magnitude: &[i16; BINS], headroom: i16, voicing: i16) -> Decision {
        if self.warmup >= 0 {
            self.warmup -= 1;
            Decision {
                active: false,
                fast: false,
                score: -4,
            }
        } else {
            let scaled = unnormalise(magnitude, headroom);
            let frame = features(&scaled, &self.fast.energy, &self.slow.energy, voicing);
            self.detector.run(&frame)
        }
    }

    /// Measure the current and tracked levels used to shape the floor.
    fn floor_levels(&mut self, magnitude: &[i16; BINS], headroom: i16, score: i16) -> FloorLevels {
        let overall = weights::overall_level(&self.fast.energy);
        let (band, summary) = self
            .levels
            .run(magnitude, &self.slow.energy, score, headroom);
        FloorLevels {
            band,
            overall,
            summary,
        }
    }

    /// Refine the per-band scale and weight used by the spectral floor.
    fn refine_floor_shape(&mut self, levels: &FloorLevels, decision: Decision) -> FloorShape {
        let mut bands = FloorBands::new(levels, decision);
        self.depths.refine(
            &mut bands.depth,
            &mut bands.mix,
            &self.fast.energy,
            &levels.band,
            &bands.weight,
            decision.active,
        );
        let (band_weight, mut band_scale) = self.shaping.finish(
            &mut bands.mix,
            &levels.band,
            levels.overall,
            levels.summary,
            decision.active,
        );
        self.scale.smooth(&mut band_scale, decision.active);
        FloorShape {
            scale: band_scale,
            weight: band_weight,
            overall_level: levels.overall,
            summary_level: levels.summary,
        }
    }

    /// Measure and refine the per-band floor shape.
    fn floor_shape(
        &mut self,
        magnitude: &[i16; BINS],
        headroom: i16,
        decision: Decision,
    ) -> FloorShape {
        let levels = self.floor_levels(magnitude, headroom, decision.score);
        self.refine_floor_shape(&levels, decision)
    }

    /// Update the suppression state and build this half-frame's spectral floor.
    fn noise_floor(
        &mut self,
        magnitude: &[i16; BINS],
        headroom: i16,
        decision: Decision,
    ) -> [i16; BINS] {
        self.fast.update(magnitude, headroom, decision.fast);
        self.slow.update(magnitude, headroom, decision.active);
        let shape = self.floor_shape(magnitude, headroom, decision);

        shaping::noise_floor(
            magnitude,
            &self.fast.energy,
            &shape.scale,
            &shape.weight,
            headroom,
            shape.overall_level,
            shape.summary_level,
        )
    }

    /// Apply the floor and turn the suppressed spectrum into overlap output.
    fn synthesise(
        &mut self,
        spectrum: &mut crate::analysis::Spectrum,
        magnitude: &[i16; BINS],
        floor: &[i16; BINS],
        headroom: i16,
    ) -> [i16; HOP] {
        shaping::suppress(spectrum, magnitude, floor, headroom);
        let mut span = shaping::inverse_transform(&spectrum.re, &spectrum.im);
        self.smoother.run(&mut span);
        self.overlap.emit(&span, headroom)
    }

    /// Run one half-frame through the whole path and return the shaping signal.
    ///
    /// `spectrum` is the complex transform of the analysis window and
    /// `magnitude` its magnitude; both come straight from the front end.
    pub fn run(
        &mut self,
        spectrum: &mut crate::analysis::Spectrum,
        magnitude: &[i16; BINS],
        headroom: i16,
        voicing: i16,
    ) -> [i16; HOP] {
        let decision = self.decide(magnitude, headroom, voicing);
        let floor = self.noise_floor(magnitude, headroom, decision);
        self.synthesise(spectrum, magnitude, &floor, headroom)
    }
}

/// Compile-time check that the shaping span is the size this module expects.
const _: () = assert!(SPAN == 256);

use crate::analysis::{self, Analysis, Direction, WINDOW};
use crate::preprocess::{self, Highpass};
use crate::{FRAME, HALF};

/// The encoder's analysis front end: mu-law samples in, shaping signal out.
#[derive(Default)]
pub struct Frontend {
    input: Highpass,
    analysis: Analysis,
    suppressor: Suppressor,
}

impl Frontend {
    /// Expand and condition one frame's worth of mu-law samples.
    ///
    /// The high-pass runs across the whole frame; everything after it works on
    /// half-frames.
    pub fn condition(&mut self, frame: &[u8; FRAME]) -> [i16; FRAME] {
        let mut linear = [0i16; FRAME];
        for (l, &b) in linear.iter_mut().zip(frame.iter()) {
            *l = preprocess::ulaw_to_linear(b);
        }
        self.input.run(&mut linear);
        linear
    }

    /// Run one conditioned half-frame through the analysis and the suppressor.
    pub fn process(&mut self, block: &[i16; HALF]) -> [i16; HOP] {
        let (mut window, headroom) = self.analysis.window(block);
        let voicing = analysis::voicing(&window);
        self.analysis.preemphasise(&mut window, headroom);

        let mut spectrum = analysis::deinterleave(&window);
        analysis::fft(&mut spectrum);
        analysis::unpack_real(&mut spectrum, Direction::Forward);
        let magnitude = analysis::magnitude(&spectrum);

        self.suppressor
            .run(&mut spectrum, &magnitude, headroom, voicing)
    }
}

/// Compile-time check that the analysis window and the shaping span agree.
const _: () = assert!(WINDOW == SPAN);