use std::f32::consts::PI;
use num_traits::{Float, FromPrimitive, One, Zero};
use num_complex::Complex;

use crate::spectrum::SpectrumError;

fn re_arrange(dat: usize, bitlength: usize) -> usize {
    let mut ret = 0;
    for i in 0..bitlength {
        if (dat & (1 << i)) != 0 {
            ret |= (1 << (bitlength - 1)) >> i;
        }
    }
    ret
}

fn re_log2n(count: u32) -> u32 {
    let mut log2n = 0;
    let mask: u32 = 0x8000_0000; // 2147483648

    for i in 0..32 {
        if ((mask >> i) & count) != 0 {
            if (mask >> i) == count {
                log2n = 31 - i;
            } else {
                log2n = 31 - i + 1;
            }
            break;
        }
    }
    log2n
}


/// Inverse FFT
pub fn ifft<T>(input: &mut Vec<Complex<T>>) -> Result<Vec<T>, SpectrumError> where 
    T: Float + FromPrimitive + Zero + Clone + One,
{
    let n = input.len();
    let mut return_data: Vec<T> = vec![T::zero(); n];
    // 生成WN 表
    let mut wn: Vec<Vec<T>> = vec![vec![T::zero(); 2]; n];
    for (i, val) in wn.iter_mut().enumerate(){
        *val = vec![
            T::from((2.0 * PI * i as f32 / n as f32).cos()).unwrap(), 
            T::from((2.0 * PI * i as f32 / n as f32).sin()).unwrap()
        ];
    }

    // 蝶形计算
    let mut steplenght = 2;
    while steplenght <= n {
        let mut step = 0;
        while step < steplenght / 2 {
            for i in 0..(n / steplenght) {
                let index0 = (n * 2 / steplenght) * step + i;
                let index1 = (n * 2 / steplenght) * step + n / steplenght + i;
                let temp_re1 = input[index0].re + input[index1].re;
                let temp_im1 = input[index0].im + input[index1].im;
                let temp_re2 = input[index0].re - input[index1].re;
                let temp_im2 = input[index0].im - input[index1].im;

                input[index1] = Complex::new(
                    temp_re2 * wn[steplenght / 2 * i][0] - temp_im2 * wn[steplenght / 2 * i][1],
                    temp_im2 * wn[steplenght / 2 * i][0] + temp_re2 * wn[steplenght / 2 * i][1]
                );

                input[index0] = Complex::new(temp_re1, temp_im1);

            }
            step += 1;
        }
        steplenght *= 2;
    }
    // 重组数据
    let bitlenghth = re_log2n(n as u32) as usize;
    let mut index;
    let mut register: Vec<Complex<T>> = vec![Complex::new(T::zero(), T::zero()); n];
    for i in 0..n {
        index = re_arrange(i, bitlenghth);
        register[i] = if index < input.len() { input[index].clone() } else { Complex::new(T::zero(), T::zero()) };
        return_data[i] = register[i].re;
    }
    // Conjugate and normalize
    Ok(return_data)
}
