// To perform a mixed radix cooley tuckey fft

use std::f32::consts::PI;

use crate::{
    complex::Complex32,
    fft::{DFTAlgorithm, FFTDirection, create_fft, dft::NaiveDFT, prime_factors, windows},
};

pub struct MixedRadixFFT {
    //size: usize, size is implicitely stored in p and q
    p: usize,
    q: usize,
    twiddle_factors: Box<[Complex32]>,

    qfft: Box<dyn DFTAlgorithm>,
    pfft: Box<dyn DFTAlgorithm>,

    staging_buffer: Box<[Complex32]>,
    pfft_input: Box<[Complex32]>,
    qfft_input: Box<[Complex32]>,
    output: Box<[Complex32]>,
}

impl DFTAlgorithm for MixedRadixFFT {
    fn create(size: usize, direction: FFTDirection) -> Self {
        let q = decide_radix_factor(size);
        let p = size / q;
        let qfft = create_fft(q, direction);
        let pfft = create_fft(p, direction);

        //let qfft = Box::new(NaiveDFT::create(q, direction));
        //let pfft = Box::new(NaiveDFT::create(p, direction));

        MixedRadixFFT {
            twiddle_factors: compute_twiddle_factors(size, direction),
            qfft,
            pfft,

            staging_buffer: vec![Complex32::zero(); size].into_boxed_slice(),
            pfft_input: vec![Complex32::zero(); p].into_boxed_slice(),
            qfft_input: vec![Complex32::zero(); q].into_boxed_slice(),
            output: vec![Complex32::zero(); size].into_boxed_slice(),
            p,
            q,
        }
    }

    fn execute(&mut self, input: &[Complex32]) {
        // Perform p ffts of size q
        for k0 in 0..self.p {
            // Copy samples into input buffer
            for k1 in 0..self.q {
                let k = k1 * self.p + k0;
                // Use output as staging buffer
                self.qfft_input[k1] = input[k];
            }

            self.qfft.execute(&self.qfft_input);

            for j0 in 0..self.q {
                // "Store j0'th of k0'th fft into staging buffer"
                self.staging_buffer[j0 * self.p + k0] =
                    self.qfft.get_output()[j0] * self.twiddle_factors[j0 * k0];
            }
        }

        // Perform q ffts of size p
        for j0 in 0..self.q {
            // Copy input
            for k0 in 0..self.p {
                // Use output as staging buffer
                self.pfft_input[k0] = self.staging_buffer[j0 * self.p + k0];
            }

            self.pfft.execute(&self.pfft_input);

            // Actually compute final output
            for j1 in 0..self.p {
                self.output[j1 * self.q + j0] = self.pfft.get_output()[j1];
            }
        }
    }

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

fn compute_twiddle_factors(size: usize, direction: FFTDirection) -> Box<[Complex32]> {
    let mut factors = vec![Complex32::zero(); size].into_boxed_slice();

    for i in 0..size {
        factors[i] = Complex32::cexp(-2. * direction.sign() * PI * i as f32 / (size as f32));
    }
    factors
}

// This will decide on a good factor to use for the mixed radix fft
fn decide_radix_factor(n: usize) -> usize {
    let factors = prime_factors(n);
    //factors.iter().for_each(|a| print!(" {a} "));
    let two_count = factors.iter().take_while(|i| **i == 2).count();

    // If there is a lot of two, we can use them as q factor to be able to use radix2 later on
    if two_count > 0 {
        return 1 << two_count;
    }

    // Otherwise take next big prime
    *factors.get(two_count).unwrap()
}
