//! # Complex - to - complex fourier space
use crate::enums::{BaseKind, TransformKind};
use crate::traits::{
    BaseElements, BaseFromOrtho, BaseGradient, BaseMatOpDiffmat, BaseMatOpLaplacian,
    BaseMatOpStencil, BaseSize, BaseTransform,
};
use crate::types::{FloatNum, ScalarNum};
use ndarray::{s, Array2};
use num_complex::Complex;
use num_traits::{One, Zero};
use rustfft::{Fft, FftPlanner};
use std::convert::TryInto;
use std::f64::consts::PI;
use std::ops::{Add, Div, Mul, Sub};
use std::sync::Arc;

/// # Container for fourier space (Complex-to-complex)
#[derive(Clone)]
pub struct FourierC2c<A> {
    /// Number of coefficients in physical space
    n: usize,
    /// Number of coefficients in spectral space
    m: usize,
    /// Handles fourier transforms
    plan_fwd: Arc<dyn Fft<A>>,
    plan_bwd: Arc<dyn Fft<A>>,
}

impl<A: FloatNum> FourierC2c<A> {
    /// Returns a new Fourier Basis for complex-to-complex transforms
    #[must_use]
    pub fn new(n: usize) -> Self {
        let mut planner = FftPlanner::<A>::new();
        Self {
            n,
            m: n,
            plan_fwd: Arc::clone(&planner.plan_fft_forward(n)),
            plan_bwd: Arc::clone(&planner.plan_fft_inverse(n)),
        }
    }

    /// Equispaced points on intervall [0, 2pi[
    ///
    /// ## Panics
    /// Float conversion fails
    #[must_use]
    #[allow(clippy::cast_precision_loss)]
    pub fn nodes(n: usize) -> Vec<A> {
        let n64 = n as f64;
        ndarray::Array1::range(0., 2. * PI, 2. * PI / n64)
            .mapv(|elem| A::from_f64(elem).unwrap())
            .to_vec()
    }

    /// Return complex wavenumber vector(0, 1, 2, -3, -2, -1)
    ///
    /// # Panics
    /// Float conversion fails
    #[must_use]
    #[allow(clippy::cast_precision_loss)]
    pub fn wavenumber(n: usize) -> Vec<Complex<A>> {
        let n2 = (n - 1) / 2 + 1;
        let mut k: Vec<Complex<A>> = vec![Complex::<A>::zero(); n];
        for (i, ki) in k.iter_mut().take(n2).enumerate() {
            ki.im = A::from(i).unwrap();
        }
        for (i, ki) in k.iter_mut().rev().take(n / 2).enumerate() {
            ki.im = A::from(-1. * (i + 1) as f64).unwrap();
        }
        k
    }
}

impl<A: FloatNum> BaseSize for FourierC2c<A> {
    /// Size in physical space
    fn len_phys(&self) -> usize {
        self.n
    }
    /// Size in spectral space
    fn len_spec(&self) -> usize {
        self.m
    }
    /// Size of orthogonal space
    fn len_orth(&self) -> usize {
        self.m
    }
}

impl<A: FloatNum> BaseElements for FourierC2c<A> {
    /// Real valued scalar type
    type RealNum = A;

    /// Return kind of base
    fn base_kind(&self) -> BaseKind {
        BaseKind::FourierC2c
    }

    /// Return kind of transform
    fn transform_kind(&self) -> TransformKind {
        TransformKind::C2c
    }

    /// Coordinates in physical space
    fn coords(&self) -> Vec<A> {
        Self::nodes(self.len_phys())
    }
}

impl<A: FloatNum> BaseMatOpDiffmat for FourierC2c<A> {
    /// Scalar type of matrix
    type NumType = Complex<A>;

    /// Explicit differential operator $ D $
    ///
    /// Matrix-based version of [`BaseGradient::gradient()`]
    ///
    /// # Panics
    /// Type conversion fails
    fn diffmat(&self, deriv: usize) -> Array2<Self::NumType> {
        assert!(deriv > 0);
        let mut mat = Array2::<Self::NumType>::zeros((self.len_spec(), self.len_spec()));
        let wavenum = Self::wavenumber(self.len_phys());
        for (l, k) in mat.diag_mut().iter_mut().zip(wavenum.iter()) {
            *l = k.powi(deriv.try_into().unwrap());
        }
        mat
    }

    /// Explicit inverse of differential operator $ D^* $
    ///
    /// Returns ``(D_pinv, I_pinv)``, where `D_pinv` is the pseudoinverse
    /// and ``I_pinv`` the corresponding pseudoidentity matrix, such
    /// that
    ///
    /// ```text
    /// D_pinv @ D = I_pinv
    /// ```
    ///
    /// Can be used as a preconditioner.
    fn diffmat_pinv(&self, deriv: usize) -> (Array2<Self::NumType>, Array2<Self::NumType>) {
        assert!(deriv > 0);
        let peye: Array2<Self::NumType> = Array2::<Self::NumType>::eye(self.m)
            .slice(s![1.., ..])
            .to_owned();
        let mut pinv = self.diffmat(deriv);
        for p in pinv.slice_mut(ndarray::s![1.., 1..]).diag_mut().iter_mut() {
            *p = Self::NumType::one() / *p;
        }
        (pinv, peye)
    }
}

impl<A: FloatNum> BaseMatOpStencil for FourierC2c<A> {
    /// Scalar type of matrix
    type NumType = A;

    /// Transformation stencil composite -> orthogonal space
    fn stencil(&self) -> Array2<Self::NumType> {
        Array2::<A>::eye(self.len_spec())
    }

    /// Inverse of transformation stencil
    fn stencil_inv(&self) -> Array2<Self::NumType> {
        Array2::<A>::eye(self.len_spec())
    }
}

impl<A: FloatNum> BaseMatOpLaplacian for FourierC2c<A> {
    /// Scalar type of laplacian matrix
    type NumType = A;

    /// Laplacian $ L $
    fn laplacian(&self) -> Array2<Self::NumType> {
        let wavenum = Self::wavenumber(self.len_phys());
        let mut lap = Array2::<A>::zeros((self.m, self.m));
        for (l, k) in lap.diag_mut().iter_mut().zip(wavenum.iter()) {
            *l = -k.im * k.im;
        }
        lap
    }

    /// Pseudoinverse matrix of Laplacian $ L^{-1} $
    ///
    /// Returns pseudoinverse and pseudoidentity,i.e
    /// ``(D_pinv, I_pinv)``
    ///
    /// ```text
    /// D_pinv @ D = I_pinv
    /// ``
    fn laplacian_pinv(&self) -> (Array2<Self::NumType>, Array2<Self::NumType>) {
        let mut d_pinv = self.laplacian();
        for p in d_pinv.slice_mut(s![1.., 1..]).diag_mut().iter_mut() {
            *p = A::one() / *p;
        }
        let i_pinv = Array2::<A>::eye(self.m);
        (d_pinv, i_pinv.slice(s![1.., ..]).to_owned())
    }
}

impl<A, T> BaseGradient<T> for FourierC2c<A>
where
    A: FloatNum,
    T: ScalarNum
        + Add<Complex<A>, Output = T>
        + Mul<Complex<A>, Output = T>
        + Div<Complex<A>, Output = T>
        + Sub<Complex<A>, Output = T>,
{
    /// Differentiate along slice
    /// ```
    /// use funspace::traits::BaseGradient;
    /// use funspace::fourier::FourierC2c;
    /// use funspace::utils::approx_eq_complex;
    /// use num_complex::Complex;
    /// use num_traits::Zero;
    /// let fo = FourierC2c::<f64>::new(5);
    /// let mut k = FourierC2c::<f64>::wavenumber(5);
    /// let expected: Vec<Complex<f64>> = k.iter().map(|x| x.powi(2)).collect();
    /// let mut outdata = vec![Complex::<f64>::zero(); 5];
    /// fo.gradient_slice(&k, &mut outdata, 1);
    /// approx_eq_complex(&outdata, &expected);
    /// ```
    ///
    /// # Panics
    /// When type conversion fails ( safe )
    fn gradient_slice(&self, indata: &[T], outdata: &mut [T], n_times: usize) {
        assert!(outdata.len() == indata.len());
        assert!(outdata.len() == self.len_spec());
        // Copy over
        for (y, x) in outdata.iter_mut().zip(indata.iter()) {
            *y = *x;
        }
        let n = self.len_spec();
        let n2 = (n - 1) / 2 + 1;
        for _ in 0..n_times {
            for (k, y) in outdata.iter_mut().take(n2).enumerate() {
                let ki: Complex<A> = Complex::<A>::new(A::zero(), A::from_usize(k).unwrap());
                *y = *y * ki;
            }
            for (k, y) in outdata.iter_mut().rev().take(n / 2).enumerate() {
                let ki: Complex<A> = Complex::<A>::new(A::zero(), A::from_usize(k).unwrap());
                *y = *y * ki;
            }
        }
    }
}

impl<A, T> BaseFromOrtho<T> for FourierC2c<A>
where
    A: FloatNum,
    T: ScalarNum,
{
    /// Composite space coefficients -> Orthogonal space coefficients
    fn to_ortho_slice(&self, indata: &[T], outdata: &mut [T]) {
        for (y, x) in outdata.iter_mut().zip(indata.iter()) {
            *y = *x;
        }
    }

    /// Orthogonal space coefficients -> Composite space coefficients
    fn from_ortho_slice(&self, indata: &[T], outdata: &mut [T]) {
        for (y, x) in outdata.iter_mut().zip(indata.iter()) {
            *y = *x;
        }
    }
}

impl<A: FloatNum> BaseTransform for FourierC2c<A> {
    type Physical = Complex<A>;

    type Spectral = Complex<A>;

    /// # Example
    /// Forward transform
    /// ```
    /// use funspace::traits::BaseTransform;
    /// use funspace::fourier::FourierC2c;
    /// use funspace::utils::approx_eq_complex;
    /// use num_complex::Complex;
    /// use num_traits::Zero;
    /// let mut fo = FourierC2c::<f64>::new(4);
    /// let indata = vec![
    ///     Complex::new(1., 1.),
    ///     Complex::new(2., 2.),
    ///     Complex::new(3., 3.),
    ///     Complex::new(4., 4.)
    /// ];
    /// let expected = vec![
    ///     Complex::new(10., 10.),
    ///     Complex::new(-4., 0.),
    ///     Complex::new(-2., -2.),
    ///     Complex::new(0., -4.)
    /// ];
    /// let mut outdata = vec![Complex::<f64>::zero(); 4];
    /// fo.forward_slice(&indata, &mut outdata);
    /// approx_eq_complex(&outdata, &expected);
    /// ```
    fn forward_slice(&self, indata: &[Self::Physical], outdata: &mut [Self::Spectral]) {
        // Check input
        assert!(indata.len() == self.len_phys());
        assert!(outdata.len() == self.len_spec());
        // Copy
        for (a, b) in outdata.iter_mut().zip(indata.iter()) {
            *a = *b;
        }
        self.plan_fwd.process(outdata);
    }

    /// # Example
    /// Backward transform
    /// ```
    /// use funspace::traits::BaseTransform;
    /// use funspace::fourier::FourierC2c;
    /// use funspace::utils::approx_eq_complex;
    /// use num_complex::Complex;
    /// use num_traits::Zero;
    /// let mut fo = FourierC2c::<f64>::new(4);
    /// let indata = vec![
    ///     Complex::new(10., 10.),
    ///     Complex::new(-4., 0.),
    ///     Complex::new(-2., -2.),
    ///     Complex::new(0., -4.)
    /// ];
    /// let expected = vec![
    ///     Complex::new(1., 1.),
    ///     Complex::new(2., 2.),
    ///     Complex::new(3., 3.),
    ///     Complex::new(4., 4.)
    /// ];
    /// let mut outdata = vec![Complex::<f64>::zero(); 4];
    /// fo.backward_slice(&indata, &mut outdata);
    /// approx_eq_complex(&outdata, &expected);
    /// ```
    fn backward_slice(&self, indata: &[Self::Spectral], outdata: &mut [Self::Physical]) {
        // Check input
        assert!(indata.len() == self.len_spec());
        assert!(outdata.len() == self.len_phys());
        // Copy and correct fft
        let cor = A::one() / A::from_usize(self.len_phys()).unwrap();
        for (a, b) in outdata.iter_mut().zip(indata.iter()) {
            a.re = cor * b.re;
            a.im = cor * b.im;
        }
        self.plan_bwd.process(outdata);
    }
}