lasprs 0.14.2

//! Common WAV types and functions shared between import and export.
use super::{Result, *};
use crate::daq::RawStreamData;
use snafu::prelude::*;

/// WAV file metadata, returned by [`read_wav_info`] (Python only).
#[derive(Debug, Clone)]
#[cfg_attr(feature = "python-bindings", gen_stub_pyclass, pyclass(from_py_object))]
pub struct WavInfo {
    /// Sample rate in Hz
    pub sample_rate: u32,
    /// Number of audio channels
    pub nchannels: usize,
    /// Bits per sample (e.g. 16, 24, 32)
    pub bits_per_sample: u16,
    /// Sample format: "int" or "float"
    pub sample_format: String,
    /// Human-readable format string, e.g. "int16", "float32"
    pub format_str: String,
    /// Total number of frames per channel
    pub nframes: usize,
    /// Duration in seconds
    pub duration: f64,
}

#[cfg(feature = "python-bindings")]
#[cfg_attr(feature = "python-bindings", gen_stub_pymethods, pymethods)]
impl WavInfo {
    #[getter]
    fn sample_rate(&self) -> u32 {
        self.sample_rate
    }
    #[getter]
    fn nchannels(&self) -> usize {
        self.nchannels
    }
    #[getter]
    fn bits_per_sample(&self) -> u16 {
        self.bits_per_sample
    }
    #[getter]
    fn sample_format(&self) -> String {
        self.sample_format.clone()
    }
    #[getter]
    fn format_str(&self) -> String {
        self.format_str.clone()
    }
    #[getter]
    fn nframes(&self) -> usize {
        self.nframes
    }
    #[getter]
    fn duration(&self) -> f64 {
        self.duration
    }
}

/// Policy for handling WAV files with different frame counts when merging
/// multiple files into a single measurement via
/// [`super::wav_import::Measurement::from_wav_files`].
#[cfg_attr(
    feature = "python-bindings",
    gen_stub_pyclass_enum,
    pyclass(eq, eq_int, from_py_object)
)]
#[derive(Debug, Clone, Copy, PartialEq)]
#[repr(u32)]
pub enum FrameMismatchPolicy {
    /// Shorter files are zero-padded at the **end** so that data is
    /// left-aligned (the beginnings of all files are time-aligned).
    AppendZeros = 0,
    /// Shorter files are zero-padded at the **beginning** so that data is
    /// right-aligned (the endings of all files are time-aligned).
    PrependZeros = 1,
}

/// Read all samples from a single `hound::WavReader` into a `RawStreamData`,
/// dispatching to the correct sample type based on `(sample_format, bits)`.
#[macro_export]
macro_rules! read_wav_samples {
    ($reader:expr, $spec_fmt:expr, $bits:expr) => {
        match ($spec_fmt, $bits) {
            (hound::SampleFormat::Int, 8) => {
                let samples = $crate::measurement::wav_common::collect_wav_samples::<i8>($reader)?;
                RawStreamData::new(samples)
            }
            (hound::SampleFormat::Int, 16) => {
                let samples = $crate::measurement::wav_common::collect_wav_samples::<i16>($reader)?;
                RawStreamData::new(samples)
            }
            (hound::SampleFormat::Int, 24) | (hound::SampleFormat::Int, 32) => {
                let samples = $crate::measurement::wav_common::collect_wav_samples::<i32>($reader)?;
                RawStreamData::new(samples)
            }
            (hound::SampleFormat::Float, 32) => {
                let samples = $crate::measurement::wav_common::collect_wav_samples::<f32>($reader)?;
                RawStreamData::new(samples)
            }
            (fmt, bits) => {
                return WAVImportSnafu {
                    msg: format!("Unsupported WAV format: {fmt:?} with {bits} bits per sample"),
                }
                .fail();
            }
        }
    };
}

/// Read all samples from multiple `hound::WavReader`s, interleave them into a
/// single buffer, and return as `RawStreamData`. Dispatches to the correct
/// sample type based on `(sample_format, bits)`.
#[macro_export]
macro_rules! read_and_interleave_wav_samples {
    ($readers:expr, $channels_per_file:expr, $frames_per_file:expr,
     $total_channels:expr, $total_frames:expr, $policy:expr,
     $spec_fmt:expr, $bits:expr) => {
        match ($spec_fmt, $bits) {
            (hound::SampleFormat::Int, 8) => {
                let samples = read_and_interleave::<i8>(
                    $readers,
                    $channels_per_file,
                    $frames_per_file,
                    $total_channels,
                    $total_frames,
                    $policy,
                )?;
                RawStreamData::new(samples)
            }
            (hound::SampleFormat::Int, 16) => {
                let samples = read_and_interleave::<i16>(
                    $readers,
                    $channels_per_file,
                    $frames_per_file,
                    $total_channels,
                    $total_frames,
                    $policy,
                )?;
                RawStreamData::new(samples)
            }
            (hound::SampleFormat::Int, 32) => {
                let samples = read_and_interleave::<i32>(
                    $readers,
                    $channels_per_file,
                    $frames_per_file,
                    $total_channels,
                    $total_frames,
                    $policy,
                )?;
                RawStreamData::new(samples)
            }
            (hound::SampleFormat::Float, 32) => {
                let samples = read_and_interleave::<f32>(
                    $readers,
                    $channels_per_file,
                    $frames_per_file,
                    $total_channels,
                    $total_frames,
                    $policy,
                )?;
                RawStreamData::new(samples)
            }
            (fmt, bits) => {
                return WAVImportSnafu {
                    msg: format!("Unsupported WAV format: {fmt:?} with {bits} bits per sample"),
                }
                .fail();
            }
        }
    };
}

/// Map a WAV spec's `(sample_format, bits_per_sample)` to a `DataType`.
pub fn wav_data_type(spec: &hound::WavSpec) -> Result<DataType> {
    match spec.sample_format {
        hound::SampleFormat::Int => match spec.bits_per_sample {
            8 => Ok(DataType::I8),
            16 => Ok(DataType::I16),
            24 => Ok(DataType::I24),
            32 => Ok(DataType::I32),
            bits => WAVImportSnafu {
                msg: format!("Unsupported bits per sample for integer WAV file: {bits}"),
            }
            .fail(),
        },
        hound::SampleFormat::Float => match spec.bits_per_sample {
            32 => Ok(DataType::F32),
            bits => WAVImportSnafu {
                msg: format!(
                    "Unsupported bits per sample for floating point WAV file: \
                     {bits}. Only 32-bit float is supported."
                ),
            }
            .fail(),
        },
    }
}

/// Build a hound `WavSpec` for the given channel count, sample rate, and
/// data type.
pub fn wav_spec(ch: usize, sr: StrictlyPositive, dt: &DataType) -> hound::WavSpec {
    let (sample_format, bits_per_sample) = match dt {
        DataType::I8 => (hound::SampleFormat::Int, 8),
        DataType::I16 => (hound::SampleFormat::Int, 16),
        DataType::I24 => (hound::SampleFormat::Int, 24),
        DataType::I32 => (hound::SampleFormat::Int, 32),
        DataType::F32 => (hound::SampleFormat::Float, 32),
        DataType::F64 => (hound::SampleFormat::Float, 64),
    };
    hound::WavSpec {
        channels: ch as u16,
        sample_rate: *sr as u32,
        bits_per_sample,
        sample_format,
    }
}

/// Read WAV file metadata using `hound`.
///
/// Returns a [`WavInfo`] struct with all relevant metadata.
#[cfg(feature = "python-bindings")]
#[pyfunction]
#[cfg_attr(feature = "python-bindings", gen_stub_pyfunction)]
pub fn read_wav_info(path: &str) -> PyResult<WavInfo> {
    use hound::WavReader;

    let reader = WavReader::open(path).map_err(|e| {
        pyo3::exceptions::PyIOError::new_err(format!("Failed to open WAV file: {e}"))
    })?;

    let spec = reader.spec();
    let total_samples = reader.len() as usize;
    let nchannels = spec.channels as usize;
    let nframes = total_samples.checked_div(nchannels).unwrap_or(0);
    let duration = if spec.sample_rate > 0 {
        nframes as f64 / spec.sample_rate as f64
    } else {
        0.0
    };

    let sample_format = match spec.sample_format {
        hound::SampleFormat::Int => "int",
        hound::SampleFormat::Float => "float",
    };

    let format_str = match spec.sample_format {
        hound::SampleFormat::Int => match spec.bits_per_sample {
            8 => "int8",
            16 => "int16",
            24 => "int24",
            32 => "int32",
            b => {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "Unsupported integer bits per sample: {b}"
                )));
            }
        },
        hound::SampleFormat::Float => match spec.bits_per_sample {
            32 => "float32",
            64 => "float64",
            b => {
                return Err(pyo3::exceptions::PyValueError::new_err(format!(
                    "Unsupported float bits per sample: {b}"
                )));
            }
        },
    };

    Ok(WavInfo {
        sample_rate: spec.sample_rate,
        nchannels,
        bits_per_sample: spec.bits_per_sample,
        sample_format: sample_format.to_string(),
        format_str: format_str.to_string(),
        nframes,
        duration,
    })
}

/// Collect all samples from a `WavReader` into a `Vec<T>`.
pub fn collect_wav_samples<T: hound::Sample>(
    reader: hound::WavReader<std::io::BufReader<std::fs::File>>,
) -> Result<Vec<T>> {
    let wav_err = |e: hound::Error| {
        WAVImportSnafu {
            msg: format!("Error reading WAV sample: {e}"),
        }
        .build()
    };
    reader
        .into_samples::<T>()
        .map(|s| s.map_err(wav_err))
        .collect::<Result<Vec<T>>>()
}