dasp-rs 0.4.0

use std::path::Path;
use std::io::Cursor;

use crate::{core::AudioData, core::AudioError};
use hound::WavReader;
use ndarray::Array2;

/// Calculates the duration of an audio signal in seconds.
///
/// # Arguments
/// * `audio` - Reference to an `AudioData` struct containing samples and sample rate
///
/// # Returns
/// Returns a `f32` representing the duration in seconds.
///
/// # Examples
/// ```
/// use dasp_rs::util::*;
/// use dasp_rs::types::*;
/// let audio = AudioData { samples: vec![0.0; 44100], sample_rate: 44100, channels: 1 };
/// let duration = get_duration(&audio);
/// assert_eq!(duration, 1.0); // 1 second
/// ```
pub fn get_duration(audio: &AudioData) -> f32 {
    audio.samples.len() as f32 / audio.sample_rate as f32
}

/// Calculates the duration of an audio file from its path.
///
/// # Arguments
/// * `path` - Path to the audio file, implementing `AsRef<Path>`
///
/// # Returns
/// Returns a `Result<f32, AudioError>` containing the duration in seconds or an error if loading fails.
///
/// # Errors
/// Returns `AudioError` if the audio file cannot be loaded.
///
/// # Examples
/// ```no_run
/// use dasp_rs::util::*;
/// use dasp_rs::types::*;
/// let duration = get_duration_from_path("test.wav");
/// // Assuming test.wav is 2 seconds long at 44100 Hz
/// assert!(duration.is_ok_and(|d| d == 2.0));
/// ```
pub fn get_duration_from_path<P: AsRef<std::path::Path>>(path: P) -> Result<f32, crate::core::AudioError> {
    let audio = crate::core::load(path, None, None, None, None)?;
    Ok(get_duration(&audio))
}

/// Converts frame indices to sample indices.
///
/// # Arguments
/// * `frames` - Array of frame indices
/// * `hop_length` - Optional hop length in samples (defaults to 512)
/// * `_n_fft` - Optional FFT size (unused, defaults to None)
///
/// # Returns
/// Returns a `Vec<usize>` containing corresponding sample indices.
///
/// # Examples
/// ```
/// use dasp_rs::util::*;
/// use dasp_rs::types::*;
/// let frames = vec![0, 1, 2];
/// let samples = frames_to_samples(&frames, None);
/// assert_eq!(samples, vec![0, 512, 1024]);
/// ```
pub fn frames_to_samples(frames: &[usize], hop_length: Option<usize>) -> Vec<usize> {
    let hop = hop_length.unwrap_or(512);
    frames.iter().map(|&f| f * hop).collect()
}

/// Converts frame indices to time values in seconds.
///
/// # Arguments
/// * `frames` - Array of frame indices
/// * `sr` - Optional sample rate in Hz (defaults to 44100)
/// * `hop_length` - Optional hop length in samples (defaults to 512)
///
/// # Returns
/// Returns a `Vec<f32>` containing corresponding time values in seconds.
///
/// # Examples
/// ```
/// use dasp_rs::util::*;
/// use dasp_rs::types::*;
/// let frames = vec![0, 1, 2];
/// let times = frames_to_time(&frames).compute();
/// assert_eq!(times, vec![0.0, 0.011609977, 0.023219954]); // Approx at 44100 Hz, hop 512
/// ```
pub fn frames_to_time(frames: &[usize]) -> FramesToTimeBuilder<'_> {
    FramesToTimeBuilder { frames, sr: 44100, hop_length: 512 }
}

/// Builder for [`frames_to_time`].
#[derive(Debug, Clone)]
pub struct FramesToTimeBuilder<'a> {
    frames: &'a [usize],
    sr: u32,
    hop_length: usize,
}

impl FramesToTimeBuilder<'_> {
    /// Set the sample rate in Hz (default: 44100).
    #[must_use]
    pub fn sample_rate(mut self, sr: u32) -> Self {
        self.sr = sr;
        self
    }

    /// Set the hop length in samples (default: 512).
    #[must_use]
    pub fn hop_length(mut self, hop_length: usize) -> Self {
        self.hop_length = hop_length;
        self
    }

    /// Convert frame indices to times in seconds.
    pub fn compute(self) -> Vec<f32> {
        self.frames
            .iter()
            .map(|&f| f as f32 * self.hop_length as f32 / self.sr as f32)
            .collect()
    }
}

/// Converts sample indices to frame indices.
///
/// # Arguments
/// * `samples` - Array of sample indices
/// * `hop_length` - Optional hop length in samples (defaults to 512)
///
/// # Returns
/// Returns a `Vec<usize>` containing corresponding frame indices (integer division).
///
/// # Examples
/// ```
/// use dasp_rs::util::*;
/// use dasp_rs::types::*;
/// let samples = vec![0, 512, 1024];
/// let frames = samples_to_frames(&samples, None);
/// assert_eq!(frames, vec![0, 1, 2]);
/// ```
pub fn samples_to_frames(samples: &[usize], hop_length: Option<usize>) -> Vec<usize> {
    let hop = hop_length.unwrap_or(512);
    samples.iter().map(|&s| s / hop).collect()
}

/// Converts sample indices to time values in seconds.
///
/// # Arguments
/// * `samples` - Array of sample indices
/// * `sr` - Optional sample rate in Hz (defaults to 44100)
///
/// # Returns
/// Returns a `Vec<f32>` containing corresponding time values in seconds.
///
/// # Examples
/// ```
/// use dasp_rs::util::*;
/// use dasp_rs::types::*;
/// let samples = vec![0, 44100];
/// let times = samples_to_time(&samples, None);
/// assert_eq!(times, vec![0.0, 1.0]);
/// ```
pub fn samples_to_time(samples: &[usize], sr: Option<u32>) -> Vec<f32> {
    let sample_rate = sr.unwrap_or(44100);
    samples.iter().map(|&s| s as f32 / sample_rate as f32).collect()
}

/// Converts time values in seconds to frame indices.
///
/// # Arguments
/// * `times` - Array of time values in seconds
/// * `sr` - Optional sample rate in Hz (defaults to 44100)
/// * `hop_length` - Optional hop length in samples (defaults to 512)
/// * `_n_fft` - Optional FFT size (unused, defaults to None)
///
/// # Returns
/// Returns a `Vec<usize>` containing corresponding frame indices.
///
/// # Examples
/// ```
/// use dasp_rs::util::*;
/// use dasp_rs::types::*;
/// let times = vec![0.0, 0.011609977];
/// let frames = time_to_frames(&times).compute();
/// assert_eq!(frames, vec![0, 1]);
/// ```
pub fn time_to_frames(times: &[f32]) -> TimeToFramesBuilder<'_> {
    TimeToFramesBuilder { times, sr: 44100, hop_length: 512 }
}

/// Builder for [`time_to_frames`].
#[derive(Debug, Clone)]
pub struct TimeToFramesBuilder<'a> {
    times: &'a [f32],
    sr: u32,
    hop_length: usize,
}

impl TimeToFramesBuilder<'_> {
    /// Set the sample rate in Hz (default: 44100).
    #[must_use]
    pub fn sample_rate(mut self, sr: u32) -> Self {
        self.sr = sr;
        self
    }

    /// Set the hop length in samples (default: 512).
    #[must_use]
    pub fn hop_length(mut self, hop_length: usize) -> Self {
        self.hop_length = hop_length;
        self
    }

    /// Convert times in seconds to frame indices.
    pub fn compute(self) -> Vec<usize> {
        self.times
            .iter()
            .map(|&t| (t * self.sr as f32 / self.hop_length as f32) as usize)
            .collect()
    }
}

/// Converts time values in seconds to sample indices.
///
/// # Arguments
/// * `times` - Array of time values in seconds
/// * `sr` - Optional sample rate in Hz (defaults to 44100)
///
/// # Returns
/// Returns a `Vec<usize>` containing corresponding sample indices.
///
/// # Examples
/// ```
/// use dasp_rs::util::*;
/// use dasp_rs::types::*;
/// let times = vec![0.0, 1.0];
/// let samples = time_to_samples(&times, None);
/// assert_eq!(samples, vec![0, 44100]);
/// ```
pub fn time_to_samples(times: &[f32], sr: Option<u32>) -> Vec<usize> {
    let sample_rate = sr.unwrap_or(44100);
    times.iter().map(|&t| (t * sample_rate as f32) as usize).collect()
}

/// Converts block indices to frame indices.
///
/// # Arguments
/// * `blocks` - Array of block indices
/// * `block_length` - Number of frames per block
///
/// # Returns
/// Returns a `Vec<usize>` containing corresponding frame indices.
///
/// # Examples
/// ```
/// use dasp_rs::util::*;
/// use dasp_rs::types::*;
/// let blocks = vec![0, 1, 2];
/// let frames = blocks_to_frames(&blocks, 10);
/// assert_eq!(frames, vec![0, 10, 20]);
/// ```
pub fn blocks_to_frames(blocks: &[usize], block_length: usize) -> Vec<usize> {
    blocks.iter().map(|&b| b * block_length).collect()
}

/// Converts block indices to sample indices.
///
/// # Arguments
/// * `blocks` - Array of block indices
/// * `block_length` - Number of frames per block
/// * `hop_length` - Optional hop length in samples (defaults to 512)
///
/// # Returns
/// Returns a `Vec<usize>` containing corresponding sample indices.
///
/// # Examples
/// ```
/// use dasp_rs::util::*;
/// use dasp_rs::types::*;
/// let blocks = vec![0, 1];
/// let samples = blocks_to_samples(&blocks, 2, None);
/// assert_eq!(samples, vec![0, 1024]); // 2 frames * 512 hop
/// ```
pub fn blocks_to_samples(blocks: &[usize], block_length: usize, hop_length: Option<usize>) -> Vec<usize> {
    let hop = hop_length.unwrap_or(512);
    blocks.iter().map(|&b| b * block_length * hop).collect()
}

/// Converts block indices to time values in seconds.
///
/// # Arguments
/// * `blocks` - Array of block indices
/// * `block_length` - Number of frames per block
/// * `hop_length` - Optional hop length in samples (defaults to 512)
/// * `sr` - Optional sample rate in Hz (defaults to 44100)
///
/// # Returns
/// Returns a `Vec<f32>` containing corresponding time values in seconds.
///
/// # Examples
/// ```
/// use dasp_rs::util::*;
/// use dasp_rs::types::*;
/// let blocks = vec![0, 1];
/// let times = blocks_to_time(&blocks, 2).compute();
/// assert_eq!(times, vec![0.0, 0.023219954]); // 2 frames * 512 hop / 44100 Hz
/// ```
pub fn blocks_to_time(blocks: &[usize], block_length: usize) -> BlocksToTimeBuilder<'_> {
    BlocksToTimeBuilder { blocks, block_length, hop_length: 512, sr: 44100 }
}

/// Builder for [`blocks_to_time`].
#[derive(Debug, Clone)]
pub struct BlocksToTimeBuilder<'a> {
    blocks: &'a [usize],
    block_length: usize,
    hop_length: usize,
    sr: u32,
}

impl BlocksToTimeBuilder<'_> {
    /// Set the hop length in samples (default: 512).
    #[must_use]
    pub fn hop_length(mut self, hop_length: usize) -> Self {
        self.hop_length = hop_length;
        self
    }

    /// Set the sample rate in Hz (default: 44100).
    #[must_use]
    pub fn sample_rate(mut self, sr: u32) -> Self {
        self.sr = sr;
        self
    }

    /// Convert block indices to times in seconds.
    pub fn compute(self) -> Vec<f32> {
        self.blocks
            .iter()
            .map(|&b| {
                b as f32 * self.block_length as f32 * self.hop_length as f32 / self.sr as f32
            })
            .collect()
    }
}

/// Generates sample indices corresponding to the columns of a 2D array.
///
/// # Arguments
/// * `X` - 2D array (typically a spectrogram)
/// * `hop_length` - Optional hop length in samples (defaults to 512)
///
/// # Returns
/// Returns a `Vec<usize>` containing sample indices for each column of `X`.
///
/// # Examples
/// ```
/// use dasp_rs::util::samples_like;
/// use ndarray::arr2;
/// let x = arr2(&[[1.0, 2.0], [3.0, 4.0]]);
/// let samples = samples_like(&x, None);
/// assert_eq!(samples, vec![0, 512]);
/// ```
pub fn samples_like(x: &Array2<f32>, hop_length: Option<usize>) -> Vec<usize> {
    let hop = hop_length.unwrap_or(512);
    (0..x.shape()[1]).map(|i| i * hop).collect()
}

/// Generates time values corresponding to the columns of a 2D array.
///
/// # Arguments
/// * `X` - 2D array (typically a spectrogram)
/// * `sr` - Optional sample rate in Hz (defaults to 44100)
/// * `hop_length` - Optional hop length in samples (defaults to 512)
///
/// # Returns
/// Returns a `Vec<f32>` containing time values in seconds for each column of `X`.
///
/// # Examples
/// ```
/// use dasp_rs::util::times_like;
/// use ndarray::arr2;
/// let x = arr2(&[[1.0, 2.0], [3.0, 4.0]]);
/// let times = times_like(&x).compute();
/// assert_eq!(times, vec![0.0, 0.011609977]); // 512 hop / 44100 Hz
/// ```
pub fn times_like(x: &Array2<f32>) -> TimesLikeBuilder<'_> {
    TimesLikeBuilder { x, sr: 44100, hop_length: 512 }
}

/// Builder for [`times_like`].
#[derive(Debug, Clone)]
pub struct TimesLikeBuilder<'a> {
    x: &'a Array2<f32>,
    sr: u32,
    hop_length: usize,
}

impl TimesLikeBuilder<'_> {
    /// Set the sample rate in Hz (default: 44100).
    #[must_use]
    pub fn sample_rate(mut self, sr: u32) -> Self {
        self.sr = sr;
        self
    }

    /// Set the hop length in samples (default: 512).
    #[must_use]
    pub fn hop_length(mut self, hop_length: usize) -> Self {
        self.hop_length = hop_length;
        self
    }

    /// Produce the per-column time values in seconds.
    pub fn compute(self) -> Vec<f32> {
        (0..self.x.shape()[1])
            .map(|i| i as f32 * self.hop_length as f32 / self.sr as f32)
            .collect()
    }
}

/// Extracts sample rate from WAV file header.
///
/// Lightweight metadata query without full sample loading.
///
/// # Parameters
/// - `path`: WAV file path (`AsRef<Path>`).
///
/// # Returns
/// - `Ok(u32)`: Sample rate in Hz.
/// - `Err(AudioError)`: I/O or format error.
/// 
/// # Example
/// ```no_run
/// use dasp_rs::util::*;
/// use dasp_rs::types::*;
/// let rate = get_samplerate("audio.wav")?;
/// assert_eq!(rate, 44100);
/// # Ok::<(), Box<dyn std::error::Error>>(())
/// ```
/// # Errors
/// Returns an error if the input is invalid (e.g., empty signal or
/// out-of-range parameters) or if the computation cannot be completed.
pub fn get_samplerate<P: AsRef<Path>>(path: P) -> Result<u32, AudioError> {
    let wav_data = std::fs::read(&path)?;
    let reader = WavReader::new(Cursor::new(wav_data))?;
    Ok(reader.spec().sample_rate)
}

#[cfg(test)]
mod tests {
    use super::*;
    use hound::{SampleFormat, WavSpec, WavWriter};
    use tempfile::NamedTempFile;

    fn write_test_wav(samples: &[f32], sample_rate: u32) -> NamedTempFile {
        let spec = WavSpec {
            channels: 1,
            sample_rate,
            bits_per_sample: 32,
            sample_format: SampleFormat::Float,
        };
        let file = NamedTempFile::new().expect("temp wav");
        let mut writer =
            WavWriter::new(std::io::BufWriter::new(file.reopen().unwrap()), spec).unwrap();
        for &s in samples {
            writer.write_sample(s).unwrap();
        }
        writer.finalize().unwrap();
        file
    }

    #[test]
    fn duration_and_frame_conversions_round_trip() {
        let audio = AudioData {
            samples: vec![0.0; 4410],
            sample_rate: 44100,
            channels: 1,
        };
        assert!((get_duration(&audio) - 0.1).abs() < 1e-6);

        let frames = vec![0, 1, 2, 3];
        let samples = frames_to_samples(&frames, Some(512));
        assert_eq!(samples, vec![0, 512, 1024, 1536]);
        assert_eq!(samples_to_frames(&samples, Some(512)), frames);

        let times = frames_to_time(&frames).sample_rate(44100).hop_length(512).compute();
        let frames_back = time_to_frames(&times).sample_rate(44100).hop_length(512).compute();
        assert_eq!(frames_back, frames);
    }

    #[test]
    fn block_and_time_mappings_align() {
        let blocks = vec![0, 1, 2];
        let frames = blocks_to_frames(&blocks, 4);
        assert_eq!(frames, vec![0, 4, 8]);

        let samples = blocks_to_samples(&blocks, 4, Some(256));
        assert_eq!(samples, vec![0, 1024, 2048]);
        let times = blocks_to_time(&blocks, 4).hop_length(256).sample_rate(44100).compute();
        assert!((times[1] - 1024.0 / 44100.0).abs() < 1e-9);
    }

    #[test]
    fn samples_and_times_like_match_dimensions() {
        let matrix = Array2::from_shape_vec((2, 3), vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        assert_eq!(samples_like(&matrix, Some(256)), vec![0, 256, 512]);
        let times = times_like(&matrix).sample_rate(48000).hop_length(480).compute();
        let expected = vec![0.0, 480.0 / 48000.0, 960.0 / 48000.0];
        for (actual, exp) in times.iter().zip(expected) {
            assert!((actual - exp).abs() < 1e-6);
        }
    }

    #[test]
    fn samplerate_reads_from_wav_header() {
        let file = write_test_wav(&[0.0, 0.0], 22_050);
        let sr = get_samplerate(file.path()).unwrap();
        assert_eq!(sr, 22_050);
    }

    #[test]
    fn duration_from_path_uses_loader() {
        let samples = vec![0.0; 44_100];
        let file = write_test_wav(&samples, 44_100);
        let duration = get_duration_from_path(file.path()).unwrap();
        assert!((duration - 1.0).abs() < 1e-6);
    }
}