audio-lab 0.1.0

use crate::signal::{Spectrogram, TimeSignal};
use ndarray::prelude::*;
use num::complex::Complex64;
use realfft::RealFftPlanner;

pub use crate::dsp::window::WindowFn;
use crate::dsp::window::generate_periodic_window;

#[derive(Debug, thiserror::Error)]
pub enum StftError {
    #[error("window_size must be > 0")]
    WindowSizeZero,
    #[error("hop_size must be > 0")]
    HopSizeZero,
    #[error("window_size ({window_size}) must be <= signal length ({signal_length})")]
    WindowLargerThanSignal {
        window_size: usize,
        signal_length: usize,
    },
}

#[derive(Debug, Clone)]
pub struct StftConfig {
    /// FFT window length in samples.
    pub window_size: usize,
    /// Step between successive frames in samples (default: `window_size / 4`, i.e. 75% overlap).
    pub hop_size: usize,
    /// Tapering window applied to each frame before the FFT.
    pub window_fn: WindowFn,
}

impl StftConfig {
    pub fn new(window_size: usize) -> Self {
        Self {
            window_size,
            hop_size: window_size / 4,
            window_fn: WindowFn::Hann,
        }
    }
}

impl Default for StftConfig {
    fn default() -> Self {
        Self::new(1024)
    }
}

/// Compute the short-time Fourier transform of a [`TimeSignal`].
///
/// Returns a [`Spectrogram`] with shape `(num_channels, num_frames, num_freq_bins)`
/// where `num_freq_bins = window_size / 2 + 1`.
///
/// The number of frames is `(num_samples - window_size) / hop_size + 1` (no padding).
pub fn stft(signal: &TimeSignal, config: &StftConfig) -> Result<Spectrogram, StftError> {
    if config.window_size == 0 {
        return Err(StftError::WindowSizeZero);
    }
    if config.hop_size == 0 {
        return Err(StftError::HopSizeZero);
    }
    let num_samples = signal.num_time_steps();
    if config.window_size > num_samples {
        return Err(StftError::WindowLargerThanSignal {
            window_size: config.window_size,
            signal_length: num_samples,
        });
    }

    let num_channels = signal.num_channels();
    let window_size = config.window_size;
    let hop_size = config.hop_size;
    let num_frames = (num_samples - window_size) / hop_size + 1;
    let num_freq_bins = window_size / 2 + 1;
    let sample_rate = signal.sample_rate();

    let window = generate_periodic_window(config.window_fn, window_size);

    let fft = RealFftPlanner::<f64>::new().plan_fft_forward(window_size);
    let mut data: Array3<Complex64> = Array3::zeros((num_channels, num_frames, num_freq_bins));

    // Buffers reused for every frame. The FFT writes straight into `data`.
    let mut windowed_frame = fft.make_input_vec();
    let mut scratch = fft.make_scratch_vec();

    for (ch_data, mut ch_out) in signal.channel_iter().zip(data.outer_iter_mut()) {
        for (frame_idx, mut frame_out) in ch_out.outer_iter_mut().enumerate() {
            let start = frame_idx * hop_size;
            azip!((wf in &mut windowed_frame, &s in ch_data.slice(s![start..start + window_size]), &w in &window) {
                *wf = s * w;
            });

            fft.process_with_scratch(
                &mut windowed_frame,
                frame_out
                    .as_slice_mut()
                    .expect("spectrogram rows are contiguous"),
                &mut scratch,
            )
            .expect("buffer lengths come from the FFT plan");
        }
    }

    let mut spectrogram =
        Spectrogram::new(data, sample_rate, window_size, hop_size, config.window_fn)
            .expect("inputs were validated above and the data has window_size / 2 + 1 bins");
    for ch in 0..signal.num_channels() {
        spectrogram.set_channel_label(ch, signal.channel_label(ch));
    }
    Ok(spectrogram)
}

#[cfg(test)]
mod tests {
    use crate::signal::{
        NoiseConfig, SineConfig, Spectrum, SweepConfig, SweepType, TimeSignal, generate_noise,
        generate_sine, generate_sweep,
    };
    use approx::assert_abs_diff_eq;
    use ndarray::arr2;

    use super::*;

    #[test]
    fn stft_errors() {
        let signal = TimeSignal::new(arr2(&[[0.0; 16]]), 16.0).unwrap();

        assert!(matches!(
            stft(
                &signal,
                &StftConfig {
                    window_size: 0,
                    hop_size: 4,
                    window_fn: WindowFn::Hann
                }
            ),
            Err(StftError::WindowSizeZero)
        ));
        assert!(matches!(
            stft(
                &signal,
                &StftConfig {
                    window_size: 8,
                    hop_size: 0,
                    window_fn: WindowFn::Hann
                }
            ),
            Err(StftError::HopSizeZero)
        ));
        assert!(matches!(
            stft(
                &signal,
                &StftConfig {
                    window_size: 32,
                    hop_size: 8,
                    window_fn: WindowFn::Hann
                }
            ),
            Err(StftError::WindowLargerThanSignal { .. })
        ));
        // A 2-sample Hann window works: its periodic form is [0, 1], not all zeros.
        let spec = stft(
            &signal,
            &StftConfig {
                window_size: 2,
                hop_size: 1,
                window_fn: WindowFn::Hann,
            },
        )
        .unwrap();
        assert!(spec.normalization().coherent_gain() > 0.0);
    }

    #[test]
    fn stft_dimensions() {
        // 32 samples, window=8, hop=4 → (32-8)/4+1 = 7 frames, 5 freq bins
        let signal = TimeSignal::new(arr2(&[[1.0; 32]]), 16.0).unwrap();
        let config = StftConfig {
            window_size: 8,
            hop_size: 4,
            window_fn: WindowFn::Rectangular,
        };
        let spec = stft(&signal, &config).unwrap();

        assert_eq!(spec.num_channels(), 1);
        assert_eq!(spec.num_frames(), 7);
        assert_eq!(spec.num_freq_bins(), 5);
        assert_eq!(spec.sample_rate(), 16.0);
        assert_eq!(spec.window_size(), 8);
        assert_eq!(spec.hop_size(), 4);
    }

    #[test]
    fn stft_frame_times_and_freq_bins() {
        let signal = TimeSignal::new(arr2(&[[0.0; 16]]), 8.0).unwrap();
        // window=8, hop=4 → (16-8)/4+1 = 3 frames
        let config = StftConfig {
            window_size: 8,
            hop_size: 4,
            window_fn: WindowFn::Rectangular,
        };
        let spec = stft(&signal, &config).unwrap();

        // frame centers: (0*4+4)/8, (1*4+4)/8, (2*4+4)/8 = 0.5, 1.0, 1.5 s
        assert_abs_diff_eq!(spec.frame_times()[0], 0.5, epsilon = 1e-10);
        assert_abs_diff_eq!(spec.frame_times()[1], 1.0, epsilon = 1e-10);
        assert_abs_diff_eq!(spec.frame_times()[2], 1.5, epsilon = 1e-10);

        // freq bins: 0, 1, 2, 3, 4 Hz  (sr=8, N=8 → bin spacing = 8/8 = 1 Hz)
        assert_abs_diff_eq!(spec.freq_bins()[0], 0.0, epsilon = 1e-10);
        assert_abs_diff_eq!(spec.freq_bins()[4], 4.0, epsilon = 1e-10);
    }

    #[test]
    fn stft_dc_signal() {
        // A constant signal of amplitude 1 with a rectangular window should give
        // DC bin = window_size and all other bins ≈ 0.
        let window_size = 8;
        let signal = TimeSignal::new(arr2(&[[1.0; 32]]), 16.0).unwrap();
        let config = StftConfig {
            window_size,
            hop_size: 4,
            window_fn: WindowFn::Rectangular,
        };
        let spec = stft(&signal, &config).unwrap();

        for frame in 0..spec.num_frames() {
            let ch_data = spec.channel(0);
            let dc = ch_data[[frame, 0]];
            assert_abs_diff_eq!(dc.re, window_size as f64, epsilon = 1e-10);
            assert_abs_diff_eq!(dc.im, 0.0, epsilon = 1e-10);
            for bin in 1..spec.num_freq_bins() {
                assert_abs_diff_eq!(ch_data[[frame, bin]].norm(), 0.0, epsilon = 1e-8);
            }
        }
    }

    #[test]
    fn stft_multi_channel() {
        let data = arr2(&[[1.0; 32], [2.0; 32]]);
        let signal = TimeSignal::new(data, 16.0).unwrap();
        let config = StftConfig {
            window_size: 8,
            hop_size: 4,
            window_fn: WindowFn::Rectangular,
        };
        let spec = stft(&signal, &config).unwrap();

        assert_eq!(spec.num_channels(), 2);
        // Channel 1 DC should be 2x channel 0 DC
        for frame in 0..spec.num_frames() {
            let dc_ch0 = spec.channel(0)[[frame, 0]].re;
            let dc_ch1 = spec.channel(1)[[frame, 0]].re;
            assert_abs_diff_eq!(dc_ch1, 2.0 * dc_ch0, epsilon = 1e-10);
        }
    }

    #[test]
    fn calibrated_spectra_for_bin_centered_sine() {
        let sample_rate = 8.0;
        let window_size = 8;
        let signal = TimeSignal::new(
            arr2(&[std::array::from_fn::<f64, 8, _>(|n| {
                (std::f64::consts::TAU * 2.0 * n as f64 / window_size as f64).sin()
            })]),
            sample_rate,
        )
        .unwrap();
        let spec = stft(
            &signal,
            &StftConfig {
                window_size,
                hop_size: window_size,
                window_fn: WindowFn::Rectangular,
            },
        )
        .unwrap();

        assert_abs_diff_eq!(spec.raw_magnitude()[[0, 0, 2]], 4.0, epsilon = 1e-10);
        assert_abs_diff_eq!(spec.amplitude_spectrum()[[0, 0, 2]], 1.0, epsilon = 1e-10);
        assert_abs_diff_eq!(spec.power_spectrum()[[0, 0, 2]], 0.5, epsilon = 1e-10);
        assert_abs_diff_eq!(
            spec.power_spectral_density()[[0, 0, 2]],
            0.5,
            epsilon = 1e-10
        );
    }

    #[test]
    fn calibrated_dc_amplitude_is_window_corrected() {
        let signal = TimeSignal::new(arr2(&[[1.0; 8]]), 8.0).unwrap();
        let spec = stft(
            &signal,
            &StftConfig {
                window_size: 8,
                hop_size: 8,
                window_fn: WindowFn::Hann,
            },
        )
        .unwrap();

        assert_abs_diff_eq!(spec.amplitude_spectrum()[[0, 0, 0]], 1.0, epsilon = 1e-10);
    }

    #[test]
    fn stft_generated_sine_has_peak_at_expected_frequency() {
        let signal = generate_sine(
            800,
            440.0,
            SineConfig {
                amplitude: 0.75,
                sample_rate: 8_000.0,
                num_channels: 1,
            },
        )
        .unwrap();
        let spec = stft(
            &signal,
            &StftConfig {
                window_size: 800,
                hop_size: 800,
                window_fn: WindowFn::Rectangular,
            },
        )
        .unwrap();

        let magnitudes = spec.amplitude_spectrum();
        let peak_bin = (0..spec.num_freq_bins())
            .max_by(|&left, &right| {
                magnitudes[[0, 0, left]]
                    .partial_cmp(&magnitudes[[0, 0, right]])
                    .unwrap()
            })
            .unwrap();

        assert_abs_diff_eq!(spec.freq_bins()[peak_bin], 440.0, epsilon = 1e-12);
        assert_abs_diff_eq!(magnitudes[[0, 0, peak_bin]], 0.75, epsilon = 1e-10);
    }

    #[test]
    fn stft_generated_sweep_shows_increasing_dominant_frequency() {
        let signal = generate_sweep(
            4_096,
            200.0..3_200.0,
            SweepConfig {
                sample_rate: 16_000.0,
                fade_out: 0,
                sweep_type: SweepType::Linear,
                ..Default::default()
            },
        )
        .unwrap();
        let spec = stft(
            &signal,
            &StftConfig {
                window_size: 512,
                hop_size: 256,
                window_fn: WindowFn::Hann,
            },
        )
        .unwrap();
        let magnitudes = spec.amplitude_spectrum();

        let dominant_frequency = |frame: usize| {
            let peak_bin = (0..spec.num_freq_bins())
                .max_by(|&left, &right| {
                    magnitudes[[0, frame, left]]
                        .partial_cmp(&magnitudes[[0, frame, right]])
                        .unwrap()
                })
                .unwrap();
            spec.freq_bins()[peak_bin]
        };

        let first_peak = dominant_frequency(0);
        let last_peak = dominant_frequency(spec.num_frames() - 1);

        assert!(first_peak < last_peak);
        assert!(first_peak >= 0.0);
        assert!(last_peak <= signal.sample_rate() / 2.0);
    }

    #[test]
    fn stft_generated_noise_has_broadband_energy() {
        let signal = generate_noise(
            2_048,
            NoiseConfig {
                sample_rate: 16_000.0,
                spectrum: Spectrum::White,
                seed: Some(123),
                ..Default::default()
            },
        )
        .unwrap();
        let spec = stft(
            &signal,
            &StftConfig {
                window_size: 256,
                hop_size: 128,
                window_fn: WindowFn::Hann,
            },
        )
        .unwrap();
        let magnitudes = spec.amplitude_spectrum();

        for frame in 0..spec.num_frames() {
            let frame_energy: f64 = (0..spec.num_freq_bins())
                .map(|bin| magnitudes[[0, frame, bin]])
                .sum();
            assert!(frame_energy > 0.0);
        }

        let active_bins = (0..spec.num_freq_bins())
            .filter(|&bin| magnitudes[[0, 0, bin]] > 1e-3)
            .count();
        assert!(active_bins > spec.num_freq_bins() / 2);
    }
}