mrrs 0.1.1

use num_complex::Complex;
use num_traits::{Float, Zero};
use std::ops::{AddAssign, Mul};

pub mod triband;

/// Naive convolution
/// Output length = signal.len() - kernel.len() + 1
pub fn convolve<TS, TK, TR>(signal: &[TS], kernel: &[TK]) -> Vec<TR>
where
    TS: Copy + Mul<TK, Output = TR>,
    TK: Copy,
    TR: Copy + Zero + AddAssign,
{
    if kernel.len() > signal.len() {
        return vec![];
    }

    let mut result = vec![TR::zero(); signal.len() - kernel.len() + 1];

    for ii in 0..result.len() {
        let mut acc = TR::zero();
        for jj in 0..kernel.len() {
            acc += signal[ii + jj] * kernel[jj];
        }
        result[ii] = acc;
    }

    result
}

/// Trait to abstract over real or complex input values.
pub trait ToComplex<T: Float> {
    fn to_complex(self) -> Complex<T>;
}

impl<T: Float> ToComplex<T> for T {
    fn to_complex(self) -> Complex<T> {
        Complex::new(self, T::zero())
    }
}

impl<T: Float> ToComplex<T> for Complex<T> {
    fn to_complex(self) -> Complex<T> {
        self
    }
}

/// Frequency shift input by Fs/4, with coefficients for up vs down phase offset (0-3).
pub fn shift_fs_over_4<T, U>(
    input: &[U],
    coeffs: &[Complex<T>],
    phase_offset: usize,
) -> Vec<Complex<T>>
where
    T: Float,
    U: Copy + ToComplex<T>,
{
    input
        .iter()
        .enumerate()
        .map(|(i, &x)| {
            let coeff = coeffs[(i + phase_offset) % 4];
            x.to_complex() * coeff
        })
        .collect()
}

/// Frequency shift input up by Fs/4, with optional phase offset (0-3).
pub fn shift_up_fs_over_4<T, U>(input: &[U], phase_offset: usize) -> Vec<Complex<T>>
where
    T: Float,
    U: Copy + ToComplex<T>,
{
    let coeffs = [
        Complex::new(T::one(), T::zero()),  //  1
        Complex::new(T::zero(), -T::one()), // -j
        Complex::new(-T::one(), T::zero()), // -1
        Complex::new(T::zero(), T::one()),  // +j
    ];
    shift_fs_over_4(input, &coeffs, phase_offset)
}

/// Frequency shift input down by Fs/4, with optional phase offset (0-3).
pub fn shift_down_fs_over_4<T, U>(input: &[U], phase_offset: usize) -> Vec<Complex<T>>
where
    T: Float,
    U: Copy + ToComplex<T>,
{
    let coeffs = [
        Complex::new(T::one(), T::zero()),  //  1
        Complex::new(T::zero(), T::one()),  // +j
        Complex::new(-T::one(), T::zero()), // -1
        Complex::new(T::zero(), -T::one()), // -j
    ];
    shift_fs_over_4(input, &coeffs, phase_offset)
}

#[cfg(test)]
mod test {
    use super::*;
    use approx::assert_relative_eq;

    #[test]
    fn test_convolve_real() {
        let signal = vec![1.0, 2.0, 3.0, 4.0, 5.0];
        let kernel = vec![0.2, 0.5, 0.2];

        let output = convolve(&signal, &kernel);
        assert_eq!(output, vec![1.8, 2.7, 3.6]);
    }
    #[test]
    fn test_convolve_complex() {
        let signal: Vec<Complex<f32>> = (1..=5).map(|x| Complex::new(x as f32, 1.0)).collect();
        let kernel = vec![0.2, 0.5, 0.2];

        let output = convolve(&signal, &kernel);
        for (x, y) in output.iter().zip(
            &vec![1.8, 2.7, 3.6]
                .into_iter()
                .map(|x| Complex::new(x as f32, 0.9))
                .collect::<Vec<Complex<f32>>>(),
        ) {
            assert_relative_eq!(x.re, y.re, epsilon = 1e-9);
            assert_relative_eq!(x.im, y.im, epsilon = 1e-9);
        }
    }
    #[test]
    fn test_shift_down_fs_over_4_real_input() {
        let input = vec![1.0f32, 2.0, 3.0, 4.0];
        let output = shift_down_fs_over_4(&input, 0);

        // Expected: input * [1, j, -1, -j]
        let expected = vec![
            Complex::new(1.0, 0.0),
            Complex::new(0.0, 2.0),
            Complex::new(-3.0, 0.0),
            Complex::new(0.0, -4.0),
        ];

        for (o, e) in output.iter().zip(expected.iter()) {
            assert_relative_eq!(o.re, e.re, epsilon = 1e-6);
            assert_relative_eq!(o.im, e.im, epsilon = 1e-6);
        }
    }

    #[test]
    fn test_shift_down_fs_over_4_complex_input_with_phase() {
        let input = vec![
            Complex::new(1.0, 1.0),
            Complex::new(2.0, 0.0),
            Complex::new(0.0, 3.0),
            Complex::new(-1.0, -1.0),
        ];
        let output = shift_down_fs_over_4(&input, 1); // start at -j

        // Expected: input * [j, -1, -j, 1]
        let expected = vec![
            Complex::new(-1.0, 1.0), // (1+i) * j = -1 + i
            Complex::new(-2.0, 0.0),
            Complex::new(3.0, 0.0), // (0+3i)*-j = 3
            Complex::new(-1.0, -1.0),
        ];

        for (o, e) in output.iter().zip(expected.iter()) {
            assert_relative_eq!(o.re, e.re, epsilon = 1e-6);
            assert_relative_eq!(o.im, e.im, epsilon = 1e-6);
        }
    }
    #[test]
    fn test_shift_up_fs_over_4_complex_input_with_phase() {
        let input = vec![
            Complex::new(1.0, 0.0),
            Complex::new(2.0, 0.0),
            Complex::new(3.0, 0.0),
            Complex::new(4.0, 0.0),
        ];
        let output = shift_up_fs_over_4(&input, 1); // start at -j

        // Expected: input * [-j, -1, j, 1]
        let expected = vec![
            Complex::new(0.0, -1.0),
            Complex::new(-2.0, 0.0),
            Complex::new(0.0, 3.0),
            Complex::new(4.0, 0.0),
        ];

        for (o, e) in output.iter().zip(expected.iter()) {
            assert_relative_eq!(o.re, e.re, epsilon = 1e-6);
            assert_relative_eq!(o.im, e.im, epsilon = 1e-6);
        }
    }
}