butterworth 0.1.0

use num_complex::{Complex, ComplexFloat};
use num_traits::identities::{Zero, One};
use crate::{filter::Cutoff};
use crate::error::{FilterError, check_sample_rate, check_cutoff};

/// Create a Butterworth filter transfer function from a given order, sample rate, and cutoff
/// frequency. The zeros, poles, and gain (zpk) form creation is based on scipy.signal.butter.
pub(crate) fn transfer_from_frequency(order: usize, sample_rate: f64, filter_type: Cutoff) ->
                                    Result<(Vec<f64>, Vec<f64>), FilterError> {
    check_sample_rate(sample_rate)?;
    match filter_type {
        Cutoff::LowPass(f) | Cutoff::HighPass(f) => {
            check_cutoff(f, sample_rate)?;
        }
        Cutoff::BandPass(f1, f2) | Cutoff::BandStop(f1, f2) => {
            check_cutoff(f1, sample_rate)?;
            check_cutoff(f2, sample_rate)?;
        }
    }
    let pi = std::f64::consts::PI;
    let m = (-(order as isize) + 1..order as isize).map(|e| e as f64)
        .collect::<Vec<f64>>().iter().enumerate().filter(|&(i, _)| i % 2 == 0)
        .map(|(_, e)| *e).collect::<Vec<f64>>();
    let mut p: Vec<Complex<f64>> = m.iter()
        .map(|m| -(Complex::i() * pi * m / (2. * order as f64)).exp()).collect();
    let mut z: Vec<Complex<f64>>;
    let mut k: f64;

    match filter_type {
        Cutoff::LowPass(mut f) => {
            f = 2. * f / sample_rate;
            let warped = 4. * (pi * f / 2.).tan();
            z = vec![];
            p = p.iter().map(|e| e * warped).collect();
            k = warped.powi(order as i32);
        }
        Cutoff::HighPass(mut f) => {
            f = 2. * f / sample_rate;
            let warped = 4. * (pi * f / 2.).tan();
            z = vec![Complex::from(0.); order];
            k = -(1. / p.iter().copied().reduce(|a, b| -a * -b)
                .expect("Could not take product of p vector")).re();
            p = p.iter().map(|e| warped / e).collect();
        }
        Cutoff::BandPass(mut f1, mut f2) => {
            f1 = 2. * f1 / sample_rate;
            let warped1 = 4. * (pi * f1 / 2.).tan();
            f2 = 2. * f2 / sample_rate;
            let warped2 = 4. * (pi * f2 / 2.).tan();
            let band_width = warped1.max(warped2) - warped1.min(warped2);
            let band_center = (warped1 * warped2).sqrt();

            z = vec![Complex::from(0.); order];
            p = p.iter().map(|e| e * band_width / 2.).collect();
            let p2 = p.clone().iter().map(|e| *e - (e.powi(2) - band_center
                .powi(2)).sqrt()).collect::<Vec<Complex<f64>>>();
            p = p.iter().map(|e| *e + (e.powi(2) - band_center.powi(2)).sqrt())
                .collect::<Vec<Complex<f64>>>();
            for i in 0..p2.len() {
                p.push(p2[i]);
            }
            k = band_width.powi(order as i32);
        }
        Cutoff::BandStop(mut f1, mut f2) => {
            f1 = 2. * f1 / sample_rate;
            let warped1 = 4. * (pi * f1 / 2.).tan();
            f2 = 2. * f2 / sample_rate;
            let warped2 = 4. * (pi * f2 / 2.).tan();
            let band_width = warped1.max(warped2) - warped1.min(warped2);
            let band_center = (warped1 * warped2).sqrt();

            z = vec![Complex::i() * band_center; order];
            for _i in 0..z.len() {
                z.push(-Complex::i() * band_center);
            }
            k = -(1. / p.iter().copied().reduce(|a, b| -a * -b)
                .expect("Could not take product of p vector")).re();
            p = p.iter().map(|e| (band_width / 2.) / e).collect();
            let p2 = p.clone().iter().map(|e| *e - (e.powi(2) - band_center
                .powi(2)).sqrt()).collect::<Vec<Complex<f64>>>();
            p = p.iter().map(|e| *e + (e.powi(2) - band_center.powi(2)).sqrt())
                .collect::<Vec<Complex<f64>>>();
            for i in 0..p2.len() {
                p.push(p2[i]);
            }
        }
    }

    let mut num = 1.;
    let mut denom = 1.;
    if z.len() != 0 {
        let num_vec: Vec<Complex<f64>> = z.iter().copied().map(|e| 4. - e).collect();
        num = num_vec.into_iter().reduce(|a, b| a * b).expect("Could not multiply z vector").re();
    }
    if p.len() != 0 {
        let denom_vec: Vec<Complex<f64>> = p.iter().copied().map(|e| 4. - e).collect();
        denom = denom_vec.into_iter().reduce(|a, b| a * b).expect("Could not multiply p vector").re();
    }

    k = k * (num / denom);

    z = z.iter().map(|e| (4. + e) / (4. - e)).collect();
    match filter_type {
        Cutoff::LowPass(_) | Cutoff::BandPass(_, _) => {
            for _i in 0..order {
                z.push(Complex::from(-1.));
            }
        }
        _ => {}
    }
    p = p.iter().map(|e| (4. + e) / (4. - e)).collect();

    Ok(convert_zeros_poles_gain_to_transfer_function(&z, &p, k))
}

/// Based on: [iir_filters](https://crates.io/crates/iir_filters),
/// [Stack Overflow discussion](https://stackoverflow.com/a/33594706)
fn polynomial_from_roots(roots: &Vec<Complex<f64>>) -> Vec<Complex<f64>> {
    let mut coefficients = vec![Complex::<f64>::zero(); roots.len()];
    coefficients.push(Complex::<f64>::one());

    for root in roots {
        for idx in 0..coefficients.len() {
            if idx == 0 {
                coefficients[idx] = coefficients[idx + 1];
            } else if idx == coefficients.len() - 1 {
                coefficients[idx] = -root * coefficients[idx];
            } else {
                coefficients[idx] = coefficients[idx + 1] - root * coefficients[idx];
            }
        }
    }

    coefficients
}

/// Convert zeros, poles, and gain to transfer function coefficients.
fn convert_zeros_poles_gain_to_transfer_function(zeros: &Vec<Complex<f64>>, poles: &Vec<Complex<f64>>,
                                                 gain: f64) -> (Vec<f64>, Vec<f64>) {
    let numerator: Vec<f64> = polynomial_from_roots(zeros).iter().map(|x| x.re * gain).collect();
    let denominator = polynomial_from_roots(poles).iter().map(|x| x.re).collect();

    (numerator, denominator)
}


#[cfg(test)]
mod test {
    use super::polynomial_from_roots;
    use num_complex::Complex;

    #[test]
    fn test_low_pass_fifth_order() {
        let (numerator, denominator) = super::transfer_from_frequency(5, 81., super::Cutoff::LowPass(1.)).unwrap();
        let expected_numerator = vec![7.76046536e-08, 3.88023268e-07, 7.76046536e-07, 7.76046536e-07, 3.88023268e-07, 7.76046536e-08];
        let expected_denominator = vec![1.00000000e+00, -4.74899593e+00, 9.02719702e+00, -8.58524339e+00, 4.08498916e+00, -7.77944391e-01];
        for i in 0..numerator.len() {
            assert!((numerator[i] - expected_numerator[i]).abs() < 1e-14);
            assert!((denominator[i] - expected_denominator[i]).abs() < 1e-6);
        }
    }

    #[test]
    fn test_low_pass_third_order() {
        let (numerator, denominator) = super::transfer_from_frequency(3, 81., super::Cutoff::LowPass(1.)).unwrap();
        let expected_numerator = vec![5.40693725e-05, 1.62208118e-04, 1.62208118e-04, 5.40693725e-05];
        let expected_denominator = vec![1.00000000e+00, -2.84489831e+00, 2.70159248e+00, -8.56261608e-01];
        for i in 0..numerator.len() {
            assert!((numerator[i] - expected_numerator[i]).abs() < 1e-12);
            assert!((denominator[i] - expected_denominator[i]).abs() < 1e-6);
        }
    }

    #[test]
    fn test_high_pass_fifth_order() {
        let (numerator, denominator) = super::transfer_from_frequency(5, 81., super::Cutoff::HighPass(1.)).unwrap();
        let expected_numerator = vec![8.82011559e-01, -4.41005780e+00, 8.82011559e+00, -8.82011559e+00, 4.41005780e+00, -8.82011559e-01];
        let expected_denominator = vec![1.00000000e+00, -4.74899593e+00, 9.02719702e+00, -8.58524339e+00, 4.08498916e+00, -7.77944391e-01];
        for i in 0..numerator.len() {
            assert!((numerator[i] - expected_numerator[i]).abs() < 1e-6);
            assert!((denominator[i] - expected_denominator[i]).abs() < 1e-6);
        }
    }

    #[test]
    fn test_high_pass_third_order() {
        let (numerator, denominator) = super::transfer_from_frequency(3, 81., super::Cutoff::HighPass(1.)).unwrap();
        let expected_numerator = vec![9.25344050e-01, -2.77603215e+00, 2.77603215e+00, -9.25344050e-01];
        let expected_denominator = vec![1.00000000e+00, -2.84489831e+00, 2.70159248e+00, -8.56261608e-01];
        for i in 0..numerator.len() {
            assert!((numerator[i] - expected_numerator[i]).abs() < 1e-6);
            assert!((denominator[i] - expected_denominator[i]).abs() < 1e-6);
        }
    }

    #[test]
    fn test_band_pass_third_order() {
        let (numerator, denominator) = super::transfer_from_frequency(3, 81., super::Cutoff::BandPass(1., 10.)).unwrap();
        let expected_numerator = vec![2.36228384e-02, 0.00000000e+00, -7.08685151e-02, 0.00000000e+00, 7.08685151e-02, 0.00000000e+00, -2.36228384e-02];
        let expected_denominator = vec![1.00000000e+00, -4.48137769e+00, 8.48254824e+00, -8.75642447e+00, 5.23320668e+00, -1.71730743e+00, 2.39473853e-01];
        for i in 0..numerator.len() {
            assert!((numerator[i] - expected_numerator[i]).abs() < 1e-6);
            assert!((denominator[i] - expected_denominator[i]).abs() < 1e-6);
        }
    }

    #[test]
    fn test_band_stop_third_order() {
        let (numerator, denominator) = super::transfer_from_frequency(3, 81., super::Cutoff::BandStop(1., 10.)).unwrap();
        let expected_numerator = vec![4.89930497e-01, -2.84783273e+00, 6.98768389e+00, -9.25944412e+00, 6.98768389e+00, -2.84783273e+00, 4.89930497e-01];
        let expected_denominator = vec![1.00000000e+00, -4.48137769e+00, 8.48254824e+00, -8.75642447e+00, 5.23320668e+00, -1.71730743e+00, 2.39473853e-01];
        for i in 0..numerator.len() {
            assert!((numerator[i] - expected_numerator[i]).abs() < 1e-6);
            assert!((denominator[i] - expected_denominator[i]).abs() < 1e-6);
        }
    }

    #[test]
    fn test_poly() {
        let zeros = vec![Complex::from(1f64), Complex::from(2f64), Complex::from(3f64)];
        let coefficients = polynomial_from_roots(&zeros);
        assert_eq!(coefficients, vec![Complex::from(1f64), Complex::from(-6f64), Complex::from(11f64), Complex::from(-6f64)])
    }

    #[test]
    fn test_poly_difficult() {
        let zeros = vec![Complex::from(1f64), Complex::from(2f64), Complex::from(3f64), Complex::from(-2f64), Complex::from(0f64)];
        let coefficients = polynomial_from_roots(&zeros);
        assert_eq!(coefficients, vec![Complex::from(1f64), Complex::from(-4f64), Complex::from(-1f64), Complex::from(16f64), Complex::from(-12f64), Complex::from(0f64)])
    }

    #[test]
    fn test_poly_imaginary() {
        let zeros = vec![Complex::i(), Complex::from(2f64), Complex::from(3f64), Complex::from(-2f64), Complex::from(0f64)];
        let coefficients = polynomial_from_roots(&zeros);
        assert_eq!(coefficients, vec![Complex::from(1f64), Complex::from(-3f64) - Complex::i(), Complex::from(-4f64) + 3. * Complex::i(), Complex::from(12f64) + 4. * Complex::i(), -12. * Complex::i(), Complex::from(0f64)])
    }
}