lasprs 0.14.2

//! Spectrogram computation: per-bin auto power of a single channel, as a
//! function of frequency and time. No time averaging is applied; each output
//! column corresponds to a single windowed FFT block (see
//! [ApsMode::Spectrogram]).

use super::*;
use crate::slm::Lref_default;
use crate::*;
use anyhow::{Result, bail};

/// Small value to avoid `log(0)` in the dB conversion.
const SMALL_NUMBER: Flt = 1e-80;

/// Result of a spectrogram computation: per-bin auto power in dB for a single
/// channel, as a function of frequency and time.
#[cfg_attr(feature = "python-bindings", gen_stub_pyclass, pyclass(from_py_object))]
#[derive(Debug, Clone)]
pub struct SpectrogramResult {
    /// Auto power per bin in dB, shape (nfreq, ntime).
    pub ap_dB: Array2<Flt>,
    /// Frequency axis [Hz], length nfreq.
    pub freq: Array1<Flt>,
    /// Time axis [s], length ntime — center time of each FFT block.
    pub t: Array1<Flt>,
}

/// Spectrogram engine. Computes the per-block auto power spectrum (in dB) of a
/// single channel. Push time data in arbitrary chunks; one output column is
/// produced per complete FFT block (with the configured overlap).
///
/// The engine expects blocks with exactly one column (one channel): the auto
/// power is always taken at channel index 0 of the computed CPS result.
pub struct SpectroGramEngine {
    /// The underlying power spectra engine, in spectrogram mode.
    aps: AvPowerSpectra,
    /// Square of the reference value for the dB conversion.
    ref_value_sq: Flt,
    /// Collected dB columns, in order.
    ap_dB: Vec<Array1<Flt>>,
    /// Center time of each collected block, relative to the start of the data
    /// fed to [SpectroGramEngine::push].
    t: Vec<Flt>,
    /// Total samples pushed across all push() calls (consumed blocks).
    samples_consumed: usize,
}

impl SpectroGramEngine {
    /// Create a new spectrogram engine.
    ///
    /// # Args
    ///
    /// * `nfft` - FFT length [samples]
    /// * `overlap` - Overlap between consecutive FFT blocks
    /// * `window` - Window applied to each block
    /// * `fs` - Sampling frequency [Hz]
    /// * `freqWeightingType` - Frequency weighting applied to the power
    /// * `ref_value` - Reference value (linear) for the dB conversion.
    ///   Defaults to [Lref_default] (20 µPa) when `None`.
    pub fn new(
        nfft: usize,
        overlap: Overlap,
        window: WindowType,
        fs: StrictlyPositive,
        freqWeightingType: FreqWeighting,
        ref_value: Option<StrictlyPositive>,
    ) -> Result<Self> {
        let settings = ApsSettingsBuilder::default()
            .mode(ApsMode::Spectrogram {})
            .overlap(overlap)
            .windowType(window)
            .freqWeightingType(freqWeightingType)
            .nfft(nfft)
            .fs(fs)
            .build()?;
        let ref_value = match ref_value {
            Some(v) => v,
            None => StrictlyPositive::new(Lref_default).expect("Lref_default must be positive"),
        };
        let ref_value_sq = *ref_value * *ref_value;
        Ok(Self {
            aps: AvPowerSpectra::new(settings),
            ref_value_sq,
            ap_dB: Vec::new(),
            t: Vec::new(),
            samples_consumed: 0,
        })
    }

    /// Push a block of time data, of shape (samples, 1). Appends one output
    /// column per complete FFT block that can be formed from the accumulated
    /// data.
    pub fn push(&mut self, block: ArrayView2<Flt>) -> Result<()> {
        let results = self.aps.compute_all(block);
        let settings = self.aps.settings();
        let hop_size = settings.get_hop_size();
        let nfft = settings.nfft;
        let fs = *settings.fs;
        for res in results {
            let ap = res.ap(0);
            let ap_dB =
                ap.mapv(|p| 10. * Flt::log10(Flt::max(p, SMALL_NUMBER) / self.ref_value_sq));
            self.ap_dB.push(ap_dB);
            // Center of block i: start = samples_consumed, center = start + nfft/2
            let center = (self.samples_consumed as Flt + nfft as Flt / 2.0) / fs;
            self.t.push(center);
            self.samples_consumed += hop_size;
        }
        Ok(())
    }

    /// Finish the computation: stack the collected columns into a 2D array and
    /// add the frequency axis.
    ///
    /// # Errors
    ///
    /// When no complete FFT block was available.
    pub fn finish(self) -> Result<SpectrogramResult> {
        if self.ap_dB.is_empty() {
            bail!("Not enough data for spectrogram: no complete FFT block available");
        }
        let ntime = self.ap_dB.len();
        let nfreq = self.ap_dB[0].len();
        let mut ap_dB = Array2::zeros((nfreq, ntime));
        for (i, col) in self.ap_dB.into_iter().enumerate() {
            ap_dB.column_mut(i).assign(&col);
        }
        let settings = self.aps.settings();
        let freq = getFreq(settings.fs, settings.nfft);
        let t = Array1::from(self.t);
        Ok(SpectrogramResult { ap_dB, freq, t })
    }
}

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

    /// Push sine data in chunks and verify spectrogram block counts.
    #[test]
    fn test_block_counts() -> anyhow::Result<()> {
        let fs: StrictlyPositive = 48000.0.try_into().unwrap();
        let freq = 1000.0;
        let duration = 2.0;
        let n_samples = (*fs * duration) as usize;
        let nfft = 4096usize;

        let signal: Vec<Flt> = (0..n_samples)
            .map(|i| {
                let t = i as Flt / *fs;
                (2.0 * std::f64::consts::PI * freq * t).sin() as Flt
            })
            .collect();

        let expected = |hop: usize| -> usize {
            if n_samples < nfft {
                0
            } else {
                1 + (n_samples - nfft) / hop
            }
        };

        let test_cases: Vec<(&str, Overlap)> = vec![
            ("0%", Overlap::NoOverlap {}),
            ("25%", Overlap::TwentyFivePercent {}),
            ("50%", Overlap::FiftyPercent {}),
            ("75%", Overlap::SeventyFivePercent {}),
            ("90%", Overlap::NinetyPercent {}),
        ];

        let chunk_size = 1024usize;

        for (name, overlap) in test_cases {
            let hop = (nfft as i64 - overlap.get_overlap_samples(nfft)) as usize;
            let mut engine = SpectroGramEngine::new(
                nfft,
                overlap,
                WindowType::Hann,
                fs,
                FreqWeighting::Z,
                None,
            )?;

            for chunk in signal.chunks(chunk_size) {
                let arr = Array2::from_shape_vec((chunk.len(), 1), chunk.to_vec())?;
                engine.push(arr.view())?;
            }

            let result = engine.finish()?;
            let got = result.t.len();
            let exp = expected(hop);
            assert_eq!(
                got, exp,
                "{name}: expected {exp} blocks, got {got}\
                 (nfft={nfft}, hop={hop}, nsamples={n_samples})",
            );

            if exp > 0 {
                assert_abs_diff_eq!(result.t[0], nfft as Flt / 2.0 / *fs, epsilon = 1e-6,);
                if exp > 1 {
                    assert_abs_diff_eq!(
                        result.t[1] - result.t[0],
                        hop as Flt / *fs,
                        epsilon = 1e-6,
                    );
                }

                // A 1 kHz sine → clear horizontal line at tone bin
                let freq_axis = &result.freq;
                let tone_hz = 1000.0;
                let tone_bin = freq_axis
                    .iter()
                    .enumerate()
                    .min_by(|(_, a), (_, b)| {
                        let da = (**a - tone_hz).abs();
                        let db = (**b - tone_hz).abs();
                        da.partial_cmp(&db).unwrap()
                    })
                    .map(|(i, _)| i)
                    .unwrap();
                // Tone bin should be loud across all time columns
                let tone_row = result.ap_dB.row(tone_bin);
                let min_dB = tone_row.fold(Flt::INFINITY, |a, &b| a.min(b));
                let max_dB = tone_row.fold(Flt::NEG_INFINITY, |a, &b| a.max(b));
                assert!(
                    min_dB > 0.0,
                    "{name}: tone bin {tone_bin} min = {min_dB} dB"
                );
                // Sine tone → equal magnitude across all blocks (horizontal line)
                let spread = max_dB - min_dB;
                assert!(
                    spread < 1.0,
                    "{name}: tone spread {spread:.3} dB (min={min_dB:.1}, max={max_dB:.1})"
                );
                // Other bins should be much quieter (skip bins near the tone
                // where spectral leakage from the Hann window is expected)
                let guard = 10usize;
                let noise_max = (0..result.ap_dB.nrows())
                    .filter(|&r| r < tone_bin.saturating_sub(guard) || r > tone_bin + guard)
                    .flat_map(|r| result.ap_dB.row(r).iter().copied().collect::<Vec<_>>())
                    .fold(Flt::NEG_INFINITY, |a, b| a.max(b));
                assert!(
                    min_dB > noise_max + 20.0,
                    "{name}: tone {min_dB} dB vs noise {noise_max} dB"
                );
            }
        }

        Ok(())
    }
}

#[cfg(feature = "python-bindings")]
#[cfg_attr(feature = "python-bindings", gen_stub_pymethods, pymethods)]
impl SpectrogramResult {
    #[getter]
    fn ap_dB<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray2<Flt>> {
        PyArray2::from_owned_array(py, self.ap_dB.clone())
    }
    #[getter]
    fn freq<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<Flt>> {
        PyArray1::from_owned_array(py, self.freq.clone())
    }
    #[getter]
    fn t<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray1<Flt>> {
        PyArray1::from_owned_array(py, self.t.clone())
    }
    fn __repr__(&self) -> String {
        format!(
            "SpectrogramResult(ap_dB: {:?}, freq: {:?}, t: {:?})",
            self.ap_dB.shape(),
            self.freq.shape(),
            self.t.shape()
        )
    }
}