oxifft 0.1.4

//! Rader's algorithm for prime-size DFT.
//!
//! Rader's algorithm converts a prime-size DFT into a cyclic convolution
//! of size p-1, which can then be computed using FFT if p-1 has good factors.
//!
//! For a prime p, the DFT can be rewritten using the primitive root g:
//! X[g^k] = x[0] + Σ_{j=0}^{p-2} x[g^{-j}] * W_p^{g^{k-j}}
//!
//! The summation is a cyclic convolution.
//!
//! Time complexity: O(p log p) for prime p (when p-1 has good factors)
//! Space complexity: O(p)

use crate::dft::problem::Sign;
use crate::kernel::{is_prime, primitive_root};
use crate::kernel::{Complex, Float};
use crate::prelude::*;

use super::bluestein::BluesteinSolver;

/// Rader's algorithm solver for prime sizes.
///
/// This solver uses the fact that for prime p, the non-zero indices
/// form a cyclic group under multiplication modulo p, generated by
/// a primitive root g.
///
/// Pre-allocates work buffers to avoid per-execution allocations.
/// Uses `Mutex` for thread-safe interior mutability with `try_lock()` fallback.
pub struct RaderSolver<T: Float> {
    /// Prime size
    p: usize,
    /// Primitive root
    g: usize,
    /// Powers of g: g^0, g^1, ..., g^(p-2) mod p
    g_powers: Vec<usize>,
    /// Inverse powers: g^(-0), g^(-1), ..., g^(-(p-2)) mod p (same as g^(p-1-k))
    g_inv_powers: Vec<usize>,
    /// Precomputed FFT of forward twiddle factors
    twiddle_fft_fwd: Vec<Complex<T>>,
    /// Precomputed FFT of backward twiddle factors
    twiddle_fft_bwd: Vec<Complex<T>>,
    /// Bluestein solver for the convolution (size p-1)
    conv_solver: BluesteinSolver<T>,
    /// Pre-allocated work buffer for reordered input
    #[cfg(feature = "std")]
    work_a: Mutex<Vec<Complex<T>>>,
    /// Pre-allocated work buffer for FFT of input
    #[cfg(feature = "std")]
    work_a_fft: Mutex<Vec<Complex<T>>>,
    /// Pre-allocated work buffer for convolution result
    #[cfg(feature = "std")]
    work_conv: Mutex<Vec<Complex<T>>>,
}

impl<T: Float> RaderSolver<T> {
    /// Create a new Rader solver for prime size p.
    ///
    /// Returns `None` if p is not prime or p < 3.
    #[must_use]
    pub fn new(p: usize) -> Option<Self> {
        if p < 3 || !is_prime(p) {
            return None;
        }

        let g = primitive_root(p)?;
        let n = p - 1; // Convolution size

        // Compute powers of g
        let mut g_powers = Vec::with_capacity(n);
        let mut g_inv_powers = Vec::with_capacity(n);

        let mut power = 1usize;
        for _ in 0..n {
            g_powers.push(power);
            power = (power * g) % p;
        }

        // g^(-k) = g^(p-1-k) mod p (since g^(p-1) = 1 mod p by Fermat's little theorem)
        for k in 0..n {
            g_inv_powers.push(g_powers[(n - k) % n]);
        }

        let conv_solver = BluesteinSolver::new(n);

        // Precompute forward twiddle factors: W_p^(g^k) for k = 0..p-2
        // W_p = e^(-2πi/p) for forward transform
        let mut twiddles_fwd = Vec::with_capacity(n);
        for k in 0..n {
            let exp = g_powers[k];
            let angle = -<T as Float>::TWO_PI * T::from_usize(exp) / T::from_usize(p);
            twiddles_fwd.push(Complex::cis(angle));
        }

        // FFT of forward twiddles
        let mut twiddle_fft_fwd = vec![Complex::zero(); n];
        conv_solver.execute(&twiddles_fwd, &mut twiddle_fft_fwd, Sign::Forward);

        // Precompute backward twiddle factors: W_p^(-g^k) for k = 0..p-2
        // W_p^(-1) = e^(+2πi/p) for backward transform
        let mut twiddles_bwd = Vec::with_capacity(n);
        for k in 0..n {
            let exp = g_powers[k];
            let angle = <T as Float>::TWO_PI * T::from_usize(exp) / T::from_usize(p);
            twiddles_bwd.push(Complex::cis(angle));
        }

        // FFT of backward twiddles
        let mut twiddle_fft_bwd = vec![Complex::zero(); n];
        conv_solver.execute(&twiddles_bwd, &mut twiddle_fft_bwd, Sign::Forward);

        Some(Self {
            p,
            g,
            g_powers,
            g_inv_powers,
            twiddle_fft_fwd,
            twiddle_fft_bwd,
            conv_solver,
            #[cfg(feature = "std")]
            work_a: Mutex::new(vec![Complex::zero(); n]),
            #[cfg(feature = "std")]
            work_a_fft: Mutex::new(vec![Complex::zero(); n]),
            #[cfg(feature = "std")]
            work_conv: Mutex::new(vec![Complex::zero(); n]),
        })
    }

    /// Solver name.
    #[must_use]
    pub fn name(&self) -> &'static str {
        "dft-rader"
    }

    /// Get the prime size.
    #[must_use]
    pub fn size(&self) -> usize {
        self.p
    }

    /// Get the primitive root used.
    #[must_use]
    #[allow(dead_code)]
    pub fn primitive_root(&self) -> usize {
        self.g
    }

    /// Check if this solver is applicable (p is prime and >= 3).
    #[must_use]
    pub fn applicable(p: usize) -> bool {
        p >= 3 && is_prime(p)
    }

    /// Execute Rader's FFT with provided work buffers.
    fn execute_with_buffers(
        &self,
        input: &[Complex<T>],
        output: &mut [Complex<T>],
        sign: Sign,
        a: &mut [Complex<T>],
        a_fft: &mut [Complex<T>],
        conv: &mut [Complex<T>],
    ) {
        let p = self.p;
        let n = p - 1;

        // Step 1: Compute X[0] = sum of all inputs
        let mut sum = Complex::zero();
        for x in input {
            sum = sum + *x;
        }

        // Step 2: Reorder input according to g^(-j)
        for j in 0..n {
            a[j] = input[self.g_inv_powers[j]];
        }

        // Step 3: FFT of reordered input
        self.conv_solver.execute(a, a_fft, Sign::Forward);

        // Step 4: Pointwise multiply with appropriate twiddle FFT
        let twiddle_fft = match sign {
            Sign::Forward => &self.twiddle_fft_fwd,
            Sign::Backward => &self.twiddle_fft_bwd,
        };

        for i in 0..n {
            a_fft[i] = a_fft[i] * twiddle_fft[i];
        }

        // Step 5: IFFT to get convolution result
        self.conv_solver.execute(a_fft, conv, Sign::Backward);

        // Normalize IFFT
        let n_inv = T::ONE / T::from_usize(n);
        for x in conv.iter_mut().take(n) {
            *x = *x * n_inv;
        }

        // Step 6: Compute output
        // X[0] = sum (already computed)
        output[0] = sum;

        // X[g^k] = x[0] + conv[k] for k = 0..p-2
        for k in 0..n {
            let idx = self.g_powers[k];
            output[idx] = input[0] + conv[k];
        }
    }

    /// Execute Rader's FFT algorithm.
    ///
    /// Uses pre-allocated work buffers when available (single-threaded case).
    /// Falls back to fresh allocation when buffers are locked (parallel execution).
    #[cfg(feature = "std")]
    pub fn execute(&self, input: &[Complex<T>], output: &mut [Complex<T>], sign: Sign) {
        let p = self.p;
        let n = p - 1;

        debug_assert_eq!(input.len(), p);
        debug_assert_eq!(output.len(), p);

        // Try to acquire all three locks. If any fails, allocate fresh buffers.
        let a_guard = self.work_a.try_lock();
        let a_fft_guard = self.work_a_fft.try_lock();
        let conv_guard = self.work_conv.try_lock();

        if let (Ok(mut a), Ok(mut a_fft), Ok(mut conv)) = (a_guard, a_fft_guard, conv_guard) {
            // Use pre-allocated buffers
            self.execute_with_buffers(input, output, sign, &mut a, &mut a_fft, &mut conv);
        } else {
            // Fallback: allocate fresh buffers (parallel execution case)
            let mut a = vec![Complex::zero(); n];
            let mut a_fft = vec![Complex::zero(); n];
            let mut conv = vec![Complex::zero(); n];
            self.execute_with_buffers(input, output, sign, &mut a, &mut a_fft, &mut conv);
        }
    }

    /// Execute Rader's FFT algorithm (no_std version - always allocates).
    #[cfg(not(feature = "std"))]
    pub fn execute(&self, input: &[Complex<T>], output: &mut [Complex<T>], sign: Sign) {
        let p = self.p;
        let n = p - 1;

        debug_assert_eq!(input.len(), p);
        debug_assert_eq!(output.len(), p);

        // no_std: always allocate fresh buffers
        let mut a = vec![Complex::zero(); n];
        let mut a_fft = vec![Complex::zero(); n];
        let mut conv = vec![Complex::zero(); n];
        self.execute_with_buffers(input, output, sign, &mut a, &mut a_fft, &mut conv);
    }

    /// Execute Rader's FFT in-place.
    pub fn execute_inplace(&self, data: &mut [Complex<T>], sign: Sign) {
        let p = self.p;
        debug_assert_eq!(data.len(), p);

        // Need temporary storage
        let input: Vec<Complex<T>> = data.to_vec();
        self.execute(&input, data, sign);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::dft::solvers::direct::DirectSolver;

    fn approx_eq(a: f64, b: f64, eps: f64) -> bool {
        (a - b).abs() < eps
    }

    fn complex_approx_eq(a: Complex<f64>, b: Complex<f64>, eps: f64) -> bool {
        approx_eq(a.re, b.re, eps) && approx_eq(a.im, b.im, eps)
    }

    #[test]
    fn test_rader_applicable() {
        assert!(!RaderSolver::<f64>::applicable(0));
        assert!(!RaderSolver::<f64>::applicable(1));
        assert!(!RaderSolver::<f64>::applicable(2));
        assert!(RaderSolver::<f64>::applicable(3));
        assert!(!RaderSolver::<f64>::applicable(4));
        assert!(RaderSolver::<f64>::applicable(5));
        assert!(!RaderSolver::<f64>::applicable(6));
        assert!(RaderSolver::<f64>::applicable(7));
        assert!(RaderSolver::<f64>::applicable(11));
        assert!(RaderSolver::<f64>::applicable(13));
    }

    #[test]
    fn test_rader_size_3() {
        let input: Vec<Complex<f64>> = (0..3).map(|i| Complex::new(f64::from(i), 0.0)).collect();
        let mut output_rader = vec![Complex::zero(); 3];
        let mut output_direct = vec![Complex::zero(); 3];

        RaderSolver::new(3)
            .unwrap()
            .execute(&input, &mut output_rader, Sign::Forward);
        DirectSolver::new().execute(&input, &mut output_direct, Sign::Forward);

        for (a, b) in output_rader.iter().zip(output_direct.iter()) {
            assert!(complex_approx_eq(*a, *b, 1e-9));
        }
    }

    #[test]
    fn test_rader_size_5() {
        let input: Vec<Complex<f64>> = (0..5)
            .map(|i| Complex::new(f64::from(i).sin(), f64::from(i).cos()))
            .collect();
        let mut output_rader = vec![Complex::zero(); 5];
        let mut output_direct = vec![Complex::zero(); 5];

        RaderSolver::new(5)
            .unwrap()
            .execute(&input, &mut output_rader, Sign::Forward);
        DirectSolver::new().execute(&input, &mut output_direct, Sign::Forward);

        for (a, b) in output_rader.iter().zip(output_direct.iter()) {
            assert!(complex_approx_eq(*a, *b, 1e-9));
        }
    }

    #[test]
    fn test_rader_size_7() {
        let input: Vec<Complex<f64>> = (0..7)
            .map(|i| Complex::new(f64::from(i), f64::from(i) * 0.5))
            .collect();
        let mut output_rader = vec![Complex::zero(); 7];
        let mut output_direct = vec![Complex::zero(); 7];

        RaderSolver::new(7)
            .unwrap()
            .execute(&input, &mut output_rader, Sign::Forward);
        DirectSolver::new().execute(&input, &mut output_direct, Sign::Forward);

        for (a, b) in output_rader.iter().zip(output_direct.iter()) {
            assert!(complex_approx_eq(*a, *b, 1e-9));
        }
    }

    #[test]
    fn test_rader_size_13() {
        let input: Vec<Complex<f64>> = (0..13)
            .map(|i| Complex::new(f64::from(i).sin(), f64::from(i).cos()))
            .collect();
        let mut output_rader = vec![Complex::zero(); 13];
        let mut output_direct = vec![Complex::zero(); 13];

        RaderSolver::new(13)
            .unwrap()
            .execute(&input, &mut output_rader, Sign::Forward);
        DirectSolver::new().execute(&input, &mut output_direct, Sign::Forward);

        for (a, b) in output_rader.iter().zip(output_direct.iter()) {
            assert!(complex_approx_eq(*a, *b, 1e-8));
        }
    }

    #[test]
    fn test_rader_inverse_recovers_input() {
        let original: Vec<Complex<f64>> = (0..11)
            .map(|i| Complex::new(f64::from(i).sin(), f64::from(i).cos()))
            .collect();
        let mut transformed = vec![Complex::zero(); 11];
        let mut recovered = vec![Complex::zero(); 11];

        let solver = RaderSolver::new(11).unwrap();
        solver.execute(&original, &mut transformed, Sign::Forward);
        solver.execute(&transformed, &mut recovered, Sign::Backward);

        // Normalize
        let n = original.len() as f64;
        for x in &mut recovered {
            *x = *x / n;
        }

        for (a, b) in original.iter().zip(recovered.iter()) {
            assert!(complex_approx_eq(*a, *b, 1e-9));
        }
    }

    #[test]
    fn test_rader_inplace() {
        let original: Vec<Complex<f64>> = (0..7).map(|i| Complex::new(f64::from(i), 0.0)).collect();

        // Out-of-place reference
        let mut out_of_place = vec![Complex::zero(); 7];
        let solver = RaderSolver::new(7).unwrap();
        solver.execute(&original, &mut out_of_place, Sign::Forward);

        // In-place
        let mut in_place = original;
        solver.execute_inplace(&mut in_place, Sign::Forward);

        for (a, b) in out_of_place.iter().zip(in_place.iter()) {
            assert!(complex_approx_eq(*a, *b, 1e-10));
        }
    }
}