lasprs 0.14.3

//! Spectrogram computation on measurement data.

use super::{error::*, *};
use crate::ps::{CPSSettings, SpectroGramEngine, SpectrogramResult};
use crate::*;
use anyhow::Result;
use snafu::prelude::*;

impl SharedMeasurement {
    /// Compute a spectrogram of the auto power of a single channel.
    ///
    /// The auto power is computed per FFT block (with the configured overlap),
    /// in dB relative to `ref_value` (defaults to the acoustic reference level
    /// of 20 µPa). Sensitivities are applied, so the levels are calibrated.
    ///
    /// # Arguments
    ///
    /// * `settings` - FFT settings (nfft, overlap, window, istart, istop)
    /// * `fw` - Frequency weighting applied to the power
    /// * `channel` - Channel to compute the auto power for
    /// * `ref_value` - Reference value (linear) for the dB conversion.
    ///   Defaults to [crate::slm::Lref_default] when `None`.
    pub fn spectrogram(
        &self,
        settings: &CPSSettings,
        fw: FreqWeighting,
        channel: usize,
        ref_value: Option<StrictlyPositive>,
    ) -> Result<SpectrogramResult> {
        let meas = self.read();
        let nchannels = meas.nchannels();
        ensure!(
            channel < nchannels,
            ChannelIdxOutOfBoundsSnafu {
                channel_idx: channel,
                max_channels: nchannels
            }
        );

        let mut engine = SpectroGramEngine::new(
            settings.nfft,
            settings.overlap,
            settings.window,
            meas.samplerate(),
            fw,
            ref_value,
        )?;

        let blockiter = meas.data_iter(Some(&[channel]), settings.istart, settings.istop)?;
        for block in blockiter {
            engine.push(block.view())?;
        }
        let mut result = engine.finish()?;

        // The engine's time axis is relative to the start of the data fed;
        // shift to the absolute start of the measurement.
        let offset = if let Some(istart) = settings.istart {
            istart as Flt / *meas.samplerate()
        } else {
            0.0
        };
        // Fix the time offset.
        result.t.mapv_inplace(|t| t + offset);
        Ok(result)
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use crate::daq::{
        DaqChannel, DaqConfig, RecordSettings, RecordStatus, Recording, StreamMgr, StreamType,
    };
    use crate::measurement::MeasurementType;
    use crate::ps::CPSSettingsBuilder;
    use crate::siggen::{SiggenCommand, SourceDescriptor};
    use approx::assert_abs_diff_eq;
    use std::time::Duration;

    /// Record a 1 kHz sine at √2 amplitude (RMS = 1 → 94 dB SPL) via loopback
    /// and verify the spectrogram structure and levels.
    #[cfg(feature = "loopback-api")]
    #[test]
    fn test_spectrogram_on_measurement() -> anyhow::Result<()> {
        let tmpdir = tempfile::tempdir()?;
        let mut smgr = StreamMgr::new_with_devices();

        let in_cfg = DaqConfig {
            api: crate::daq::DaqApiDescriptor::Loopback,
            device_name: "Loopback".into(),
            inchannel_config: vec![
                DaqChannel::defaultAudio("ch0"),
                DaqChannel::defaultAudio("ch1"),
            ],
            outchannel_config: vec![],
            dtype: crate::daq::DataType::F32,
            sampleRateIndex: 0,
            framesPerBlockIndex: 0,
            monitorOutput: false,
        };
        let out_cfg = {
            let mut cfg = in_cfg.clone();
            cfg.inchannel_config = vec![];
            cfg.outchannel_config = vec![
                DaqChannel::defaultAudio("ch0"),
                DaqChannel::defaultAudio("ch1"),
            ];
            cfg
        };

        smgr.startStream(StreamType::Input, &in_cfg)?;
        smgr.startStream(StreamType::Output, &out_cfg)?;
        smgr.setSiggenSource(SourceDescriptor::Sine {
            frequency: (1000.).try_into().unwrap(),
        })?;
        smgr.siggenCommand(SiggenCommand::SetMuteAllChannels { mute: false })?;
        // Set amplitude to √2 → RMS = 1 → 94 dB SPL
        smgr.siggenCommand(SiggenCommand::SetAllGains {
            g: sq2.try_into().unwrap(),
        })?;

        let duration = Duration::from_secs(1);
        let rec_settings = RecordSettings::new(
            "spec_test",
            false,
            Some(duration),
            None,
            None,
            Some(MeasurementType::NotSpecific {}),
            false,
            Some(tmpdir.path()),
            vec![],
        )?;

        let mut recording = Recording::new(rec_settings, &mut smgr)?;
        loop {
            if matches!(recording.status(), RecordStatus::Finished { .. }) {
                break;
            }
            std::thread::sleep(Duration::from_millis(10));
        }
        recording.stop();

        let m = recording.getMeasurement().unwrap();
        drop(recording);
        drop(smgr);

        let fs = *m.read().samplerate();
        let nfft = 1024usize;
        let settings = CPSSettingsBuilder::default()
            .nfft(nfft)
            .overlap(Overlap::default())
            .build()
            .unwrap();

        let res = m.spectrogram(&settings, FreqWeighting::Z, 0, None)?;

        // Single-sided frequency axis: nfft/2 + 1 bins
        assert_eq!(res.freq.len(), nfft / 2 + 1);
        assert_abs_diff_eq!(res.freq[1] - res.freq[0], fs / nfft as Flt, epsilon = 1e-3);

        // Time axis: block center times, hop = nfft/2 for the default 50%
        // overlap. The first block starts at t=0, so its center is at nfft/2
        // samples.
        assert!(!res.t.is_empty(), "Should have at least one time block");
        assert_abs_diff_eq!(res.t[0], (nfft / 2) as Flt / fs, epsilon = 1e-6);
        if res.t.len() > 1 {
            let hop = nfft / 2;
            assert_abs_diff_eq!(res.t[1] - res.t[0], hop as Flt / fs, epsilon = 1e-6);
        }

        // ap_dB shape: (nfreq, ntime)
        assert_eq!(res.ap_dB.shape(), [nfft / 2 + 1, res.t.len()]);

        // Total power per time column: the window is power-normalized, so the
        // sum of per-bin powers equals the signal power (RMS² = 1 → 94 dB SPL).
        // The dB values are already relative to ref², so sum the converted
        // linear powers directly.
        let last_col = res.ap_dB.column(res.t.len() - 1);
        let total_pwr: Flt = last_col.mapv(|v| Flt::powf(10., v / 10.)).sum();
        let total_dB = 10. * Flt::log10(total_pwr);
        assert_abs_diff_eq!(total_dB, 93.98, epsilon = 0.5);

        // The tone should clearly be present: the maximum bin level of the
        // last column must be well above the broadband level.
        let max_bin_dB = last_col.fold(Flt::NEG_INFINITY, |a, &b| a.max(b));
        assert!(
            max_bin_dB > 90.,
            "Expected a clear tone at ~94 dB, got max bin {max_bin_dB} dB"
        );

        Ok(())
    }

    /// Generate a sine wave WAV, import it, and verify spectrogram block
    /// counts are correct across different overlap settings.
    #[test]
    fn test_spectrogram_block_counts() -> anyhow::Result<()> {
        use hound::{SampleFormat, WavSpec, WavWriter};

        let dir = tempfile::tempdir()?;
        let dirpath = dir.path();
        let wav_path = dirpath.join("sine.wav");

        let fs = 48000u32;
        let duration_secs = 2.0;
        let n_samples = (fs as f64 * duration_secs) as usize;
        let freq = 1000.0f32;

        {
            let spec = WavSpec {
                channels: 1,
                sample_rate: fs,
                bits_per_sample: 32,
                sample_format: SampleFormat::Float,
            };
            let mut writer = WavWriter::create(&wav_path, spec)?;
            for i in 0..n_samples {
                let t = i as f32 / fs as f32;
                let val = (2.0 * std::f32::consts::PI * freq * t).sin();
                writer.write_sample(val)?;
            }
            writer.finalize()?;
        }

        let h5_path = dirpath.join("sine_h5");
        let meas = Measurement::from_wav(&wav_path, Some(&h5_path), None, None, None, None, None)?;

        // let nfft = 4096usize;
        let nfft = 48000;
        // Expected blocks: 1 + floor((n_samples - nfft) / hop_size)
        let expected_blocks = |hop: usize| -> usize {
            if n_samples < nfft {
                0
            } else {
                1 + (n_samples - nfft) / hop
            }
        };

        let test_cases: Vec<(Overlap, usize)> = vec![
            (Overlap::NoOverlap {}, nfft),
            (Overlap::TwentyFivePercent {}, nfft - nfft / 4),
            (Overlap::FiftyPercent {}, nfft / 2),
            (Overlap::SeventyFivePercent {}, nfft / 4),
            (Overlap::NinetyPercent {}, nfft / 10),
        ];

        for (overlap, hop) in test_cases {
            let settings = CPSSettingsBuilder::default()
                .nfft(nfft)
                .overlap(overlap)
                .build()
                .unwrap();
            let res = meas.spectrogram(&settings, FreqWeighting::Z, 0, None)?;
            let expected = expected_blocks(hop);
            assert_eq!(
                res.t.len(),
                expected,
                "Overlap {:?}: expected {} blocks (hop={}), got {}",
                overlap,
                expected,
                hop,
                res.t.len(),
            );
            // Verify the first block center is at nfft/2 / fs
            if expected > 0 {
                let first_center = nfft as Flt / 2.0 / fs as Flt;
                assert_abs_diff_eq!(res.t[0], first_center, epsilon = 1e-6);
            }
        }

        Ok(())
    }
}