use crate::error::Error;
use itertools::Itertools;
use ndarray::{Array, Array1, ArrayD, AsArray, Axis, Dimension, RemoveAxis, concatenate, s};
use ndrustfft::{FftHandler, Normalization, ndfft_par, ndifft_par};
use num::Complex;
use num::traits::FloatConst;
use std::ops::MulAssign;

/// Fourier transform
pub fn fft<'a, A, D>(array: A) -> Result<Array<Complex<f64>, D>, Error>
where
    A: AsArray<'a, f64, D>,
    D: Dimension,
{
    let mut input = array.into().mapv(|i| Complex::new(i, 0.0));
    let shape = input.shape().to_vec();
    let mut tmp = ArrayD::zeros(shape.clone()).into_dimensionality()?;
    for (i, s) in shape.iter().enumerate() {
        let handler = FftHandler::new(*s).normalization(Normalization::None);
        ndfft_par(&input.view(), &mut tmp.view_mut(), &handler, i);
        std::mem::swap(&mut input, &mut tmp);
    }
    Ok(input)
}

/// inverse Fourier transform
pub fn ifft<'a, A, D>(array: A) -> Result<Array<Complex<f64>, D>, Error>
where
    A: AsArray<'a, Complex<f64>, D>,
    D: Dimension,
{
    let mut input = array.into().to_owned();
    let shape = input.shape().to_vec();
    let mut tmp = ArrayD::zeros(shape.clone()).into_dimensionality()?;
    for (i, s) in shape.iter().enumerate() {
        let handler = FftHandler::new(*s).normalization(Normalization::None);
        ndifft_par(&input.view(), &mut tmp.view_mut(), &handler, i);
        std::mem::swap(&mut input, &mut tmp);
    }
    Ok(input / Complex::from(shape.iter().product::<usize>() as f64))
}

pub fn fft_freq(size: usize) -> Array1<f64> {
    // let s = size as isize;
    // let h = s / 2;
    // (-h..h).map(|i| (i % s - s) as f64 / s as f64).collect()

    let val = 1.0 / size as f64;
    let mut results = Array1::zeros(size);
    let n = (size - 1) / 2 + 1;
    let p1 = Array1::range(0.0, n as f64, 1.0);
    results.slice_mut(s![..n]).assign(&p1);
    let p2 = Array1::range(-((size / 2) as f64), 0.0, 1.0);
    results.slice_mut(s![n..]).assign(&p2);
    results * val
}

pub fn fft_shift<'a, A, T, D>(x: A) -> Array<T, D>
where
    A: AsArray<'a, T, D>,
    T: 'a + Clone,
    D: Dimension + RemoveAxis,
{
    let mut x = x.into().to_owned();
    let shift = x.shape().iter().map(|i| i / 2).collect::<Vec<_>>();
    for (i, s) in shift.iter().enumerate() {
        let (a, b) = x.view().split_at(Axis(i), *s);
        x = concatenate(Axis(i), &[b, a]).unwrap();
    }
    x
}

pub fn gaussian(x: &[f64], mu: f64, sigma: f64) -> Vec<f64> {
    let a = 2.0 * sigma.powi(2);
    let b = (a * f64::PI()).sqrt();
    x.iter()
        .map(|i| (-(i - mu).powi(2) / a).exp() / b)
        .collect()
}

/// Gaussian kernel in frequency space
pub fn gaussian_kernel<D>(shape: &[usize], sigma: &[f64]) -> Result<Array<Complex<f64>, D>, Error>
where
    D: Dimension,
{
    let mut g = Array::ones(shape).into_dimensionality()?;
    for (i, (s, t)) in shape.iter().zip_eq(sigma.iter()).enumerate() {
        let a = 0.5 / t.powi(2);
        let f = (-(f64::PI() * fft_freq(*s)).powi(2) / a)
            .exp()
            .mapv(Complex::from);
        for mut lane in g.lanes_mut(Axis(i)) {
            lane.mul_assign(&f);
        }
    }
    Ok(g)
}

/// smooth an array using a Gaussian kernel
pub fn gaussian_smooth<'a, A, D>(array: A, sigma: &[f64]) -> Result<Array<f64, D>, Error>
where
    A: AsArray<'a, f64, D>,
    D: Dimension,
{
    let array = array.into().to_owned();
    let shape = array.shape();
    let kernel = gaussian_kernel::<D>(shape, sigma)?;
    Ok(ifft((kernel * fft(array.view())?).view())?.mapv(|i| i.re))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::julia_image;
    use ndarray::{Ix1, array};
    use tiffwrite::IJTiffFile;

    #[test]
    fn smooth() -> Result<(), Box<dyn std::error::Error>> {
        let a = array![0.0, 0.0, 1.0, 0.0, 1.0, 2.0, 1.0, 0.0, 1.0, 0.0, 0.0];
        let b = gaussian_smooth(a.view(), &[1.4])?;
        let c = array![
            0.13775053, 0.26443112, 0.45036765, 0.67622086, 0.93678676, 1.06888618, 0.93678676,
            0.67622086, 0.45036765, 0.26443112, 0.13775053
        ];
        debug_assert!(b.iter().zip_eq(c.iter()).all(|(x, y)| (x - y).abs() < 1e-8));
        Ok(())
    }

    #[test]
    fn smooth2() -> Result<(), Box<dyn std::error::Error>> {
        let im_a = julia_image(
            &[60, 80],
            &[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
            &[29.5, 39.5],
            &[-0.8, 0.156],
        );
        let im_b = gaussian_smooth(im_a.mapv(|i| i as f64).view(), &[4.0, 4.0])?;
        let mut t = IJTiffFile::new(std::env::home_dir().unwrap().join("tmp/julia.tif"))?;
        t.save(im_a.mapv(|i| i as u32).view(), 0, 0, 0)?;
        t.save(im_b.mapv(|i| i as u32).view(), 1, 0, 0)?;
        Ok(())
    }


    #[test]
    fn fft_test() -> Result<(), Box<dyn std::error::Error>> {
        let x = Array1::linspace(0.0, 1.0, 11).to_vec();
        let y = gaussian(x.as_slice(), 0.4, 0.1);
        let j = fft(Array1::from_vec(y).view())?;
        let k = array![
            Complex::new(9.99998512536899, 0.0),
            Complex::new(-5.562906924752932, -6.419930454175756),
            Complex::new(-0.7410757333589753, 5.154238692693016),
            Complex::new(1.9379781953109205, -1.2454762622580575),
            Complex::new(-0.7086874511902606, -0.20810344788799306),
            Complex::new(0.08206001372382088, 0.1796510861631586),
            Complex::new(0.08206001372382088, -0.1796510861631586),
            Complex::new(-0.7086874511902606, 0.20810344788799306),
            Complex::new(1.9379781953109205, 1.2454762622580575),
            Complex::new(-0.7410757333589753, -5.154238692693016),
            Complex::new(-5.562906924752932, 6.419930454175756),
        ];
        debug_assert!(
            j.iter()
                .zip_eq(k.iter())
                .all(|(x, y)| (x.re - y.re) < 1e-8 && (x.im - y.im) < 1e-8)
        );
        Ok(())
    }

    #[test]
    fn ifft_test() -> Result<(), Box<dyn std::error::Error>> {
        let x = Array1::linspace(0.0, 1.0, 11).to_vec();
        let y = gaussian(x.as_slice(), 0.4, 0.1)
            .into_iter()
            .map(Complex::from)
            .collect::<Vec<_>>();
        let j = ifft(Array1::from_vec(y).view())?;
        let k = array![
            Complex::new(0.9090895568517264, 0.0),
            Complex::new(-0.5057188113411757, 0.5836300412887051),
            Complex::new(-0.06737052121445229, -0.4685671538811833),
            Complex::new(0.17617983593735642, 0.1132251147507325),
            Complex::new(-0.06442613192638733, 0.018918495262544823),
            Complex::new(0.00746000124762008, -0.01633191692392351),
            Complex::new(0.00746000124762008, 0.01633191692392351),
            Complex::new(-0.06442613192638733, -0.018918495262544823),
            Complex::new(0.17617983593735642, -0.1132251147507325),
            Complex::new(-0.06737052121445229, 0.4685671538811833),
            Complex::new(-0.5057188113411757, -0.5836300412887051),
        ];
        debug_assert!(
            j.into_iter()
                .zip_eq(k.into_iter())
                .all(|(x, y)| ((x.re - y.re).abs() < 1e-8) && ((x.im - y.im).abs() < 1e-8))
        );
        Ok(())
    }

    #[test]
    fn fft_freq_test() -> Result<(), Box<dyn std::error::Error>> {
        let f = fft_freq(11);
        let g = vec![
            0.0,
            0.09090909090909091,
            0.18181818181818182,
            0.2727272727272727,
            0.36363636363636365,
            0.4545454545454546,
            -0.4545454545454546,
            -0.36363636363636365,
            -0.2727272727272727,
            -0.18181818181818182,
            -0.09090909090909091,
        ];
        debug_assert!(f.iter().zip_eq(g.iter()).all(|(x, y)| (x - y).abs() < 1e-8));
        Ok(())
    }

    #[test]
    fn gaussian_kernel_test() -> Result<(), Box<dyn std::error::Error>> {
        let x = gaussian_kernel::<Ix1>(&[11], &[1.5])?.mapv(|i| i.re);
        println!("x = {:?}", x.to_vec());
        let y = array![
            1.00000000e+00,
            6.92774033e-01,
            2.30338431e-01,
            3.67556755e-02,
            2.81491714e-03,
            1.03464181e-04,
            1.03464181e-04,
            2.81491714e-03,
            3.67556755e-02,
            2.30338431e-01,
            6.92774033e-01
        ];
        debug_assert!(y.iter().zip_eq(x.iter()).all(|(x, y)| (x - y).abs() < 1e-8));
        Ok(())
    }

    #[test]
    fn fft_shift_test() -> Result<(), Box<dyn std::error::Error>> {
        let x = Array1::from_iter(0..10);
        let y = fft_shift(x.view());
        let z = array![5, 6, 7, 8, 9, 0, 1, 2, 3, 4];
        debug_assert_eq!(y, z);
        Ok(())
    }

    #[test]
    fn gaussian_test() -> Result<(), Box<dyn std::error::Error>> {
        let x = Array1::linspace(0.0, 1.0, 11).to_vec();
        let y = gaussian(x.as_slice(), 0.4, 0.1);
        let z = vec![
            1.33830226e-03,
            4.43184841e-02,
            5.39909665e-01,
            2.41970725e+00,
            3.98942280e+00,
            2.41970725e+00,
            5.39909665e-01,
            4.43184841e-02,
            1.33830226e-03,
            1.48671951e-05,
            6.07588285e-08,
        ];
        debug_assert!(y.iter().zip_eq(z.iter()).all(|(x, y)| (x - y).abs() < 1e-8));
        Ok(())
    }
}