pub mod dft;
pub mod mixed_radix;
pub mod rader;
pub mod rader2;
pub mod radix2;

use crate::{
    complex::Complex32,
    fft::{
        dft::NaiveDFT, mixed_radix::MixedRadixFFT, rader2::Rader2FFT,
        radix2::Radix2FFT,
    }, windows,
};

pub type FFTWindow = fn(f32) -> f32;

#[derive(Copy, Clone)]
pub enum FFTDirection {
    Forward,
    Inverse,
}

impl FFTDirection {
    fn sign(&self) -> f32 {
        match self {
            FFTDirection::Forward => 1.,
            FFTDirection::Inverse => -1.,
        }
    }
}

pub trait DFTAlgorithm {
    fn create(size: usize, direction: FFTDirection) -> Self
    where
        Self: Sized;

    fn execute(&mut self, input: &[Complex32]);

    fn get_output(&self) -> &[Complex32];
}

pub fn create_fft(size: usize, direction: FFTDirection) -> Box<dyn DFTAlgorithm> {
    if size <= 16 {
        //println!("Naive {size}");
        return Box::new(NaiveDFT::create(size, direction));
    }
    if size.count_ones() == 1 {
        //println!("Radix 2 {size}");
        return Box::new(Radix2FFT::create(size, direction));
    }

    if is_prime(size) {
        //println!("Prime rader {size}");
        return Box::new(Rader2FFT::create(size, direction));
        //return Box::new(NaiveDFT::create(size, direction));
    }

    //println!("Mixed radix {size}");
    Box::new(MixedRadixFFT::create(size, direction))
    //Box::new(NaiveDFT::create(size, direction))
}

pub struct FFT {
    fft: Box<dyn DFTAlgorithm>,
    size: usize,
    window: FFTWindow,
    input_buffer: Box<[Complex32]>,
}

impl FFT {
    pub fn new(size: usize, window: FFTWindow) -> FFT {
        FFT {
            fft: create_fft(size, FFTDirection::Forward),
            window,
            size,
            input_buffer: vec![Complex32::zero(); size].into(),
        }
    }

    pub fn new_inv(size: usize) -> FFT {
        FFT {
            fft: create_fft(size, FFTDirection::Inverse),
            window: windows::rectangular,
            size,
            input_buffer: vec![Complex32::zero(); size].into(),
        }
    }

    pub fn execute(&mut self, input: &[Complex32]) {
        self.input_buffer
            .iter_mut()
            .zip(input.iter())
            .enumerate()
            .for_each(|(i, (x, y))| *x = *y * (self.window)(i as f32 / self.size as f32));

        self.fft.execute(&self.input_buffer);
    }

    pub fn get_output(&self) -> &[Complex32] {
        self.fft.get_output()
    }
}

// Utilities
pub fn prime_factors(n: usize) -> Vec<usize> {
    let mut factors = vec![];

    let mut num = n;
    // Divide num successively
    while num != 1 {
        // Try divisors from 2 up to n (included)
        for i in 2..n + 1 {
            // if i divides num, it is a prime factor (if it wasn't, then i would have prime
            // factors that would divide into num before i)
            if num % i == 0 {
                factors.push(i);
                num /= i;
                break;
            }
        }
    }

    // If n = 1 then it does not have any prime factors
    // The prime factor decomposition theorem states that any integer
    // greater than TWO has a unique decomposition

    factors
}

pub fn is_prime(n: usize) -> bool {
    prime_factors(n).len() == 1
}
