spectrs 0.1.3

use rayon::prelude::*;
use rustfft::{FftPlanner, num_complex::Complex};
use std::f32::consts::PI;

// Different spectrogram types
#[derive(Debug, Clone, Copy)]
pub enum SpectrogramType {
    Magnitude,
    Power,
}

/// Create Hann window, see e.g. https://en.wikipedia.org/wiki/Hann_function
fn create_hann_window(length: usize) -> Vec<f32> {
    (0..length)
        .map(|i| 0.5 * (1.0 - (2.0 * PI * i as f32 / (length - 1) as f32).cos()))
        .collect()
}

/// Compute the spectrogram (single-threaded)
/// n_samples: number of samples in each Fast Fourier Transform (FFT) window
/// hop_length: stride between windows, i.e. number of samples between successive FFT frames
/// win_length: number of samples in the window function applied before FFT
/// Pad with zeros if needed. This is because usually win_length < n_samples
/// and the missing are just zeros (in this case complex zeros)
pub fn compute_spectrogram(
    audio: &[f32],
    n_samples: usize,
    hop_length: usize,
    win_length: usize,
    center: bool,
    spectrogram_type: SpectrogramType,
) -> Vec<Vec<f32>> {
    // Set-up FFT
    let mut planner = FftPlanner::<f32>::new();
    let fft = planner.plan_fft_forward(n_samples);

    // Choose the transformation function to create the spectrogram
    let transform_fn: fn(&Complex<f32>) -> f32 = match spectrogram_type {
        SpectrogramType::Magnitude => |c| c.norm(),
        SpectrogramType::Power => |c| c.norm_sqr(),
    };

    // Create (Hann) window
    let window = create_hann_window(win_length);

    // Determine the number of frames
    let n_frames = (audio.len().saturating_sub(win_length)) / hop_length + 1;

    // Determine number of frequency bins
    let n_freq_bins = n_samples / 2 + 1;

    // Directly create spectrogram in [freq][time] format (no transpose needed)
    let mut spectrogram = vec![vec![0.0f32; n_frames]; n_freq_bins];

    // Sequential loop over frames
    for frame_idx in 0..n_frames {
        // Determine start and end sample for each frame
        let start = frame_idx * hop_length;
        let end = (start + win_length).clamp(0, audio.len());

        // Skip if start is beyond the end of the file
        if start > audio.len() {
            continue;
        }

        // Init buffer to be filled with windowed audio
        let mut frame = vec![Complex::<f32>::new(0.0, 0.0); n_samples];

        // Add an offset if the window needs to be centered
        let centering_offset = if center {
            (n_samples - win_length) / 2_usize
        } else {
            0_usize
        };

        // Window & copy into complex buffer
        let src = &audio[start..end];
        let win = &window[..src.len()];
        for (dst, (&s, &w)) in frame
            .iter_mut()
            .skip(centering_offset)
            .zip(src.iter().zip(win.iter()))
        {
            dst.re = s * w; // Convolve audio and window
            dst.im = 0.0; // No imaginary part
        }

        // Run FFT
        fft.process(&mut frame);

        // Store positive freqs only and apply transformation fn
        for (freq_idx, c) in frame.iter().take(n_freq_bins).enumerate() {
            spectrogram[freq_idx][frame_idx] = transform_fn(c);
        }
    }

    spectrogram
}

/// Compute the spectrogram (parallelized with rayon)
/// n_samples: number of samples in each Fast Fourier Transform (FFT) window
/// hop_length: stride between windows, i.e. number of samples between successive FFT frames
/// win_length: number of samples in the window function applied before FFT
/// Pad with zeros if needed. This is because usually win_length < n_samples
/// and the missing are just zeros (in this case complex zeros)
pub fn par_compute_spectrogram(
    audio: &[f32],
    n_samples: usize,
    hop_length: usize,
    win_length: usize,
    center: bool,
    spectrogram_type: SpectrogramType,
) -> Vec<Vec<f32>> {
    // Set-up FFT
    let mut planner = FftPlanner::<f32>::new();
    let fft = planner.plan_fft_forward(n_samples);

    // Choose the transformation function to create the spectrogram
    let transform_fn: fn(&Complex<f32>) -> f32 = match spectrogram_type {
        SpectrogramType::Magnitude => |c| c.norm(),
        SpectrogramType::Power => |c| c.norm_sqr(),
    };

    // Create (Hann) window
    let window = create_hann_window(win_length);

    // Determine the number of frames
    let n_frames = (audio.len().saturating_sub(win_length)) / hop_length + 1;

    // Determine number of frequency bins
    let n_freq_bins = n_samples / 2 + 1;

    // Frame-major spectrogram for safe parallel writes: spectrogram[frame][freq]
    // Eventually to be transposed
    let mut transposed_spectrogram = vec![vec![0.0f32; n_freq_bins]; n_frames];

    // Parallel loop over frames
    transposed_spectrogram
        .par_iter_mut() // Auto-parallelize with rayon
        .enumerate() // Extract frame idx
        .for_each(|(frame_idx, out_row)| {
            // Determine start and end sample for each frame, recalling that hop_length is a stride
            // If the end is after the end of the audio it might still be good (depending on start, see after)
            let start = frame_idx * hop_length;
            let end = (start + win_length).clamp(0, audio.len());

            // Start is beyond the end of the file
            if start > audio.len() {
                return;
            }

            // Init thread-local buffers to be filled with windowed audio
            let mut frame = vec![Complex::<f32>::new(0.0, 0.0); n_samples];

            // Add an offset if the window needs to be centered
            let centering_offset = if center {
                (n_samples - win_length) / 2_usize
            } else {
                0_usize
            };

            // Window & copy into complex buffer
            let src = &audio[start..end];
            let win = &window[..src.len()];
            for (dst, (&s, &w)) in frame
                .iter_mut()
                .skip(centering_offset)
                .zip(src.iter().zip(win.iter()))
            {
                dst.re = s * w; // Convolve audio and window
                dst.im = 0.0; // No imaginary part
            }

            // Run FFT
            fft.process(&mut frame);

            // Store positive freqs only and apply transformation fn depending on request
            for (k, c) in frame.iter().take(n_freq_bins).enumerate() {
                out_row[k] = transform_fn(c);
            }
        });

    // If your downstream expects [freq][frame], transpose once (cache-friendly)
    let mut spectrogram = vec![vec![0.0f32; n_frames]; n_freq_bins];
    for (t, row) in transposed_spectrogram.into_iter().enumerate() {
        for (f, v) in row.into_iter().enumerate() {
            spectrogram[f][t] = v;
        }
    }
    spectrogram
}