embedded-nn 0.2.1

//! Bridge from `embedded-dsp`'s complex FFT output into this crate's
//! quantized tensors. Available with the `embedded-dsp` feature.
//!
//! The common on-MCU audio/sensor classification pipeline is: sample ->
//! FFT (`embedded_dsp::transform`) -> per-bin power spectrum -> quantize
//! into an `s8`/`s16` feature vector -> feed into this crate's inference
//! ops (`fully_connected`, `convolution`, ...). This module covers the
//! last step: turning `embedded_dsp::Complex<f32>` FFT bins into the
//! quantized tensors the rest of this crate expects.

use embedded_dsp::types::Complex;

/// Per-bin power (squared magnitude) of a complex spectrum bin: `re^2 + im^2`.
///
/// Squared magnitude (rather than magnitude) avoids a `sqrt`, so this
/// bridge stays free of any float-math-backend requirement beyond basic
/// arithmetic.
#[inline]
fn power(c: Complex<f32>) -> f32 {
    c.real * c.real + c.imag * c.imag
}

/// Round-half-away-from-zero and shift by `zero_point`, without relying on
/// `f32::round` (not available in `core` on all targets without `std`/`libm`).
#[inline]
fn quantize_affine(value: f32, scale: f32, zero_point: i32) -> i32 {
    let scaled = value / scale;
    let rounded = scaled + if scaled >= 0.0 { 0.5 } else { -0.5 };
    (rounded as i32).wrapping_add(zero_point)
}

/// Quantize a complex spectrum's per-bin power into a signed 8-bit tensor.
///
/// Uses the affine scheme `q = round(power / scale) + zero_point`,
/// saturated to `i8`'s range — the same convention
/// [`crate::types::PerTensorQuantParams`] assumes elsewhere in this crate,
/// but applied to the raw `f32` power values a caller gets straight out of
/// an FFT rather than an already-quantized tensor.
///
/// Writes `spectrum.len().min(out.len())` bins and returns that count.
pub fn quantize_power_spectrum_s8(
    spectrum: &[Complex<f32>],
    scale: f32,
    zero_point: i32,
    out: &mut [i8],
) -> usize {
    let len = spectrum.len().min(out.len());
    for i in 0..len {
        let q = quantize_affine(power(spectrum[i]), scale, zero_point);
        out[i] = q.clamp(i8::MIN as i32, i8::MAX as i32) as i8;
    }
    len
}

/// Quantize a complex spectrum's per-bin power into a signed 16-bit tensor.
///
/// See [`quantize_power_spectrum_s8`] for the quantization scheme.
pub fn quantize_power_spectrum_s16(
    spectrum: &[Complex<f32>],
    scale: f32,
    zero_point: i32,
    out: &mut [i16],
) -> usize {
    let len = spectrum.len().min(out.len());
    for i in 0..len {
        let q = quantize_affine(power(spectrum[i]), scale, zero_point);
        out[i] = q.clamp(i16::MIN as i32, i16::MAX as i32) as i16;
    }
    len
}

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

    #[test]
    fn quantizes_known_values_s8() {
        let spectrum = [Complex::new(3.0f32, 4.0), Complex::new(0.0, 0.0)];
        // power = 3^2 + 4^2 = 25.0, and 0.0
        let mut out = [0i8; 2];
        let n = quantize_power_spectrum_s8(&spectrum, 1.0, -10, &mut out);
        assert_eq!(n, 2);
        assert_eq!(out[0], 15); // 25/1.0 - 10 = 15
        assert_eq!(out[1], -10);
    }

    #[test]
    fn saturates_s8_upper_bound() {
        let spectrum = [Complex::new(1000.0f32, 0.0)];
        let mut out = [0i8; 1];
        quantize_power_spectrum_s8(&spectrum, 1.0, 0, &mut out);
        assert_eq!(out[0], i8::MAX);
    }

    #[test]
    fn saturates_s8_lower_bound() {
        let spectrum = [Complex::new(0.0f32, 0.0)];
        let mut out = [0i8; 1];
        quantize_power_spectrum_s8(&spectrum, 1.0, -1000, &mut out);
        assert_eq!(out[0], i8::MIN);
    }

    #[test]
    fn quantizes_known_values_s16() {
        let spectrum = [Complex::new(3.0f32, 4.0)];
        let mut out = [0i16; 1];
        quantize_power_spectrum_s16(&spectrum, 0.5, 0, &mut out);
        assert_eq!(out[0], 50); // 25 / 0.5 = 50
    }

    #[test]
    fn truncates_to_shorter_output_buffer() {
        let spectrum = [
            Complex::new(1.0f32, 0.0),
            Complex::new(2.0, 0.0),
            Complex::new(3.0, 0.0),
        ];
        let mut out = [0i8; 2];
        let n = quantize_power_spectrum_s8(&spectrum, 1.0, 0, &mut out);
        assert_eq!(n, 2);
    }
}