// Implementation of raders's fft for prime sized ffts

use std::{f32::consts::PI, ops::Deref};

use super::mixed_radix;
use crate::{
    complex::Complex32,
    fft::{DFT, FFTDirection, create_fft, dft::NaiveDFT, is_prime, windows},
};

pub struct RaderFFT {
    permutations: Box<[usize]>,
    convolution_op: Box<[Complex32]>,
    staging_buffer: Box<[Complex32]>,
    inv_fft: Box<dyn DFT>,
    conv_fft: Box<dyn DFT>,

    size: usize,
}

impl DFT for RaderFFT {
    fn create(size: usize, direction: FFTDirection) -> Self
    where
        Self: Sized,
    {
        assert!(is_prime(size));
        let g = compute_prime_primitive_root(size);
        let permutations: Box<[usize]> = (0..(size - 1)).map(|i| exp_mod(g, i + 1, size)).collect();

        let mut conv_fft = create_fft(size - 1, FFTDirection::Forward);
        //let mut conv_fft = create_fft(size - 1);
        let mut convolution_op = vec![Complex32::zero(); size - 1];
        let conv_fft_input: Vec<Complex32> = (0..(size - 1))
            .map(|i| {
                Complex32::cexp(
                    -2. * direction.sign() * PI * (permutations[i] as f32) / (size as f32),
                )
            })
            .collect();
        conv_fft.execute(&conv_fft_input, &mut convolution_op, windows::rectangular);

        RaderFFT {
            permutations,
            convolution_op: convolution_op.into(),
            staging_buffer: vec![Complex32::zero(); size - 1].into(),
            inv_fft: create_fft(size - 1, FFTDirection::Inverse),
            conv_fft,

            size,
        }
    }

    fn execute(&mut self, input: &[Complex32], output: &mut [Complex32], window: fn(f32) -> f32) {
        // Compute fft of input signal
        for i in 0..(self.size - 1) {
            let k = self.permutations[self.size - 1 - i - 1];
            self.staging_buffer[i] = input[k] * window(k as f32 / (self.size as f32));
        }

        self.conv_fft
            .execute(&self.staging_buffer, output, windows::rectangular);

        for i in 0..(self.size - 1) {
            self.staging_buffer[i] = output[i] * self.convolution_op[i];
        }

        self.inv_fft
            .execute(&self.staging_buffer, output, windows::rectangular);

        for i in 0..(self.size - 1) {
            let k = self.permutations[i];
            self.staging_buffer[k - 1] = output[i];
        }

        output[0] = input[0] * window(0.0);
        for i in 0..(self.size - 1) {
            output[i + 1] = (self.staging_buffer[i] / (self.size - 1) as f32) + input[0];
            output[0] = output[0] + (input[i + 1] * window((i + 1) as f32 / self.size as f32));
        }
    }
}

pub fn compute_prime_primitive_root(n: usize) -> usize {
    assert!(is_prime(n));

    let phi = n - 1; // Euler's totient for n prime

    // Test all candidates
    for i in 1..(n + 1) {
        // Find multiplicative order of i
        let mut val = i;
        let mut order = 1;
        for j in 0..n {
            if val == 1 {
                break;
            }
            val = (val * i) % n;
            order += 1;
        }

        if order == phi {
            return i;
        }
    }

    unreachable!("Prime must have primitive root");
}

pub fn exp_mod(mut n: usize, mut exp: usize, m: usize) -> usize {
    if m == 1 {
        return 0;
    }

    n %= m;
    let mut r = 1;
    while exp > 0 {
        if exp % 2 == 1 {
            r = (r * n) % m;
        }

        n = (n * n) % m;
        exp >>= 1;
    }

    r
}
