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