helena 0.1.0

//! Dependency-free time-domain quality metrics for decoded waveforms.
//!
//! These operate on [`Pcm`], whose invariants (a nonzero [`SampleRate`], whole
//! interleaved frames, and *finite samples*) already hold from construction.
//! So the non-finite-sample rejections the previous `Waveform`-based metrics
//! carried are gone here — a NaN or infinite sample is unrepresentable in a
//! [`Pcm`], not something these folds re-scan for — and the old zero-`sample_rate`
//! special cases die with the sentinel. Only the metrics' own preconditions
//! (a non-empty clip, matching shapes for a pair) remain to check.

use std::num::NonZeroUsize;

use crate::{Error, Pcm, Result};

/// Full-scale magnitude for `f32` PCM.
const FULL_SCALE: f32 = 1.0;
const ENERGY_EPS: f64 = 1e-20;
const SI_SNR_DB_CAP: f64 = 120.0;

/// Parameter-free time-domain quality statistics of a decoded [`Pcm`] clip
/// (PRD §16.2 audio-quality metrics).
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct WaveformQuality {
    clipping_rate: f32,
    peak: f32,
    rms: f32,
}

impl WaveformQuality {
    /// Measure `pcm`. Requires a non-empty clip; finiteness is a [`Pcm`]
    /// invariant, so it is not re-checked here.
    pub fn measure(pcm: &Pcm) -> Result<Self> {
        let samples = non_empty_samples(pcm)?;
        let n = samples.len() as f64;
        let mut clipped = 0usize;
        let mut peak = 0.0f32;
        let mut sum_sq = 0.0f64;
        for &s in samples {
            let mag = s.abs();
            if mag >= FULL_SCALE {
                clipped += 1;
            }
            peak = peak.max(mag);
            sum_sq += (s as f64) * (s as f64);
        }
        Ok(Self {
            clipping_rate: (clipped as f64 / n) as f32,
            peak,
            rms: (sum_sq / n).sqrt() as f32,
        })
    }

    /// Fraction of samples at or beyond full scale (`|s| >= 1.0`).
    pub fn clipping_rate(&self) -> f32 {
        self.clipping_rate
    }

    /// Largest sample magnitude.
    pub fn peak(&self) -> f32 {
        self.peak
    }

    /// Root-mean-square amplitude over all samples.
    pub fn rms(&self) -> f32 {
        self.rms
    }
}

/// Noise floor: the lowest per-window RMS amplitude over consecutive,
/// non-overlapping `window_frames`-frame windows (PRD §16.2).
pub fn noise_floor(pcm: &Pcm, window_frames: NonZeroUsize) -> Result<f32> {
    let samples = non_empty_samples(pcm)?;
    let channels = pcm.channels();
    let window = window_frames.get();
    let frames = samples.len() / channels;
    if frames < window {
        return Err(Error::validation(format!(
            "noise-floor window of {window} frames exceeds the clip's {frames} frames"
        )));
    }
    let stride = window * channels;
    let mut floor = f64::INFINITY;
    for w in 0..frames / window {
        let block = &samples[w * stride..(w + 1) * stride];
        let sum_sq: f64 = block.iter().map(|&s| (s as f64) * (s as f64)).sum();
        floor = floor.min((sum_sq / stride as f64).sqrt());
    }
    Ok(floor as f32)
}

/// Time-domain signal preservation between a reference and reconstruction.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct SignalPreservation {
    zero_lag_correlation: f64,
    max_abs_correlation: f64,
    max_abs_correlation_lag_frames: i64,
    scale_invariant_snr_db: Option<f64>,
}

impl SignalPreservation {
    /// Measure same-shaped clips with a default ±5ms correlation search window.
    pub fn measure(reference: &Pcm, candidate: &Pcm) -> Result<Self> {
        let frames = validate_pair(reference, candidate)?;
        let max_lag_frames = (reference.rate().get() as usize / 200)
            .max(1)
            .min(frames.saturating_sub(1));
        Self::measure_with_max_lag_frames(reference, candidate, max_lag_frames)
    }

    /// Pearson correlation at zero lag.
    pub fn zero_lag_correlation(self) -> f64 {
        self.zero_lag_correlation
    }

    /// Largest absolute Pearson correlation over the bounded lag search.
    pub fn max_abs_correlation(self) -> f64 {
        self.max_abs_correlation
    }

    /// Lag in frames at which [`Self::max_abs_correlation`] was observed.
    pub fn max_abs_correlation_lag_frames(self) -> i64 {
        self.max_abs_correlation_lag_frames
    }

    /// Scale-invariant SNR in dB, capped to ±120 dB.
    pub fn scale_invariant_snr_db(self) -> Option<f64> {
        self.scale_invariant_snr_db
    }

    fn measure_with_max_lag_frames(
        reference: &Pcm,
        candidate: &Pcm,
        max_lag_frames: usize,
    ) -> Result<Self> {
        let frames = validate_pair(reference, candidate)?;
        let channels = reference.channels();
        let max_lag_frames = max_lag_frames.min(frames.saturating_sub(1));
        let zero_lag_correlation = correlation_at_lag(reference, candidate, channels, frames, 0);
        let mut best_abs = zero_lag_correlation.abs();
        let mut best_lag = 0i64;
        for lag in -(max_lag_frames as i64)..=(max_lag_frames as i64) {
            let corr = correlation_at_lag(reference, candidate, channels, frames, lag);
            let abs = corr.abs();
            if abs > best_abs || (abs == best_abs && lag.abs() < best_lag.abs()) {
                best_abs = abs;
                best_lag = lag;
            }
        }
        Ok(Self {
            zero_lag_correlation,
            max_abs_correlation: best_abs,
            max_abs_correlation_lag_frames: best_lag,
            scale_invariant_snr_db: scale_invariant_snr_db(reference, candidate),
        })
    }
}

/// A [`Pcm`] is finite and whole-framed by construction; only the metric's own
/// non-empty precondition remains (RMS and correlation over zero samples are
/// undefined). A silent clip of zero frames is representable, so this stays a
/// runtime check rather than a type invariant.
fn non_empty_samples(pcm: &Pcm) -> Result<&[f32]> {
    if pcm.samples().is_empty() {
        return Err(Error::validation("waveform has no samples"));
    }
    Ok(pcm.samples())
}

fn validate_pair(reference: &Pcm, candidate: &Pcm) -> Result<usize> {
    let reference_samples = non_empty_samples(reference)?;
    let candidate_samples = non_empty_samples(candidate)?;
    if reference.rate() != candidate.rate()
        || reference.channels() != candidate.channels()
        || reference.frames() != candidate.frames()
    {
        return Err(Error::validation(
            "signal-preservation waveforms must share sample rate, channels, and frame count",
        ));
    }
    if reference_samples.len() != candidate_samples.len() {
        return Err(Error::validation(
            "signal-preservation waveforms must have equal sample counts",
        ));
    }
    Ok(reference.frames())
}

fn correlation_at_lag(
    reference: &Pcm,
    candidate: &Pcm,
    channels: usize,
    frames: usize,
    lag_frames: i64,
) -> f64 {
    let lag_abs = lag_frames.unsigned_abs() as usize;
    if lag_abs >= frames {
        return 0.0;
    }
    let overlap_frames = frames - lag_abs;
    let (reference_start, candidate_start) = if lag_frames >= 0 {
        (lag_abs, 0)
    } else {
        (0, lag_abs)
    };
    let samples = overlap_frames * channels;
    let mut reference_sum = 0.0f64;
    let mut candidate_sum = 0.0f64;
    for offset in 0..overlap_frames {
        let reference_base = (reference_start + offset) * channels;
        let candidate_base = (candidate_start + offset) * channels;
        for channel in 0..channels {
            reference_sum += f64::from(reference.samples()[reference_base + channel]);
            candidate_sum += f64::from(candidate.samples()[candidate_base + channel]);
        }
    }

    let reference_mean = reference_sum / samples as f64;
    let candidate_mean = candidate_sum / samples as f64;
    let mut dot = 0.0f64;
    let mut reference_energy = 0.0f64;
    let mut candidate_energy = 0.0f64;
    for offset in 0..overlap_frames {
        let reference_base = (reference_start + offset) * channels;
        let candidate_base = (candidate_start + offset) * channels;
        for channel in 0..channels {
            let reference =
                f64::from(reference.samples()[reference_base + channel]) - reference_mean;
            let candidate =
                f64::from(candidate.samples()[candidate_base + channel]) - candidate_mean;
            dot += reference * candidate;
            reference_energy += reference * reference;
            candidate_energy += candidate * candidate;
        }
    }
    normalized_dot(dot, reference_energy, candidate_energy)
}

fn scale_invariant_snr_db(reference: &Pcm, candidate: &Pcm) -> Option<f64> {
    let (dot, reference_energy, candidate_energy) =
        centered_pair_stats(reference.samples(), candidate.samples());
    if reference_energy <= ENERGY_EPS {
        return None;
    }
    if candidate_energy <= ENERGY_EPS {
        return Some(-SI_SNR_DB_CAP);
    }

    let projection_power = (dot * dot) / reference_energy;
    let residual_power = (candidate_energy - projection_power).max(0.0);
    if projection_power <= ENERGY_EPS {
        return Some(-SI_SNR_DB_CAP);
    }
    if residual_power <= ENERGY_EPS {
        return Some(SI_SNR_DB_CAP);
    }
    let db = 10.0 * (projection_power / residual_power).log10();
    Some(db.clamp(-SI_SNR_DB_CAP, SI_SNR_DB_CAP))
}

fn centered_pair_stats(reference: &[f32], candidate: &[f32]) -> (f64, f64, f64) {
    let n = reference.len() as f64;
    let reference_mean = reference.iter().map(|&s| f64::from(s)).sum::<f64>() / n;
    let candidate_mean = candidate.iter().map(|&s| f64::from(s)).sum::<f64>() / n;
    let mut dot = 0.0f64;
    let mut reference_energy = 0.0f64;
    let mut candidate_energy = 0.0f64;
    for (&reference, &candidate) in reference.iter().zip(candidate) {
        let reference = f64::from(reference) - reference_mean;
        let candidate = f64::from(candidate) - candidate_mean;
        dot += reference * candidate;
        reference_energy += reference * reference;
        candidate_energy += candidate * candidate;
    }
    (dot, reference_energy, candidate_energy)
}

fn normalized_dot(dot: f64, reference_energy: f64, candidate_energy: f64) -> f64 {
    if reference_energy <= ENERGY_EPS || candidate_energy <= ENERGY_EPS {
        0.0
    } else {
        (dot / (reference_energy * candidate_energy).sqrt()).clamp(-1.0, 1.0)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{ChannelLayout, SampleRate};
    use proptest::prelude::*;

    fn rate(hz: u32) -> SampleRate {
        SampleRate::try_from(hz).unwrap()
    }

    /// Build a mono clip at `hz`, panicking on the (test-controlled) invariants.
    fn mono(hz: u32, samples: Vec<f32>) -> Pcm {
        Pcm::new(rate(hz), ChannelLayout::Mono, samples).unwrap()
    }

    /// Build a stereo clip at `hz`.
    fn stereo(hz: u32, samples: Vec<f32>) -> Pcm {
        Pcm::new(rate(hz), ChannelLayout::Stereo, samples).unwrap()
    }

    fn window(n: usize) -> NonZeroUsize {
        NonZeroUsize::new(n).expect("test window is non-zero")
    }

    #[test]
    fn quality_of_clean_clip() {
        let w = mono(16_000, vec![0.5, -0.5, 0.5, -0.5]);
        let q = WaveformQuality::measure(&w).unwrap();
        assert_eq!(q.clipping_rate(), 0.0);
        assert_eq!(q.peak(), 0.5);
        assert!((q.rms() - 0.5).abs() < 1e-6);
    }

    #[test]
    fn clipping_rate_counts_full_scale_samples() {
        let w = mono(16_000, vec![1.0, 0.2, -1.5, -0.1]);
        let q = WaveformQuality::measure(&w).unwrap();
        assert_eq!(q.clipping_rate(), 0.5);
        assert_eq!(q.peak(), 1.5);
    }

    #[test]
    fn quality_rejects_empty_clip() {
        // Zero channels, partial frames, and non-finite samples are now
        // unrepresentable: `Pcm::new` rejects them at construction, so those
        // former measurement rejections have no `Pcm` to test. An empty clip is
        // still representable (a zero-frame silence), and quality is undefined on
        // it, so that precondition remains a measurement-time check.
        assert!(WaveformQuality::measure(&mono(16_000, vec![])).is_err());
    }

    #[test]
    fn noise_floor_is_quietest_window() {
        let w = mono(16_000, vec![1.0, 0.0, 0.0, 0.0]);
        assert_eq!(noise_floor(&w, window(1)).unwrap(), 0.0);
        assert!((noise_floor(&w, window(4)).unwrap() - 0.5).abs() < 1e-6);
    }

    #[test]
    fn noise_floor_respects_frame_interleaving() {
        let w = stereo(16_000, vec![1.0, -1.0, 0.0, 0.0]);
        assert_eq!(noise_floor(&w, window(1)).unwrap(), 0.0);
    }

    #[test]
    fn noise_floor_rejects_oversized_window() {
        let w = mono(16_000, vec![0.1, 0.2]);
        assert!(noise_floor(&w, window(3)).is_err());
    }

    #[test]
    fn signal_preservation_scores_identical_clip() {
        let samples = pseudo_signal(512);
        let reference = mono(48_000, samples.clone());
        let candidate = mono(48_000, samples);

        let signal = SignalPreservation::measure(&reference, &candidate).unwrap();

        assert!(signal.zero_lag_correlation() > 0.999_999);
        assert!(signal.max_abs_correlation() > 0.999_999);
        assert_eq!(signal.max_abs_correlation_lag_frames(), 0);
        assert_eq!(signal.scale_invariant_snr_db(), Some(SI_SNR_DB_CAP));
    }

    #[test]
    fn signal_preservation_flags_equal_energy_orthogonal_noise() {
        let reference = mono(48_000, vec![1.0, -1.0, 1.0, -1.0]);
        let candidate = mono(48_000, vec![1.0, 1.0, -1.0, -1.0]);

        let signal =
            SignalPreservation::measure_with_max_lag_frames(&reference, &candidate, 0).unwrap();

        assert!(signal.zero_lag_correlation().abs() < 1e-12);
        assert!(signal.max_abs_correlation() < 1e-12);
        assert_eq!(signal.scale_invariant_snr_db(), Some(-SI_SNR_DB_CAP));
    }

    #[test]
    fn signal_preservation_finds_small_lagged_matches() {
        let samples = pseudo_signal(1024);
        let mut delayed = vec![0.0; 3];
        delayed.extend(samples.iter().take(samples.len() - 3));
        let reference = mono(48_000, samples);
        let candidate = mono(48_000, delayed);

        let signal =
            SignalPreservation::measure_with_max_lag_frames(&reference, &candidate, 8).unwrap();

        assert!(signal.max_abs_correlation() > 0.999_999);
        assert_eq!(signal.max_abs_correlation_lag_frames(), -3);
    }

    #[test]
    fn signal_preservation_rejects_mismatched_pairs() {
        let reference = mono(48_000, vec![0.0, 1.0]);
        let candidate = mono(44_100, vec![0.0, 1.0]);

        assert!(SignalPreservation::measure(&reference, &candidate).is_err());
    }

    fn pseudo_signal(len: usize) -> Vec<f32> {
        let mut state = 0xC0FFEEu32;
        (0..len)
            .map(|_| {
                state = state.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
                (((state >> 8) as f32 / 16_777_215.0) * 2.0 - 1.0) * 0.25
            })
            .collect()
    }

    proptest! {
        #[test]
        fn quality_invariants(samples in proptest::collection::vec(-2.0f32..2.0, 1..64)) {
            let w = mono(16_000, samples);
            let q = WaveformQuality::measure(&w).unwrap();
            prop_assert!(q.rms() <= q.peak() + 1e-6);
            prop_assert!((0.0..=1.0).contains(&q.clipping_rate()));
        }

        #[test]
        fn noise_floor_at_most_rms(samples in proptest::collection::vec(-2.0f32..2.0, 1..64)) {
            let w = mono(16_000, samples);
            let rms = WaveformQuality::measure(&w).unwrap().rms();
            let floor = noise_floor(&w, window(1)).unwrap();
            prop_assert!(floor <= rms + 1e-6, "floor {floor} vs rms {rms}");
        }

        #[test]
        fn signal_correlations_stay_bounded(
            reference in proptest::collection::vec(-2.0f32..2.0, 1..64),
            candidate in proptest::collection::vec(-2.0f32..2.0, 1..64),
        ) {
            let len = reference.len().min(candidate.len());
            let reference = mono(16_000, reference[..len].to_vec());
            let candidate = mono(16_000, candidate[..len].to_vec());
            let signal = SignalPreservation::measure_with_max_lag_frames(
                &reference,
                &candidate,
                0,
            )
            .unwrap();
            prop_assert!((-1.0..=1.0).contains(&signal.zero_lag_correlation()));
            prop_assert!((0.0..=1.0).contains(&signal.max_abs_correlation()));
            if let Some(si_snr) = signal.scale_invariant_snr_db() {
                prop_assert!((-SI_SNR_DB_CAP..=SI_SNR_DB_CAP).contains(&si_snr));
            }
        }
    }
}