//! An audio sample rate conversion library for Rust.
//!
//! This library provides resamplers to process audio in chunks.
//! The ratio between input and output sample rates is completely free.
//! Implementations are available that accept a fixed length input
//! while returning a variable length output, and vice versa.
//! The resampling is based on band-limited interpolation using sinc
//! interpolation filters. The sinc interpolation upsamples by an adjustable factor,
//! and then the new sample points are calculated by interpolating between these points.
//!
//! ## Documentation
//!
//! The full documentation can be generated by rustdoc. To generate and view it run:
//! ```text
//! cargo doc --open
//! ```
//!
//! ## Example
//! Resample a single chunk of a dummy audio file from 44100 to 48000 Hz.
//! See also the "fixedin64" example that can be used to process a file from disk.
//! ```
//! use rubato::{Resampler, SincFixedIn, InterpolationType, InterpolationParameters, WindowFunction};
//! let params = InterpolationParameters {
//!     sinc_len: 256,
//!     f_cutoff: 0.95,
//!     interpolation: InterpolationType::Nearest,
//!     oversampling_factor: 160,
//!     window: WindowFunction::BlackmanHarris2,
//! };
//! let mut resampler = SincFixedIn::<f64>::new(
//!     48000 as f64 / 44100 as f64,
//!     params,
//!     1024,
//!     2,
//! );
//!
//! let waves_in = vec![vec![0.0f64; 1024];2];
//! let waves_out = resampler.process(&waves_in).unwrap();
//! ```
//!
//! ## Compatibility
//!
//! The `rubato` crate only depends on the `num` crate and should work with any rustc version that crate supports.

mod helpers;
mod interpolation;

use crate::helpers::*;
use crate::interpolation::*;
use num::traits::Float;
use std::error;
use std::fmt;

#[macro_use]
extern crate log;

type Res<T> = Result<T, Box<dyn error::Error>>;

/// Custom error returned by resamplers
#[derive(Debug)]
pub struct ResamplerError {
    desc: String,
}

impl fmt::Display for ResamplerError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self.desc)
    }
}

impl error::Error for ResamplerError {
    fn description(&self) -> &str {
        &self.desc
    }
}

impl ResamplerError {
    pub fn new(desc: &str) -> Self {
        ResamplerError {
            desc: desc.to_owned(),
        }
    }
}

/// Different window functions that can be used to window the sinc function.
#[derive(Debug)]
pub enum WindowFunction {
    /// Blackman. Intermediate rolloff and intermediate attenuation.
    Blackman,
    /// Squared Blackman. Slower rolloff but better attenuation than Blackman.
    Blackman2,
    /// Blackman-Harris. Slow rolloff but good attenuation.
    BlackmanHarris,
    /// Squared Blackman-Harris. Slower rolloff but better attenuation than Blackman-Harris.
    BlackmanHarris2,
    /// Hann, fast rolloff but not very high attenuation
    Hann,
    /// Squared Hann, slower rolloff and higher attenuation than simple Hann
    Hann2,
}

/// A struct holding the parameters for interpolation.
#[derive(Debug)]
pub struct InterpolationParameters {
    /// Length of the windowed sinc interpolation filter.
    /// Higher values can allow a higher cut-off frequency leading to less high frequency roll-off
    /// at the expense of higher cpu usage. 256 is a good starting point.
    pub sinc_len: usize,
    /// Relative cutoff frequency of the sinc interpolation filter
    /// (relative to the lowest one of fs_in/2 or fs_out/2). Start at 0.95, and increase if needed.
    pub f_cutoff: f32,
    /// The number of intermediate points go use for interpolation.
    /// Higher values use more memory for storing the sinc filters.
    /// Only the points actually needed are calculated dusing processing
    /// so a larger number does not directly lead to higher cpu usage.
    /// But keeping it down helps in keeping the sincs in the cpu cache. Start at 128.
    pub oversampling_factor: usize,
    /// Interpolation type, see `InterpolationType`
    pub interpolation: InterpolationType,
    /// Window function to use.
    pub window: WindowFunction,
}

/// Interpolation methods that can be selected. For asynchronous interpolation where the
/// ratio between inut and output sample rates can be any number, it's not possible to
/// pre-calculate all the needed interpolation filters.
/// Instead they have to be computed as needed, which becomes impractical since the
/// sincs are very expensive to generate in terms of cpu time.
/// It's more efficient to combine the sinc filters with some other interpolation technique.
/// Then sinc filters are used to provide a fixed number of interpolated points between input samples,
/// and then the new value is calculated by interpolation between those points.

#[derive(Debug)]
pub enum InterpolationType {
    /// For cubic interpolation, the four nearest intermediate points are calculated
    /// using sinc interpolation.
    /// Then a cubic polynomial is fitted to these points, and is then used to calculate the new sample value.
    /// The computation time as about twice the one for linear interpolation,
    /// but it requires much fewer intermediate points for a good result.
    Cubic,
    /// With linear interpolation the new sample value is calculated by linear interpolation
    /// between the two nearest points.
    /// This requires two intermediate points to be calcuated using sinc interpolation,
    /// and te output is a weighted average of these two.
    /// This is relatively fast, but needs a large number of intermediate points to
    /// push the resampling artefacts below the noise floor.
    Linear,
    /// The Nearest mode doesn't do any interpolation, but simply picks the nearest intermediate point.
    /// This is useful when the nearest point is actually the correct one, for example when upsampling by a factor 2,
    /// like 48kHz->96kHz.
    /// Then setting the oversampling_factor to 2, and using Nearest mode,
    /// no unneccesary computations are performed and the result is the same as for synchronous resampling.
    /// This also works for other ratios that can be expressed by a fraction. For 44.1kHz -> 48 kHz,
    /// setting oversampling_factor to 160 gives the desired result (since 48kHz = 160/147 * 44.1kHz).
    Nearest,
}

/// A resampler that accepts a fixed number of audio chunks for input
/// and returns a variable number of frames.
///
/// The resampling is done by creating a number of intermediate points (defined by oversampling_factor)
/// by sinc interpolation. The new samples are then calculated by interpolating between these points.
pub struct SincFixedIn<T: Float> {
    nbr_channels: usize,
    chunk_size: usize,
    oversampling_factor: usize,
    last_index: f64,
    resample_ratio: f64,
    resample_ratio_original: f64,
    sinc_len: usize,
    sincs: Vec<Vec<T>>,
    buffer: Vec<Vec<T>>,
    interpolation: InterpolationType,
}

/// A resampler that return a fixed number of audio chunks.
/// The number of input frames required is given by the frames_needed function.
///
/// The resampling is done by creating a number of intermediate points (defined by oversampling_factor)
/// by sinc interpolation. The new samples are then calculated by interpolating between these points.
pub struct SincFixedOut<T: Float> {
    nbr_channels: usize,
    chunk_size: usize,
    needed_input_size: usize,
    oversampling_factor: usize,
    last_index: f64,
    current_buffer_fill: usize,
    resample_ratio: f64,
    resample_ratio_original: f64,
    sinc_len: usize,
    sincs: Vec<Vec<T>>,
    buffer: Vec<Vec<T>>,
    interpolation: InterpolationType,
}

/// A resampler that us used to resample a chunk of audio to a new sample rate.
/// The rate can be adjusted as required.
pub trait Resampler<T: Float> {
    /// Resample a chunk of audio. Input and output data is stored in a vector,
    /// where each element contains a vector with all samples for a single channel.
    fn process(&mut self, wave_in: &[Vec<T>]) -> Res<Vec<Vec<T>>>;

    /// Update the resample ratio. New value must be within +-10% of the original one.
    fn set_resample_ratio(&mut self, new_ratio: f64) -> Res<()>;

    /// Update the resample ratio relative to the original one. Must be in the range 0.9 - 1.1.
    fn set_resample_ratio_relative(&mut self, rel_ratio: f64) -> Res<()>;

    /// Query for the number of frames needed for the next call to "process".
    fn nbr_frames_needed(&self) -> usize;
}

impl<T: Float> Resampler<T> for SincFixedIn<T> {
    /// Resample a chunk of audio. The input length is fixed, and the output varies in length.
    /// # Errors
    ///
    /// The function returns an error if the length of the input data is not equal
    /// to the number of channels and chunk size defined when creating the instance.
    fn process(&mut self, wave_in: &[Vec<T>]) -> Res<Vec<Vec<T>>> {
        if wave_in.len() != self.nbr_channels {
            return Err(Box::new(ResamplerError::new(
                "Wrong number of channels in input",
            )));
        }
        if wave_in[0].len() != self.chunk_size {
            return Err(Box::new(ResamplerError::new(
                "Wrong number of frames in input",
            )));
        }
        let end_idx = self.chunk_size as isize - (self.sinc_len as isize + 1);
        //update buffer with new data
        for wav in self.buffer.iter_mut() {
            for idx in 0..(2 * self.sinc_len) {
                wav[idx] = wav[idx + self.chunk_size];
            }
        }
        for (chan, wav) in wave_in.iter().enumerate() {
            for (idx, sample) in wav.iter().enumerate() {
                self.buffer[chan][idx + 2 * self.sinc_len] = *sample;
            }
        }

        let mut idx = self.last_index;
        let t_ratio = 1.0 / self.resample_ratio as f64;

        let mut wave_out =
            vec![
                vec![T::zero(); (self.chunk_size as f64 * self.resample_ratio + 10.0) as usize];
                self.nbr_channels
            ];
        let mut n = 0;

        match self.interpolation {
            InterpolationType::Cubic => {
                let mut points = vec![T::zero(); 4];
                let mut nearest = vec![(0isize, 0isize); 4];
                while idx < end_idx as f64 {
                    idx += t_ratio;
                    get_nearest_times_4(idx, self.oversampling_factor as isize, &mut nearest);
                    let frac = idx * self.oversampling_factor as f64
                        - (idx * self.oversampling_factor as f64).floor();
                    let frac_offset = T::from(frac).unwrap();
                    for (chan, buf) in self.buffer.iter().enumerate() {
                        for (n, p) in nearest.iter().zip(points.iter_mut()) {
                            *p = get_sinc_interpolated(
                                &buf,
                                &self.sincs,
                                (n.0 + 2 * self.sinc_len as isize) as usize,
                                n.1 as usize,
                            );
                        }
                        wave_out[chan][n] = interp_cubic(frac_offset, &points);
                    }
                    n += 1;
                }
            }
            InterpolationType::Linear => {
                let mut points = vec![T::zero(); 2];
                let mut nearest = vec![(0isize, 0isize); 2];
                while idx < end_idx as f64 {
                    idx += t_ratio;
                    get_nearest_times_2(idx, self.oversampling_factor as isize, &mut nearest);
                    let frac = idx * self.oversampling_factor as f64
                        - (idx * self.oversampling_factor as f64).floor();
                    let frac_offset = T::from(frac).unwrap();
                    for (chan, buf) in self.buffer.iter().enumerate() {
                        for (n, p) in nearest.iter().zip(points.iter_mut()) {
                            *p = get_sinc_interpolated(
                                &buf,
                                &self.sincs,
                                (n.0 + 2 * self.sinc_len as isize) as usize,
                                n.1 as usize,
                            );
                        }
                        wave_out[chan][n] = interp_lin(frac_offset, &points);
                    }
                    n += 1;
                }
            }
            InterpolationType::Nearest => {
                let mut point;
                let mut nearest;
                while idx < end_idx as f64 {
                    idx += t_ratio;
                    nearest = get_nearest_time(idx, self.oversampling_factor as isize);
                    for (chan, buf) in self.buffer.iter().enumerate() {
                        point = get_sinc_interpolated(
                            &buf,
                            &self.sincs,
                            (nearest.0 + 2 * self.sinc_len as isize) as usize,
                            nearest.1 as usize,
                        );
                        wave_out[chan][n] = point;
                    }
                    n += 1;
                }
            }
        }

        // store last index for next iteration
        self.last_index = idx - self.chunk_size as f64;
        for w in wave_out.iter_mut() {
            w.truncate(n);
        }
        trace!(
            "Resampling, {} frames in, {} frames out",
            wave_in[0].len(),
            wave_out[0].len()
        );
        Ok(wave_out)
    }

    /// Update the resample ratio. New value must be within +-10% of the original one
    fn set_resample_ratio(&mut self, new_ratio: f64) -> Res<()> {
        trace!("Change resample ratio to {}", new_ratio);
        if (new_ratio / self.resample_ratio_original > 0.9)
            && (new_ratio / self.resample_ratio_original < 1.1)
        {
            self.resample_ratio = new_ratio;
            Ok(())
        } else {
            Err(Box::new(ResamplerError::new(
                "New resample ratio is too far off from original",
            )))
        }
    }
    /// Update the resample ratio relative to the original one
    fn set_resample_ratio_relative(&mut self, rel_ratio: f64) -> Res<()> {
        let new_ratio = self.resample_ratio_original * rel_ratio;
        self.set_resample_ratio(new_ratio)
    }

    /// Query for the number of frames needed for the next call to "process".
    /// Will always return the chunk_size defined when creating the instance.
    fn nbr_frames_needed(&self) -> usize {
        self.chunk_size
    }
}

impl<T: Float> SincFixedIn<T> {
    /// Create a new SincFixedIn
    ///
    /// Parameters are:
    /// - `resample_ratio`: Ratio between output and input sample rates.
    /// - `parameters`: Parameters for interpolation, see `InterpolationParameters`
    /// - `chunk_size`: size of input data in frames
    /// - `nbr_channels`: number of channels in input/output
    pub fn new(
        resample_ratio: f64,
        parameters: InterpolationParameters,
        chunk_size: usize,
        nbr_channels: usize,
    ) -> Self {
        debug!(
            "Create new SincFixedIn, ratio: {}, chunk_size: {}, channels: {}, parameters: {:?}",
            resample_ratio, chunk_size, nbr_channels, parameters
        );
        let sinc_cutoff = if resample_ratio >= 1.0 {
            parameters.f_cutoff
        } else {
            parameters.f_cutoff * resample_ratio as f32
        };
        let sincs = make_sincs(
            parameters.sinc_len,
            parameters.oversampling_factor,
            sinc_cutoff,
            parameters.window,
        );
        let buffer = vec![vec![T::zero(); chunk_size + 2 * parameters.sinc_len]; nbr_channels];
        SincFixedIn {
            nbr_channels,
            chunk_size,
            oversampling_factor: parameters.oversampling_factor,
            last_index: -((parameters.sinc_len / 2) as f64),
            resample_ratio,
            resample_ratio_original: resample_ratio,
            sinc_len: parameters.sinc_len,
            sincs,
            buffer,
            interpolation: parameters.interpolation,
        }
    }
}

impl<T: Float> SincFixedOut<T> {
    /// Create a new SincFixedOut
    ///
    /// Parameters are:
    /// - `resample_ratio`: Ratio between output and input sample rates.
    /// - `parameters`: Parameters for interpolation, see `InterpolationParameters`
    /// - `chunk_size`: size of output data in frames
    /// - `nbr_channels`: number of channels in input/output
    pub fn new(
        resample_ratio: f64,
        parameters: InterpolationParameters,
        chunk_size: usize,
        nbr_channels: usize,
    ) -> Self {
        debug!(
            "Create new SincFixedOut, ratio: {}, chunk_size: {}, channels: {}, parameters: {:?}",
            resample_ratio, chunk_size, nbr_channels, parameters
        );
        let sinc_cutoff = if resample_ratio >= 1.0 {
            parameters.f_cutoff
        } else {
            parameters.f_cutoff * resample_ratio as f32
        };
        let sincs = make_sincs(
            parameters.sinc_len,
            parameters.oversampling_factor,
            sinc_cutoff,
            parameters.window,
        );
        let needed_input_size =
            (chunk_size as f64 / resample_ratio).ceil() as usize + 2 + parameters.sinc_len / 2;
        let buffer = vec![
            vec![T::zero(); 3 * needed_input_size / 2 + 2 * parameters.sinc_len];
            nbr_channels
        ];
        SincFixedOut {
            nbr_channels,
            chunk_size,
            needed_input_size,
            oversampling_factor: parameters.oversampling_factor,
            last_index: -((parameters.sinc_len / 2) as f64),
            current_buffer_fill: needed_input_size,
            resample_ratio,
            resample_ratio_original: resample_ratio,
            sinc_len: parameters.sinc_len,
            sincs,
            buffer,
            interpolation: parameters.interpolation,
        }
    }
}

impl<T: Float> Resampler<T> for SincFixedOut<T> {
    /// Query for the number of frames needed for the next call to "process".
    fn nbr_frames_needed(&self) -> usize {
        self.needed_input_size
    }

    /// Update the resample ratio. New value must be within +-10% of the original one
    fn set_resample_ratio(&mut self, new_ratio: f64) -> Res<()> {
        trace!("Change resample ratio to {}", new_ratio);
        if (new_ratio / self.resample_ratio_original > 0.9)
            && (new_ratio / self.resample_ratio_original < 1.1)
        {
            self.resample_ratio = new_ratio;
            self.needed_input_size = (self.last_index as f32
                + self.chunk_size as f32 / self.resample_ratio as f32
                + self.sinc_len as f32)
                .ceil() as usize
                + 2;
            Ok(())
        } else {
            Err(Box::new(ResamplerError::new(
                "New resample ratio is too far off from original",
            )))
        }
    }

    /// Update the resample ratio relative to the original one
    fn set_resample_ratio_relative(&mut self, rel_ratio: f64) -> Res<()> {
        let new_ratio = self.resample_ratio_original * rel_ratio;
        self.set_resample_ratio(new_ratio)
    }

    /// Resample a chunk of audio. The required input length is provided by
    /// the "nbr_frames_required" function, and the output length is fixed.
    /// # Errors
    ///
    /// The function returns an error if the length of the input data is not
    /// equal to the number of channels defined when creating the instance,
    /// and the number of audio frames given by "nbr_frames"required".
    fn process(&mut self, wave_in: &[Vec<T>]) -> Res<Vec<Vec<T>>> {
        //update buffer with new data
        if wave_in.len() != self.nbr_channels {
            return Err(Box::new(ResamplerError::new(
                "Wrong number of channels in input",
            )));
        }
        if wave_in[0].len() != self.needed_input_size {
            return Err(Box::new(ResamplerError::new(
                "Wrong number of frames in input",
            )));
        }
        for wav in self.buffer.iter_mut() {
            for idx in 0..(2 * self.sinc_len) {
                wav[idx] = wav[idx + self.current_buffer_fill];
            }
        }
        self.current_buffer_fill = wave_in[0].len();
        for (chan, wav) in wave_in.iter().enumerate() {
            for (idx, sample) in wav.iter().enumerate() {
                self.buffer[chan][idx + 2 * self.sinc_len] = *sample;
            }
        }

        let mut idx = self.last_index;
        let t_ratio = 1.0 / self.resample_ratio as f64;

        let mut wave_out = vec![vec![T::zero(); self.chunk_size]; self.nbr_channels];

        match self.interpolation {
            InterpolationType::Cubic => {
                let mut points = vec![T::zero(); 4];
                let mut nearest = vec![(0isize, 0isize); 4];
                for n in 0..self.chunk_size {
                    idx += t_ratio;
                    get_nearest_times_4(idx, self.oversampling_factor as isize, &mut nearest);
                    let frac = idx * self.oversampling_factor as f64
                        - (idx * self.oversampling_factor as f64).floor();
                    let frac_offset = T::from(frac).unwrap();
                    for (chan, buf) in self.buffer.iter().enumerate() {
                        for (n, p) in nearest.iter().zip(points.iter_mut()) {
                            *p = get_sinc_interpolated(
                                &buf,
                                &self.sincs,
                                (n.0 + 2 * self.sinc_len as isize) as usize,
                                n.1 as usize,
                            );
                        }
                        wave_out[chan][n] = interp_cubic(frac_offset, &points);
                    }
                }
            }
            InterpolationType::Linear => {
                let mut points = vec![T::zero(); 2];
                let mut nearest = vec![(0isize, 0isize); 2];
                for n in 0..self.chunk_size {
                    idx += t_ratio;
                    get_nearest_times_2(idx, self.oversampling_factor as isize, &mut nearest);
                    let frac = idx * self.oversampling_factor as f64
                        - (idx * self.oversampling_factor as f64).floor();
                    let frac_offset = T::from(frac).unwrap();
                    for (chan, buf) in self.buffer.iter().enumerate() {
                        for (n, p) in nearest.iter().zip(points.iter_mut()) {
                            *p = get_sinc_interpolated(
                                &buf,
                                &self.sincs,
                                (n.0 + 2 * self.sinc_len as isize) as usize,
                                n.1 as usize,
                            );
                        }
                        wave_out[chan][n] = interp_lin(frac_offset, &points);
                    }
                }
            }
            InterpolationType::Nearest => {
                let mut point;
                let mut nearest;
                for n in 0..self.chunk_size {
                    idx += t_ratio;
                    nearest = get_nearest_time(idx, self.oversampling_factor as isize);
                    for (chan, buf) in self.buffer.iter().enumerate() {
                        point = get_sinc_interpolated(
                            &buf,
                            &self.sincs,
                            (nearest.0 + 2 * self.sinc_len as isize) as usize,
                            nearest.1 as usize,
                        );
                        wave_out[chan][n] = point;
                    }
                }
            }
        }

        // store last index for next iteration
        //trace!("idx {}, fill{}", idx, self.current_buffer_fill);
        self.last_index = idx - self.current_buffer_fill as f64;
        //let next_last_index = self.last_index as f64 + self.chunk_size as f64 / self.resample_ratio as f64 + self.sinc_len as f64;
        //let needed_with_margin = next_last_index + (self.sinc_len) as f64;
        self.needed_input_size = (self.last_index as f32
            + self.chunk_size as f32 / self.resample_ratio as f32
            + self.sinc_len as f32)
            .ceil() as usize
            + 2;
        //self.needed_input_size = ((self.chunk_size as f32 + self.last_index as f32 + (self.sinc_len) as f32)/ self.resample_ratio).ceil() as usize + 2;
        //self.needed_input_size = (self.needed_input_size as isize
        //    + self.last_index.round() as isize
        //    + self.sinc_len as isize) as usize + 2;
        trace!(
            "Resampling, {} frames in, {} frames out. Next needed length: {} frames, last index {}",
            wave_in[0].len(),
            wave_out[0].len(),
            self.needed_input_size,
            self.last_index
        );
        Ok(wave_out)
    }
}

#[cfg(test)]
mod tests {
    use crate::InterpolationParameters;
    use crate::InterpolationType;
    use crate::Resampler;
    use crate::WindowFunction;
    use crate::{SincFixedIn, SincFixedOut};

    #[test]
    fn make_resampler_fi() {
        let params = InterpolationParameters {
            sinc_len: 64,
            f_cutoff: 0.95,
            interpolation: InterpolationType::Cubic,
            oversampling_factor: 16,
            window: WindowFunction::BlackmanHarris2,
        };
        let mut resampler = SincFixedIn::<f64>::new(1.2, params, 1024, 2);
        let waves = vec![vec![0.0f64; 1024]; 2];
        let out = resampler.process(&waves).unwrap();
        assert_eq!(out.len(), 2);
        assert!(out[0].len() > 1150 && out[0].len() < 1250);
    }

    #[test]
    fn make_resampler_fo() {
        let params = InterpolationParameters {
            sinc_len: 64,
            f_cutoff: 0.95,
            interpolation: InterpolationType::Cubic,
            oversampling_factor: 16,
            window: WindowFunction::BlackmanHarris2,
        };
        let mut resampler = SincFixedOut::<f64>::new(1.2, params, 1024, 2);
        let frames = resampler.nbr_frames_needed();
        println!("{}", frames);
        assert!(frames > 800 && frames < 900);
        let waves = vec![vec![0.0f64; frames]; 2];
        let out = resampler.process(&waves).unwrap();
        assert_eq!(out.len(), 2);
        assert_eq!(out[0].len(), 1024);
    }
}