use num::Complex;
use num::Float;
use num::FromPrimitive;
use num::One;
use num::Zero;
use num::complex::ComplexFloat;
use rustfft::FftNum;
use rustfft::FftPlanner;
use std::f64::consts::PI;
use std::iter::Sum;
use std::ops::Add;
use std::ops::Div;
use std::ops::Mul;

use crate::filtering::history_buf::HistoryBuf;
use crate::map;
use crate::units::DigitalFrequency;


/// Represents a finite impulse response as a vector
/// of values in time.
///
/// Convention
/// indices : 0 ----------------- fir.0.len - 1
/// time    :   ---------------->
///
/// For a reverb ir for example the clap would be at index 0
/// and the reverb tail towards the end of the vector.
pub struct Fir<T>(pub Vec<T>);

impl<T> Fir<Complex<T>>
where
    T: FftNum + Float + Clone,
{
    /// Synthesizes an impulse response from a transfer function using an inverse fourrier
    /// transform.
    ///
    /// The input units are thus :
    /// Transfer function :
    /// start                          center                    end
    /// [                                                          ]
    /// frequency :
    /// 0                              pi                         2*pi
    /// =                              pi, -pi                     0
    ///                          nyquist's frequency
    /// (As the frequencies are periodic)
    ///
    pub fn from_transfer_function(tf: impl AsRef<[Complex<T>]>) -> Fir<Complex<T>>
    {
        let mut planner = FftPlanner::new();
        let tf_len = tf.as_ref().len();
        let ifft = planner.plan_fft_inverse(tf.as_ref().len());

        let mut fir = tf.as_ref().to_vec();
        ifft.process(fir.as_mut_slice());

        let mut shifted_fir = vec![];
        for i in 0..tf_len
        {
            let k = (tf_len - (tf_len / 2) + i) % tf_len;
            shifted_fir.push(fir[k]);
        }

        Fir(shifted_fir)
    }

    /// Creates a low pass filter with the following ideal transfer function using the ifft method:
    ///
    /// ________________                      ________________
    ///                 |____________________|
    /// 0           cuttoff               -cuttoff            0
    ///
    pub fn lowpass(cutoff: DigitalFrequency, length: usize) -> Fir<Complex<T>>
    {
        let mut tf = vec![Complex::<T>::zero(); length];

        let cutoff_bin = map(cutoff.as_rad(), 0., 2. * PI, 0., length as f64).floor() as usize;
        for i in 0..cutoff_bin
        {
            tf[i] = Complex::<T>::one();
            tf[length - i - 1] = Complex::<T>::one();
        }

        Self::from_transfer_function(tf)
    }
}

impl<T> Fir<T>
where
    T: ComplexFloat + Div<T::Real, Output = T> + Copy + Sum,
    T::Real: Float,
{
    /// Returns the same impulse response
    /// normalized by the length of the sum of the vectors.
    pub fn normalized(mut self) -> Self
    {
        let sum: T = self.0.iter().copied().sum();
        let len = Float::sqrt(sum.im() * sum.im() + sum.re() * sum.re());

        self.0.iter_mut().for_each(|x| *x = *x / len);
        self
    }
}

impl<T> Fir<T>
where
    T: ComplexFloat + Div<T::Real, Output = T>,
    T::Real: Float + FromPrimitive,
{
    /// Returns the same impulse response
    /// normalized by the length of the impulse response.
    pub fn normalized_len(mut self) -> Self
    {
        let len = T::Real::from_usize(self.0.len()).unwrap();
        self.0.iter_mut().for_each(|x| *x = *x / len);
        self
    }
}

impl<T> Fir<T>
where
    T: ComplexFloat + Div<T::Real, Output = T> + Copy + Add<T, Output = T>,
    T::Real: Float,
{
    /// Returns the same impulse response
    /// normalized by the energy or the sum of the squares of the magnitues
    /// of the impulse response
    pub fn normalized_sqr(mut self) -> Self
    {
        let sum = self
            .0
            .iter()
            .copied()
            .map(|x| x.abs() * x.abs())
            .reduce(|x, y| x + y)
            .unwrap();

        self.0.iter_mut().for_each(|x| *x = *x / sum);
        self
    }
}

impl<T> Fir<T>
where
    T: ComplexFloat + FromPrimitive,
{
    /// Creates a square unit impulse response :
    /// a vector of length `length` filled with ones
    pub fn square(length: usize) -> Self
    {
        Self((0..length).map(|_| T::one()).collect())
    }

    /// Creates a simple proportional integral (PI) loop impulse response :
    ///
    /// FIR:
    /// ```text
    /// _ ...................................  Kp (`proportional_gain`)
    ///  |
    ///  |
    ///  |
    ///  |____________________________  ...... Ki (`integral_gain`)
    ///                               |
    ///
    /// 0 ------------------------- `length`
    /// ```
    ///
    pub fn proportional_integral(length: usize, proportional_gain: T, integral_gain: T) -> Self
    {
        Self(
            (0..length)
                .map(|i| {
                    if i == 0
                    {
                        proportional_gain + integral_gain
                    }
                    else
                    {
                        integral_gain
                    }
                })
                .collect(),
        )
    }

    /// Creates a root raised cosine (RRC) FIR of length `length`
    /// with the given roll off factor.
    ///
    /// The corresponding RC (convolution of this filter with itself)
    /// has its zero crossing every `symbol_length` samples (except at 0).
    pub fn root_raised_cosine(length: usize, roll_off: f64, symbol_length: usize) -> Self
    {
        Self(
            (0..length)
                .map(|i| {
                    let t = map(
                        i as f64,
                        0.,
                        length as f64,
                        -(length as f64) * 0.5,
                        length as f64 * 0.5,
                    ) / symbol_length as f64;
                    T::from_f64(root_raised_cosine(t, roll_off, 1.)).unwrap()
                })
                .collect(),
        )
    }
}

/// A simple convolutional finite impulse response filter
pub struct FirFilter<F, T, O>
where
    F: Mul<T, Output = O>,
    O: Add<O, Output = O> + Sum + Clone + Zero,
{
    fir: Vec<F>,
    //taps: VecDeque<T>,
    taps: HistoryBuf<T>
}

impl<F, T, O> FirFilter<F, T, O>
where
    T: Clone + Zero,
    F: Mul<T, Output = O> + Clone,
    O: Add<O, Output = O> + Sum + Clone + Zero,
{
    /// Creates a filter with the given impulse response
    pub fn new(impulse_response: Fir<F>) -> Self
    {
        let len = impulse_response.0.len();
        Self {
            fir: impulse_response.0,
            taps: HistoryBuf::new(T::zero(), len),
        }
    }

    /// Gets the next output given an input sample.
    ///
    /// At the beginning, the delay line starts with zeroes.
    pub fn next(&mut self, input: T) -> O
    {
        self.taps.push(input);
        let taps = self.taps.as_slice();
        Self::dot_prod(&self.fir, taps) 
    }

    pub fn dot_prod(a: &[F], b: &[T]) -> O
    {
        assert_eq!(a.len(), b.len());

        let mut sum: [_; 4] = [O::zero(), O::zero(), O::zero(), O::zero()];

        let (a_chunks, a_remainder) = a.as_chunks::<4>();
        let (b_chunks, b_remainder) = b.as_chunks::<4>();

        for (x, y) in a_chunks.iter().zip(b_chunks.iter())

        {
                sum[0] = sum[0].clone() + x[0].clone() * y[0].clone();
                sum[1] = sum[1].clone() + x[1].clone() * y[1].clone();
                sum[2] = sum[2].clone() + x[2].clone() * y[2].clone();
                sum[3] = sum[3].clone() + x[3].clone() * y[3].clone();
        }

        let mut sum = sum[0].clone() + sum[1].clone() + sum[2].clone() + sum[3].clone();
        for (x, y) in a_remainder.iter().zip(b_remainder.iter())
        {
            sum = sum + x.clone() * y.clone();
        }

        sum
    }
}

// Completely stolen from sdrpp code
pub fn estimate_fir_length(transition_width: f64, sample_rate: f64) -> f64
{
    3.8 * sample_rate / transition_width
}

/// Root raised cosine function
pub fn root_raised_cosine(t: f64, beta: f64, symbol_time: f64) -> f64
{
    let eps = 1e-8;

    if t.abs() < eps
    {
        // t = 0 special case
        return (1.0 / symbol_time.sqrt()) * (1.0 + beta * (4.0 / PI - 1.0));
    }

    if beta > 0.0 && (t.abs() - symbol_time / (4.0 * beta)).abs() < eps
    {
        // t = ±T / (4β) special case
        let term1 = (1.0 + 2.0 / PI) * (PI / (4.0 * beta)).sin();
        let term2 = (1.0 - 2.0 / PI) * (PI / (4.0 * beta)).cos();
        return (beta / (symbol_time.sqrt() * 2.0_f64.sqrt())) * (term1 + term2);
    }

    // General case
    let numerator = (PI * t * (1.0 - beta) / symbol_time).sin()
        + 4.0 * beta * t / symbol_time * (PI * t * (1.0 + beta) / symbol_time).cos();

    let denominator = PI * t * (1.0 - (4.0 * beta * t / symbol_time).powi(2)) / symbol_time;

    (1.0 / symbol_time.sqrt()) * (numerator / denominator)
}
