grafos_dsp/

stages.rs

//! Built-in DSP processing stages.

extern crate alloc;
use alloc::vec;
use alloc::vec::Vec;

use grafos_std::error::FabricError;

use crate::fft;
use crate::types::{Block, Complex, Sample};

/// Trait for a DSP processing stage that transforms blocks.
pub trait DspStage {
    /// Process a block of samples, returning a transformed block.
    fn process(&mut self, block: &Block) -> Result<Block, FabricError>;
}

/// Forward FFT stage: real time-domain samples to complex frequency-domain.
///
/// The output block stores interleaved `[re, im, re, im, ...]` pairs.
/// The block size must be a power of 2.
///
/// When the `gpu` feature is enabled, set `gpu` to `true` to dispatch FFT
/// through `grafos-tensor`'s GPU path. Falls back to CPU if the tensor
/// cannot be placed on GPU.
pub struct FftStage {
    #[cfg(feature = "gpu")]
    gpu: bool,
}

impl FftStage {
    /// Create a CPU-only FFT stage.
    pub fn new() -> Self {
        Self {
            #[cfg(feature = "gpu")]
            gpu: false,
        }
    }

    /// Create an FFT stage with GPU dispatch enabled.
    ///
    /// When the `gpu` feature is not compiled in, this is identical to `new()`.
    pub fn with_gpu() -> Self {
        Self {
            #[cfg(feature = "gpu")]
            gpu: true,
        }
    }

    fn process_cpu(&self, block: &Block) -> Result<Block, FabricError> {
        let n = block.frames();
        if n == 0 || !n.is_power_of_two() {
            return Err(FabricError::CapacityExceeded);
        }

        let channels = block.channels as usize;
        let mut output_data = Vec::with_capacity(n * 2 * channels);

        for ch in 0..channels {
            let channel_data: Vec<f32> = (0..n).map(|i| block.data[i * channels + ch]).collect();
            let freq = fft::fft_real(&channel_data);
            for c in &freq {
                output_data.push(c.re);
                output_data.push(c.im);
            }
        }

        Ok(Block {
            data: output_data,
            sample_rate: block.sample_rate,
            channels: block.channels,
        })
    }

    #[cfg(feature = "gpu")]
    fn process_gpu(&self, block: &Block) -> Result<Block, FabricError> {
        let n = block.frames();
        if n == 0 || !n.is_power_of_two() {
            return Err(FabricError::CapacityExceeded);
        }

        let channels = block.channels as usize;
        let mut output_data = Vec::with_capacity(n * 2 * channels);

        for ch in 0..channels {
            let channel_data: Vec<f32> = (0..n).map(|i| block.data[i * channels + ch]).collect();
            let tensor = grafos_tensor::FabricTensor::from_slice(&[n], &channel_data)?;
            let gpu_tensor = match tensor.to_gpu() {
                Ok(t) => t,
                Err(_) => {
                    // GPU unavailable, fall back to CPU for this channel
                    let freq = fft::fft_real(&channel_data);
                    for c in &freq {
                        output_data.push(c.re);
                        output_data.push(c.im);
                    }
                    continue;
                }
            };
            let freq_tensor = gpu_tensor.fft()?;
            let cpu_result = freq_tensor.to_cpu()?;
            output_data.extend_from_slice(cpu_result.as_slice());
        }

        Ok(Block {
            data: output_data,
            sample_rate: block.sample_rate,
            channels: block.channels,
        })
    }
}

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

impl DspStage for FftStage {
    fn process(&mut self, block: &Block) -> Result<Block, FabricError> {
        #[cfg(feature = "gpu")]
        if self.gpu {
            return self.process_gpu(block);
        }
        self.process_cpu(block)
    }
}

/// Inverse FFT stage: complex frequency-domain to real time-domain.
///
/// Expects input with interleaved `[re, im, re, im, ...]` pairs per channel.
///
/// When the `gpu` feature is enabled, set `gpu` to `true` to dispatch IFFT
/// through `grafos-tensor`'s GPU path. Falls back to CPU if the tensor
/// cannot be placed on GPU.
pub struct IfftStage {
    #[cfg(feature = "gpu")]
    gpu: bool,
}

impl IfftStage {
    /// Create a CPU-only IFFT stage.
    pub fn new() -> Self {
        Self {
            #[cfg(feature = "gpu")]
            gpu: false,
        }
    }

    /// Create an IFFT stage with GPU dispatch enabled.
    ///
    /// When the `gpu` feature is not compiled in, this is identical to `new()`.
    pub fn with_gpu() -> Self {
        Self {
            #[cfg(feature = "gpu")]
            gpu: true,
        }
    }

    fn process_cpu(&self, block: &Block) -> Result<Block, FabricError> {
        let channels = block.channels as usize;
        if channels == 0 {
            return Err(FabricError::CapacityExceeded);
        }

        let total_per_channel = block.data.len() / channels;
        if !total_per_channel.is_multiple_of(2) {
            return Err(FabricError::CapacityExceeded);
        }
        let n_complex = total_per_channel / 2;
        if !n_complex.is_power_of_two() {
            return Err(FabricError::CapacityExceeded);
        }

        let mut output_data = vec![0.0f32; n_complex * channels];

        for ch in 0..channels {
            let base = ch * total_per_channel;
            let complex_data: Vec<Complex> = (0..n_complex)
                .map(|i| Complex::new(block.data[base + i * 2], block.data[base + i * 2 + 1]))
                .collect();
            let time_domain = fft::ifft_real(&complex_data);
            for (i, &s) in time_domain.iter().enumerate() {
                output_data[i * channels + ch] = s;
            }
        }

        Ok(Block {
            data: output_data,
            sample_rate: block.sample_rate,
            channels: block.channels,
        })
    }

    #[cfg(feature = "gpu")]
    fn process_gpu(&self, block: &Block) -> Result<Block, FabricError> {
        let channels = block.channels as usize;
        if channels == 0 {
            return Err(FabricError::CapacityExceeded);
        }

        let total_per_channel = block.data.len() / channels;
        if total_per_channel % 2 != 0 {
            return Err(FabricError::CapacityExceeded);
        }
        let n_complex = total_per_channel / 2;
        if !n_complex.is_power_of_two() {
            return Err(FabricError::CapacityExceeded);
        }

        let mut output_data = vec![0.0f32; n_complex * channels];

        for ch in 0..channels {
            let base = ch * total_per_channel;
            let freq_data: Vec<f32> = block.data[base..base + total_per_channel].to_vec();
            let tensor = grafos_tensor::FabricTensor::from_slice(&[total_per_channel], &freq_data)?;
            let gpu_tensor = match tensor.to_gpu() {
                Ok(t) => t,
                Err(_) => {
                    // GPU unavailable, fall back to CPU for this channel
                    let complex_data: Vec<Complex> = (0..n_complex)
                        .map(|i| {
                            Complex::new(block.data[base + i * 2], block.data[base + i * 2 + 1])
                        })
                        .collect();
                    let time_domain = fft::ifft_real(&complex_data);
                    for (i, &s) in time_domain.iter().enumerate() {
                        output_data[i * channels + ch] = s;
                    }
                    continue;
                }
            };
            let time_tensor = gpu_tensor.ifft()?;
            let cpu_result = time_tensor.to_cpu()?;
            let recovered = cpu_result.as_slice();
            for (i, &s) in recovered.iter().enumerate() {
                output_data[i * channels + ch] = s;
            }
        }

        Ok(Block {
            data: output_data,
            sample_rate: block.sample_rate,
            channels: block.channels,
        })
    }
}

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

impl DspStage for IfftStage {
    fn process(&mut self, block: &Block) -> Result<Block, FabricError> {
        #[cfg(feature = "gpu")]
        if self.gpu {
            return self.process_gpu(block);
        }
        self.process_cpu(block)
    }
}

/// FIR (Finite Impulse Response) filter stage.
///
/// Direct-form convolution with configurable coefficients.
/// Maintains an overlap buffer for continuous stream processing.
pub struct FirFilterStage {
    coefficients: Vec<f32>,
    /// Per-channel overlap buffer from previous blocks.
    overlap: Vec<Vec<f32>>,
    initialized: bool,
}

impl FirFilterStage {
    /// Create a new FIR filter with the given coefficients.
    pub fn new(coefficients: Vec<f32>) -> Self {
        Self {
            coefficients,
            overlap: Vec::new(),
            initialized: false,
        }
    }
}

impl DspStage for FirFilterStage {
    fn process(&mut self, block: &Block) -> Result<Block, FabricError> {
        let channels = block.channels as usize;
        let frames = block.frames();
        let order = self.coefficients.len();

        if !self.initialized {
            self.overlap = (0..channels).map(|_| vec![0.0; order - 1]).collect();
            self.initialized = true;
        }

        let mut output_data = vec![0.0f32; block.data.len()];

        for ch in 0..channels {
            let input: Vec<f32> = (0..frames).map(|i| block.data[i * channels + ch]).collect();

            // Prepend overlap from previous block
            let mut extended = self.overlap[ch].clone();
            extended.extend_from_slice(&input);

            for i in 0..frames {
                let mut sum = 0.0f32;
                for (j, &coeff) in self.coefficients.iter().enumerate() {
                    sum += coeff * extended[i + order - 1 - j];
                }
                output_data[i * channels + ch] = sum;
            }

            // Save overlap for next block
            let overlap_start = input.len().saturating_sub(order - 1);
            self.overlap[ch] = input[overlap_start..].to_vec();
            // Pad with zeros if input was shorter than order-1
            while self.overlap[ch].len() < order - 1 {
                self.overlap[ch].insert(0, 0.0);
            }
        }

        Ok(Block {
            data: output_data,
            sample_rate: block.sample_rate,
            channels: block.channels,
        })
    }
}

/// IIR (Infinite Impulse Response) filter stage.
///
/// Direct-form II transposed implementation with configurable
/// feedforward (b) and feedback (a) coefficients.
pub struct IirFilterStage {
    /// Feedforward coefficients (b[0], b[1], ..., b[M]).
    b: Vec<f32>,
    /// Feedback coefficients (a[1], a[2], ..., a[N]).
    /// a[0] is assumed to be 1.0 and is not stored.
    a: Vec<f32>,
    /// Per-channel delay line.
    delay: Vec<Vec<f32>>,
    initialized: bool,
}

impl IirFilterStage {
    /// Create a new IIR filter.
    ///
    /// `b` are the feedforward coefficients, `a` are the feedback coefficients.
    /// `a[0]` is assumed to be 1.0 and should not be included in `a`.
    pub fn new(b: Vec<f32>, a: Vec<f32>) -> Self {
        Self {
            b,
            a,
            delay: Vec::new(),
            initialized: false,
        }
    }
}

impl DspStage for IirFilterStage {
    fn process(&mut self, block: &Block) -> Result<Block, FabricError> {
        let channels = block.channels as usize;
        let frames = block.frames();
        let order = core::cmp::max(self.b.len(), self.a.len() + 1);

        if !self.initialized {
            self.delay = (0..channels).map(|_| vec![0.0; order]).collect();
            self.initialized = true;
        }

        let mut output_data = vec![0.0f32; block.data.len()];

        for ch in 0..channels {
            let delay = &mut self.delay[ch];

            for i in 0..frames {
                let x = block.data[i * channels + ch];

                // Direct-form II transposed
                let y = self.b.first().copied().unwrap_or(0.0) * x + delay[0];

                // Update delay line
                for j in 0..order - 1 {
                    let b_term = if j + 1 < self.b.len() {
                        self.b[j + 1] * x
                    } else {
                        0.0
                    };
                    let a_term = if j < self.a.len() { self.a[j] * y } else { 0.0 };
                    let next = if j + 1 < delay.len() {
                        delay[j + 1]
                    } else {
                        0.0
                    };
                    delay[j] = b_term - a_term + next;
                }

                output_data[i * channels + ch] = y;
            }
        }

        Ok(Block {
            data: output_data,
            sample_rate: block.sample_rate,
            channels: block.channels,
        })
    }
}

/// Gain stage: multiply all samples by a constant.
pub struct GainStage {
    gain: f32,
}

impl GainStage {
    pub fn new(gain: f32) -> Self {
        Self { gain }
    }
}

impl DspStage for GainStage {
    fn process(&mut self, block: &Block) -> Result<Block, FabricError> {
        let data: Vec<Sample> = block.data.iter().map(|&s| s * self.gain).collect();
        Ok(Block {
            data,
            sample_rate: block.sample_rate,
            channels: block.channels,
        })
    }
}

/// Mixer stage: sums multiple input blocks into one.
///
/// Call `add_input` for each input block, then `mix` to produce the sum.
pub struct MixerStage {
    accumulated: Option<Block>,
}

impl MixerStage {
    pub fn new() -> Self {
        Self { accumulated: None }
    }

    /// Add an input block to the mix.
    pub fn add_input(&mut self, block: &Block) -> Result<(), FabricError> {
        match &mut self.accumulated {
            None => {
                self.accumulated = Some(block.clone());
            }
            Some(acc) => {
                if acc.data.len() != block.data.len() {
                    return Err(FabricError::CapacityExceeded);
                }
                for (a, &b) in acc.data.iter_mut().zip(block.data.iter()) {
                    *a += b;
                }
            }
        }
        Ok(())
    }

    /// Produce the mixed output and reset internal state.
    pub fn mix(&mut self) -> Result<Block, FabricError> {
        self.accumulated.take().ok_or(FabricError::CapacityExceeded)
    }
}

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

/// DspStage implementation for MixerStage that passes blocks through
/// (accumulating for later mixing).
impl DspStage for MixerStage {
    fn process(&mut self, block: &Block) -> Result<Block, FabricError> {
        self.add_input(block)?;
        Ok(block.clone())
    }
}

/// Resample stage: sample rate conversion using linear interpolation.
pub struct ResampleStage {
    target_rate: u32,
}

impl ResampleStage {
    pub fn new(target_rate: u32) -> Self {
        Self { target_rate }
    }
}

impl DspStage for ResampleStage {
    fn process(&mut self, block: &Block) -> Result<Block, FabricError> {
        if block.sample_rate == self.target_rate {
            return Ok(block.clone());
        }

        let channels = block.channels as usize;
        let in_frames = block.frames();
        if in_frames == 0 || channels == 0 {
            return Ok(Block {
                data: Vec::new(),
                sample_rate: self.target_rate,
                channels: block.channels,
            });
        }

        let ratio = self.target_rate as f64 / block.sample_rate as f64;
        let out_frames = (in_frames as f64 * ratio).round() as usize;
        let mut output_data = vec![0.0f32; out_frames * channels];

        for ch in 0..channels {
            let input: Vec<f32> = (0..in_frames)
                .map(|i| block.data[i * channels + ch])
                .collect();

            for i in 0..out_frames {
                let src_pos = i as f64 / ratio;
                let src_idx = src_pos as usize;
                let frac = src_pos - src_idx as f64;

                let s0 = input[src_idx.min(in_frames - 1)];
                let s1 = input[(src_idx + 1).min(in_frames - 1)];
                output_data[i * channels + ch] = s0 + (s1 - s0) * frac as f32;
            }
        }

        Ok(Block {
            data: output_data,
            sample_rate: self.target_rate,
            channels: block.channels,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloc::vec;

    fn approx_eq(a: f32, b: f32, tol: f32) -> bool {
        (a - b).abs() < tol
    }

    fn mono_block(data: Vec<f32>, sample_rate: u32) -> Block {
        Block {
            data,
            sample_rate,
            channels: 1,
        }
    }

    fn stereo_block(data: Vec<f32>, sample_rate: u32) -> Block {
        Block {
            data,
            sample_rate,
            channels: 2,
        }
    }

    // --- FftStage ---

    #[test]
    fn fft_stage_power_of_two() {
        let mut stage = FftStage::new();
        let block = mono_block(vec![1.0, 0.0, -1.0, 0.0], 44100);
        let result = stage.process(&block).unwrap();
        // 4 input samples -> 8 output values (interleaved re/im)
        assert_eq!(result.data.len(), 8);
        assert_eq!(result.sample_rate, 44100);
        assert_eq!(result.channels, 1);
    }

    #[test]
    fn fft_stage_rejects_non_power_of_two() {
        let mut stage = FftStage::new();
        let block = mono_block(vec![1.0, 2.0, 3.0], 44100);
        assert!(stage.process(&block).is_err());
    }

    #[test]
    fn fft_stage_rejects_empty() {
        let mut stage = FftStage::new();
        let block = mono_block(vec![], 44100);
        assert!(stage.process(&block).is_err());
    }

    #[test]
    fn fft_stage_dc_signal() {
        let mut stage = FftStage::new();
        // Constant signal: all energy should be in bin 0
        let block = mono_block(vec![1.0; 4], 44100);
        let result = stage.process(&block).unwrap();
        // Bin 0: re should be 4.0 (sum of all samples), im should be ~0
        assert!(approx_eq(result.data[0], 4.0, 1e-4));
        assert!(approx_eq(result.data[1], 0.0, 1e-4));
    }

    #[test]
    fn fft_stage_default() {
        let stage = FftStage::default();
        let mut stage = stage;
        let block = mono_block(vec![1.0, 0.0, 1.0, 0.0], 48000);
        assert!(stage.process(&block).is_ok());
    }

    #[test]
    fn fft_stage_stereo() {
        let mut stage = FftStage::new();
        // 2 channels, 4 frames = 8 interleaved samples
        let block = stereo_block(vec![1.0, 2.0, 0.0, 0.0, -1.0, -2.0, 0.0, 0.0], 44100);
        let result = stage.process(&block).unwrap();
        // Each channel: 4 frames -> 4 complex bins -> 8 floats per channel = 16 total
        assert_eq!(result.data.len(), 16);
        assert_eq!(result.channels, 2);
    }

    // --- IfftStage ---

    #[test]
    fn ifft_stage_rejects_zero_channels() {
        let mut stage = IfftStage::new();
        let block = Block {
            data: vec![1.0, 0.0, 1.0, 0.0],
            sample_rate: 44100,
            channels: 0,
        };
        assert!(stage.process(&block).is_err());
    }

    #[test]
    fn ifft_stage_rejects_odd_count() {
        let mut stage = IfftStage::new();
        // 3 values for 1 channel is not a valid re/im pair count
        let block = mono_block(vec![1.0, 0.0, 1.0], 44100);
        assert!(stage.process(&block).is_err());
    }

    #[test]
    fn ifft_stage_rejects_non_power_of_two_complex_count() {
        let mut stage = IfftStage::new();
        // 6 floats = 3 complex, but 3 is not a power of 2
        let block = mono_block(vec![1.0, 0.0, 1.0, 0.0, 1.0, 0.0], 44100);
        assert!(stage.process(&block).is_err());
    }

    #[test]
    fn ifft_stage_default() {
        let stage = IfftStage::default();
        let mut stage = stage;
        // 4 complex bins = 8 floats
        let block = mono_block(vec![4.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], 44100);
        let result = stage.process(&block).unwrap();
        assert_eq!(result.data.len(), 4);
    }

    #[test]
    fn fft_ifft_roundtrip_stage() {
        let mut fft_stage = FftStage::new();
        let mut ifft_stage = IfftStage::new();
        let original = mono_block(vec![1.0, 2.0, 3.0, 4.0], 48000);
        let freq = fft_stage.process(&original).unwrap();
        let recovered = ifft_stage.process(&freq).unwrap();
        assert_eq!(recovered.data.len(), 4);
        for i in 0..4 {
            assert!(
                approx_eq(recovered.data[i], original.data[i], 1e-4),
                "sample {i}: expected {}, got {}",
                original.data[i],
                recovered.data[i]
            );
        }
    }

    // --- FirFilterStage ---

    #[test]
    fn fir_identity_filter() {
        // Single coefficient [1.0] = identity
        let mut stage = FirFilterStage::new(vec![1.0]);
        let block = mono_block(vec![1.0, 2.0, 3.0, 4.0], 44100);
        let result = stage.process(&block).unwrap();
        for i in 0..4 {
            assert!(
                approx_eq(result.data[i], block.data[i], 1e-6),
                "sample {i}: expected {}, got {}",
                block.data[i],
                result.data[i]
            );
        }
    }

    #[test]
    fn fir_gain_filter() {
        // [2.0] = multiply by 2
        let mut stage = FirFilterStage::new(vec![2.0]);
        let block = mono_block(vec![1.0, 2.0, 3.0], 44100);
        let result = stage.process(&block).unwrap();
        assert!(approx_eq(result.data[0], 2.0, 1e-6));
        assert!(approx_eq(result.data[1], 4.0, 1e-6));
        assert!(approx_eq(result.data[2], 6.0, 1e-6));
    }

    #[test]
    fn fir_moving_average() {
        // 2-tap moving average: [0.5, 0.5]
        let mut stage = FirFilterStage::new(vec![0.5, 0.5]);
        let block = mono_block(vec![0.0, 1.0, 0.0, 1.0], 44100);
        let result = stage.process(&block).unwrap();
        // y[0] = 0.5*0 + 0.5*0 = 0 (overlap is zero)
        assert!(approx_eq(result.data[0], 0.0, 1e-6));
        // y[1] = 0.5*1 + 0.5*0 = 0.5
        assert!(approx_eq(result.data[1], 0.5, 1e-6));
        // y[2] = 0.5*0 + 0.5*1 = 0.5
        assert!(approx_eq(result.data[2], 0.5, 1e-6));
        // y[3] = 0.5*1 + 0.5*0 = 0.5
        assert!(approx_eq(result.data[3], 0.5, 1e-6));
    }

    #[test]
    fn fir_overlap_across_blocks() {
        // Verify state carries across blocks
        let mut stage = FirFilterStage::new(vec![0.5, 0.5]);
        let block1 = mono_block(vec![1.0, 0.0], 44100);
        let block2 = mono_block(vec![0.0, 1.0], 44100);

        let _ = stage.process(&block1).unwrap();
        let result2 = stage.process(&block2).unwrap();
        // First sample of block2: y[0] = 0.5*0 + 0.5*0 = 0.0 (overlap from block1 tail)
        // The overlap from block1 is [0.0] (last sample), so:
        // y[0] = 0.5*0.0 + 0.5*0.0 = 0.0
        assert!(approx_eq(result2.data[0], 0.0, 1e-6));
    }

    #[test]
    fn fir_stereo() {
        let mut stage = FirFilterStage::new(vec![1.0]);
        // 2 channels, 2 frames: [L0, R0, L1, R1]
        let block = stereo_block(vec![1.0, 2.0, 3.0, 4.0], 44100);
        let result = stage.process(&block).unwrap();
        assert_eq!(result.data.len(), 4);
        assert!(approx_eq(result.data[0], 1.0, 1e-6)); // L0
        assert!(approx_eq(result.data[1], 2.0, 1e-6)); // R0
        assert!(approx_eq(result.data[2], 3.0, 1e-6)); // L1
        assert!(approx_eq(result.data[3], 4.0, 1e-6)); // R1
    }

    // --- IirFilterStage ---

    #[test]
    fn iir_passthrough() {
        // b=[1.0], a=[] => y = x (passthrough)
        let mut stage = IirFilterStage::new(vec![1.0], vec![]);
        let block = mono_block(vec![1.0, 2.0, 3.0, 4.0], 44100);
        let result = stage.process(&block).unwrap();
        for i in 0..4 {
            assert!(
                approx_eq(result.data[i], block.data[i], 1e-6),
                "sample {i}: expected {}, got {}",
                block.data[i],
                result.data[i]
            );
        }
    }

    #[test]
    fn iir_gain() {
        // b=[2.0], a=[] => y = 2*x
        let mut stage = IirFilterStage::new(vec![2.0], vec![]);
        let block = mono_block(vec![1.0, 0.5, 0.25], 44100);
        let result = stage.process(&block).unwrap();
        assert!(approx_eq(result.data[0], 2.0, 1e-6));
        assert!(approx_eq(result.data[1], 1.0, 1e-6));
        assert!(approx_eq(result.data[2], 0.5, 1e-6));
    }

    #[test]
    fn iir_first_order_feedback() {
        // b=[1.0], a=[−0.5] => y[n] = x[n] + 0.5*y[n−1]
        let mut stage = IirFilterStage::new(vec![1.0], vec![-0.5]);
        let block = mono_block(vec![1.0, 0.0, 0.0, 0.0], 44100);
        let result = stage.process(&block).unwrap();
        // y[0] = 1.0 + 0 = 1.0
        assert!(approx_eq(result.data[0], 1.0, 1e-6));
        // y[1] = 0.0 + 0.5*1.0 = 0.5
        assert!(approx_eq(result.data[1], 0.5, 1e-6));
        // y[2] = 0.0 + 0.5*0.5 = 0.25
        assert!(approx_eq(result.data[2], 0.25, 1e-6));
        // y[3] = 0.0 + 0.5*0.25 = 0.125
        assert!(approx_eq(result.data[3], 0.125, 1e-6));
    }

    #[test]
    fn iir_stereo() {
        // b=[1.0], a=[] passthrough on stereo
        let mut stage = IirFilterStage::new(vec![1.0], vec![]);
        let block = stereo_block(vec![1.0, 10.0, 2.0, 20.0], 44100);
        let result = stage.process(&block).unwrap();
        assert!(approx_eq(result.data[0], 1.0, 1e-6));
        assert!(approx_eq(result.data[1], 10.0, 1e-6));
        assert!(approx_eq(result.data[2], 2.0, 1e-6));
        assert!(approx_eq(result.data[3], 20.0, 1e-6));
    }

    #[test]
    fn iir_state_across_blocks() {
        // b=[1.0], a=[-0.5] with impulse then zeros across two blocks
        let mut stage = IirFilterStage::new(vec![1.0], vec![-0.5]);
        let block1 = mono_block(vec![1.0, 0.0], 44100);
        let result1 = stage.process(&block1).unwrap();
        assert!(approx_eq(result1.data[0], 1.0, 1e-6));
        assert!(approx_eq(result1.data[1], 0.5, 1e-6));

        // Second block continues the decay
        let block2 = mono_block(vec![0.0, 0.0], 44100);
        let result2 = stage.process(&block2).unwrap();
        assert!(approx_eq(result2.data[0], 0.25, 1e-6));
        assert!(approx_eq(result2.data[1], 0.125, 1e-6));
    }

    // --- GainStage ---

    #[test]
    fn gain_stage_unity() {
        let mut stage = GainStage::new(1.0);
        let block = mono_block(vec![1.0, 2.0, 3.0], 44100);
        let result = stage.process(&block).unwrap();
        assert_eq!(result.data, block.data);
    }

    #[test]
    fn gain_stage_amplify() {
        let mut stage = GainStage::new(3.0);
        let block = mono_block(vec![1.0, -1.0, 0.5], 44100);
        let result = stage.process(&block).unwrap();
        assert!(approx_eq(result.data[0], 3.0, 1e-6));
        assert!(approx_eq(result.data[1], -3.0, 1e-6));
        assert!(approx_eq(result.data[2], 1.5, 1e-6));
    }

    #[test]
    fn gain_stage_attenuate() {
        let mut stage = GainStage::new(0.5);
        let block = mono_block(vec![4.0, 2.0], 44100);
        let result = stage.process(&block).unwrap();
        assert!(approx_eq(result.data[0], 2.0, 1e-6));
        assert!(approx_eq(result.data[1], 1.0, 1e-6));
    }

    #[test]
    fn gain_stage_zero() {
        let mut stage = GainStage::new(0.0);
        let block = mono_block(vec![100.0, -50.0], 44100);
        let result = stage.process(&block).unwrap();
        assert!(approx_eq(result.data[0], 0.0, 1e-6));
        assert!(approx_eq(result.data[1], 0.0, 1e-6));
    }

    #[test]
    fn gain_stage_negative() {
        let mut stage = GainStage::new(-1.0);
        let block = mono_block(vec![1.0, -2.0], 44100);
        let result = stage.process(&block).unwrap();
        assert!(approx_eq(result.data[0], -1.0, 1e-6));
        assert!(approx_eq(result.data[1], 2.0, 1e-6));
    }

    #[test]
    fn gain_stage_preserves_metadata() {
        let mut stage = GainStage::new(2.0);
        let block = Block {
            data: vec![1.0],
            sample_rate: 96000,
            channels: 3,
        };
        let result = stage.process(&block).unwrap();
        assert_eq!(result.sample_rate, 96000);
        assert_eq!(result.channels, 3);
    }

    // --- MixerStage ---

    #[test]
    fn mixer_single_input() {
        let mut mixer = MixerStage::new();
        let block = mono_block(vec![1.0, 2.0, 3.0], 44100);
        mixer.add_input(&block).unwrap();
        let result = mixer.mix().unwrap();
        assert_eq!(result.data, block.data);
    }

    #[test]
    fn mixer_two_inputs() {
        let mut mixer = MixerStage::new();
        let a = mono_block(vec![1.0, 2.0, 3.0], 44100);
        let b = mono_block(vec![10.0, 20.0, 30.0], 44100);
        mixer.add_input(&a).unwrap();
        mixer.add_input(&b).unwrap();
        let result = mixer.mix().unwrap();
        assert!(approx_eq(result.data[0], 11.0, 1e-6));
        assert!(approx_eq(result.data[1], 22.0, 1e-6));
        assert!(approx_eq(result.data[2], 33.0, 1e-6));
    }

    #[test]
    fn mixer_mismatched_lengths() {
        let mut mixer = MixerStage::new();
        let a = mono_block(vec![1.0, 2.0], 44100);
        let b = mono_block(vec![1.0, 2.0, 3.0], 44100);
        mixer.add_input(&a).unwrap();
        assert!(mixer.add_input(&b).is_err());
    }

    #[test]
    fn mixer_empty_mix() {
        let mut mixer = MixerStage::new();
        assert!(mixer.mix().is_err());
    }

    #[test]
    fn mixer_resets_after_mix() {
        let mut mixer = MixerStage::new();
        let block = mono_block(vec![1.0], 44100);
        mixer.add_input(&block).unwrap();
        let _ = mixer.mix().unwrap();
        // After mix(), internal state is cleared
        assert!(mixer.mix().is_err());
    }

    #[test]
    fn mixer_default() {
        let mut mixer = MixerStage::default();
        let block = mono_block(vec![5.0], 44100);
        mixer.add_input(&block).unwrap();
        let result = mixer.mix().unwrap();
        assert!(approx_eq(result.data[0], 5.0, 1e-6));
    }

    #[test]
    fn mixer_dsp_stage_passes_through() {
        let mut mixer = MixerStage::new();
        let block = mono_block(vec![1.0, 2.0], 44100);
        // DspStage::process accumulates and returns a clone
        let result = mixer.process(&block).unwrap();
        assert_eq!(result.data, block.data);
        // The accumulated value is available via mix()
        let mixed = mixer.mix().unwrap();
        assert_eq!(mixed.data, block.data);
    }

    // --- ResampleStage ---

    #[test]
    fn resample_same_rate() {
        let mut stage = ResampleStage::new(44100);
        let block = mono_block(vec![1.0, 2.0, 3.0, 4.0], 44100);
        let result = stage.process(&block).unwrap();
        assert_eq!(result.data, block.data);
        assert_eq!(result.sample_rate, 44100);
    }

    #[test]
    fn resample_upsample_2x() {
        let mut stage = ResampleStage::new(96000);
        let block = mono_block(vec![0.0, 1.0], 48000);
        let result = stage.process(&block).unwrap();
        assert_eq!(result.sample_rate, 96000);
        // 2 frames at 48k -> 4 frames at 96k
        assert_eq!(result.data.len(), 4);
        // First sample should be ~0.0
        assert!(approx_eq(result.data[0], 0.0, 1e-4));
        // Last sample should be ~1.0
        assert!(approx_eq(result.data[3], 1.0, 1e-4));
    }

    #[test]
    fn resample_downsample_2x() {
        let mut stage = ResampleStage::new(22050);
        let block = mono_block(vec![0.0, 0.25, 0.5, 0.75], 44100);
        let result = stage.process(&block).unwrap();
        assert_eq!(result.sample_rate, 22050);
        // 4 frames at 44100 -> 2 frames at 22050
        assert_eq!(result.data.len(), 2);
    }

    #[test]
    fn resample_empty_block() {
        let mut stage = ResampleStage::new(96000);
        let block = mono_block(vec![], 44100);
        let result = stage.process(&block).unwrap();
        assert!(result.data.is_empty());
        assert_eq!(result.sample_rate, 96000);
    }

    #[test]
    fn resample_stereo() {
        let mut stage = ResampleStage::new(96000);
        // 2 channels, 2 frames at 48k
        let block = stereo_block(vec![0.0, 10.0, 1.0, 20.0], 48000);
        let result = stage.process(&block).unwrap();
        assert_eq!(result.sample_rate, 96000);
        assert_eq!(result.channels, 2);
        // 2 frames -> 4 frames, 2 channels = 8 samples
        assert_eq!(result.data.len(), 8);
    }

    #[test]
    fn resample_preserves_dc() {
        // A constant signal should remain constant after resampling
        let mut stage = ResampleStage::new(96000);
        let block = mono_block(vec![5.0; 4], 48000);
        let result = stage.process(&block).unwrap();
        for &sample in &result.data {
            assert!(approx_eq(sample, 5.0, 1e-4));
        }
    }
}