use std::{f32::consts::PI, vec};
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
}

/// FFT using radix-2 Cooley-Tukey algorithm
pub fn fft<T>(input: &Vec<T>) -> Result<Vec<Complex<T>>, SpectrumError> where 
    T: Float + FromPrimitive + Zero + Clone + One,
{
    let bits: usize = re_log2n(input.len() as u32) as usize;
    let n = 0x01 << bits;
    let mut register: Vec<Complex<T>> = vec![Complex::new(T::zero(), T::zero()); n];
    for i in 0..n {
        let j = re_arrange(i, bits);
        register[i] = if j < input.len() {
            Complex::new(input[j], T::zero())
        } else {
            Complex::new(T::zero(), T::zero())
        };
    }
    // 生成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 < n / steplenght {
            for i in 0..(steplenght / 2) {
                let index0 = steplenght * step + i;
                let index1 = steplenght * step + i + steplenght / 2;
                let temp_re = register[index1].re * wn[n / steplenght * i][0] - register[index1].im * wn[n / steplenght * i][1];
                let temp_im = register[index1].im * wn[n / steplenght * i][0] + register[index1].re * wn[n / steplenght * i][1];

                register[index1] = Complex::new(register[index0].re - temp_re, register[index0].im - temp_im);
                register[index0] = Complex::new(register[index0].re + temp_re, register[index0].im + temp_im);

            }
              
            step += 1;
        }
        steplenght *= 2;
    }
    Ok(register)
}