// Cooley-Tukey algorithm

use crate::complex::Complex32;
use crate::fft::{DFT, FFTDirection};
use std::f32::consts::PI;

pub struct Radix2FFT {
    direction: FFTDirection,
    size: usize,
    length: usize,
}

impl DFT for Radix2FFT {
    // Size as power of two
    fn create(size: usize, direction: FFTDirection) -> Self {
        if !is_power_of_two(size) {
            panic!("Tried to create a Radix2 FFT with a non power of two size.");
        }

        Radix2FFT {
            size: size.ilog2() as usize,
            direction,
            length: size,
        }
    }

    fn execute(&mut self, input: &[Complex32], output: &mut [Complex32], window: fn(f32) -> f32) {
        // Reorder samples
        for (i, x) in output.iter_mut().enumerate() {
            let k = reverse_bits(i, self.size as u32);
            *x = input[k] * window(k as f32 / self.size as f32);
        }

        for step in 1..(self.size + 1) {
            let pol_length = 2usize.pow(step as u32);
            let mid_point = pol_length / 2;
            for s in (0..(self.length / pol_length)).map(|i| i * pol_length) {
                for i in 0..mid_point {
                    // Compute current polynomial at each unit root
                    let a = output[s + i];
                    let b = output[s + i + mid_point];
                    let angle = -2. * self.direction.sign() * PI * (i as f32) / (pol_length as f32);
                    let phasor = Complex32::cexp(angle);
                    output[i + s] = a + phasor * b;
                    output[i + s + mid_point] = a - phasor * b;
                }
            }
        }
    }
}

// Utilities
pub fn reverse_bits(n: usize, bit_count: u32) -> usize {
    let mut num = n;
    let mut output = 0usize;
    for _ in 0..bit_count {
        output <<= 1;
        output |= num & 1;
        num >>= 1;
    }

    output
}

fn is_power_of_two(n: usize) -> bool {
    if n == 0 {
        return false;
    }

    let mut num = n;
    while num != 1 {
        if num % 2 != 0 {
            return false;
        }
        num /= 2;
    }
    return true;
}