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

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

use crate::{
    complex::Complex32,
    fft::{
        create_fft, is_prime, DFTAlgorithm, FFTDirection
    },
};

pub struct Rader2FFT {
    permutations: Box<[usize]>,
    convolution_operand: Box<[Complex32]>,

    convolution_fft_input: Box<[Complex32]>,
    convolution_ifft: Box<dyn DFTAlgorithm>,
    convolution_fft: Box<dyn DFTAlgorithm>,

    output: Box<[Complex32]>,

    sub_size: usize,
    size: usize,
}

impl DFTAlgorithm for Rader2FFT {
    fn create(size: usize, direction: FFTDirection) -> Self
    where
        Self: Sized,
    {
        assert!(is_prime(size));

        // Primitive root and its powers
        let g = compute_prime_primitive_root(size);
        let permutations: Box<[usize]> = (0..(size - 1)).map(|i| exp_mod(g, i + 1, size)).collect();
        let sub_size = next_pow2(2 * size - 3);

        // Compute fourrier transform of twiddle factors
        let twiddle_factors = (0..sub_size)
            .map(|i| {
                Complex32::cexp(
                    -2. * PI * direction.sign() * (permutations[i % (size - 1)] as f32)
                        / (size as f32),
                )
            })
            .collect::<Vec<Complex32>>();

        let mut convolution_fft = create_fft(sub_size, FFTDirection::Forward);
        convolution_fft.execute(&twiddle_factors);
        Rader2FFT {
            permutations,
            convolution_operand: convolution_fft.get_output().iter().copied().collect(),

            convolution_fft,
            convolution_ifft: create_fft(sub_size, FFTDirection::Inverse),
            convolution_fft_input: vec![Complex32::zero(); sub_size].into(),

            output: vec![Complex32::zero(); size].into(),
            size,
            sub_size,
        }
    }

    fn execute(&mut self, input: &[Complex32]) {
        // Compute fft of input signal

        self.convolution_fft_input[0] = input[self.permutations[self.size - 2]];
        for i in 0..(self.sub_size - self.size + 1) {
            self.convolution_fft_input[i + 1] = Complex32::zero();
        }
        for i in 1..(self.size - 1) {
            // reverse sequence
            let k = self.permutations[self.size - 1 - i - 1];
            self.convolution_fft_input[i + self.sub_size - self.size + 1] = input[k];
        }

        self.convolution_fft.execute(&self.convolution_fft_input);

        // Compute convolution by multiplying in freq domain
        for i in 0..self.sub_size {
            // Using output as staging buffer
            self.convolution_fft_input[i] =
                self.convolution_fft.get_output()[i] * self.convolution_operand[i];
        }

        self.convolution_ifft.execute(&self.convolution_fft_input);

        self.output[0] = Complex32::zero();
        for x in input {
            self.output[0] = self.output[0] + *x;
        }

        for i in 0..(self.size - 1) {
            // Actually compute the output
            let k = self.permutations[i];
            self.output[k] =
                (self.convolution_ifft.get_output()[i] / (self.sub_size) as f32) + input[0];
        }
    }

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

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 _ 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
}

pub fn next_pow2(mut n: usize) -> usize {
    if n.count_ones() == 1 {
        n
    } else {
        let mut p = 0;
        while n > 0 {
            n >>= 1;
            p += 1;
        }
        1 << p
    }
}