use crate::util::*;
use num::complex::Complex;
use num::traits::{Float, FloatConst};

fn fft_butterfly_radix_2_unsafe<T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
    stride: usize,
    big_n: usize,
    twiddles: &[Complex<T>],
    _: &ExecutionStage<T>,
) {
    let input_ptr = input.as_ptr();
    let output_ptr = output.as_mut_ptr();
    for start_idx in 0..stride {
        for k in 0..big_n / 2 {
            unsafe {
                // Collect inputs.
                let p = *input_ptr.add(start_idx + 2 * k * stride);
                let q = *input_ptr.add(start_idx + (2 * k + 1) * stride);
                // Collect relevant twiddles.
                let ot = twiddles.get_unchecked(k * stride);
                // Combine.
                let r = q * ot;
                *output_ptr.add(start_idx + k * stride) = p + r;
                *output_ptr.add(start_idx + (k + big_n / 2) * stride) = p - r;
            }
        }
    }
}

fn fft_butterfly_radix_2_s0_unsafe<T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
    _: &Matrix<Complex<T>>,
) {
    let input_ptr = input.as_ptr();
    let output_ptr = output.as_mut_ptr();
    let stride = input.len() / 2;
    let big_n = 2;
    for start_idx in 0..stride {
        for k in 0..big_n / 2 {
            unsafe {
                // Collect inputs.
                let p = *input_ptr.add(start_idx + 2 * k * stride);
                let q = *input_ptr.add(start_idx + (2 * k + 1) * stride);
                *output_ptr.add(start_idx + k * stride) = p + q;
                *output_ptr.add(start_idx + (k + big_n / 2) * stride) = p - q;
            }
        }
    }
}

fn radix_3_dft<T: Float + FloatConst>(
    input: &[Complex<T>; 3],
    _: &Matrix<Complex<T>>,
) -> [Complex<T>; 3] {
    const N: usize = 3;
    let zero = Complex::new(T::zero(), T::zero());
    let mut result = [zero; N];
    let half = T::from(0.5).unwrap();
    let sqrt_3_over_2 =
        T::from(0.8660254037844386467637231707529361834714026269051903140279034897).unwrap();
    let sum = input[1] + input[2];
    let diff = input[1] - input[2];
    result[0] = input[0] + input[1] + input[2];
    result[1] = input[0] - sum.scale(half) - mul_i(diff.scale(sqrt_3_over_2));
    result[2] = input[0] - sum.scale(half) + mul_i(diff.scale(sqrt_3_over_2));
    result
}

fn fft_butterfly_radix_3_unsafe<T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
    stride: usize,
    big_n: usize,
    twiddles: &[Complex<T>],
    stage: &ExecutionStage<T>,
) {
    let input_ptr = input.as_ptr();
    let output_ptr = output.as_mut_ptr();
    const N: usize = 3;
    for start_idx in 0..stride {
        for k in 0..big_n / N {
            unsafe {
                // Collect inputs.
                let inputs: [Complex<T>; N] =
                    std::array::from_fn(|idx| *input_ptr.add(start_idx + (N * k + idx) * stride));
                // Collect relevant twiddles.
                let twiddles: [&Complex<T>; N] =
                    std::array::from_fn(|idx| twiddles.get_unchecked(idx * k * stride));
                // Precompute a term used in the inner loop.
                let it: [Complex<T>; N] = std::array::from_fn(|idx| {
                    if idx == 0 {
                        inputs[idx]
                    } else {
                        inputs[idx] * twiddles[idx]
                    }
                });
                let result = radix_3_dft(&it, &stage.outer_terms);
                for row in 0..N {
                    *output_ptr.add(start_idx + (k + row * big_n / N) * stride) = result[row];
                }
            }
        }
    }
}

fn fft_butterfly_radix_3_s0_unsafe<T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
    outer_terms: &Matrix<Complex<T>>,
) {
    const N: usize = 3;
    let input_ptr = input.as_ptr();
    let output_ptr = output.as_mut_ptr();
    let stride = input.len() / N;
    let big_n = N;
    for start_idx in 0..stride {
        for k in 0..big_n / N {
            unsafe {
                // Collect inputs.
                let inputs: [Complex<T>; N] =
                    std::array::from_fn(|idx| *input_ptr.add(start_idx + (N * k + idx) * stride));
                let result = radix_3_dft(&inputs, outer_terms);
                for row in 0..N {
                    *output_ptr.add(start_idx + (k + row * big_n / N) * stride) = result[row];
                }
            }
        }
    }
}

fn fft_butterfly_radix_9_unsafe<T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
    stride: usize,
    big_n: usize,
    twiddles: &[Complex<T>],
    stage: &ExecutionStage<T>,
) {
    let input_ptr = input.as_ptr();
    let output_ptr = output.as_mut_ptr();
    const N: usize = 9;
    for start_idx in 0..stride {
        for k in 0..big_n / N {
            unsafe {
                let inputs: [Complex<T>; N] =
                    std::array::from_fn(|idx| *input_ptr.add(start_idx + (N * k + idx) * stride));
                let twiddles: [&Complex<T>; N] =
                    std::array::from_fn(|idx| twiddles.get_unchecked(idx * k * stride));
                let it: [Complex<T>; N] = std::array::from_fn(|idx| {
                    if idx == 0 {
                        inputs[idx]
                    } else {
                        inputs[idx] * twiddles[idx]
                    }
                });

                let b00 = [it[0], it[3], it[6]];
                let b01 = [it[1], it[4], it[7]];
                let b02 = [it[2], it[5], it[8]];

                let b00_out = radix_3_dft(&b00, &stage.outer_terms);
                let b01_out = radix_3_dft(&b01, &stage.outer_terms);
                let b02_out = radix_3_dft(&b02, &stage.outer_terms);

                // Recall the general way that outer terms work:
                // First row is all 1s.
                // Second row rotates 1/r clockwise.
                // Third row rotates 2/r clockwise.
                // outer_terms already takes into account the way rotations
                // accelerate each row, so we just have to index into the right
                // row/col.
                let b10 = [b00_out[0], b01_out[0], b02_out[0]];
                let b11 = [
                    b00_out[1],
                    b01_out[1] * stage.outer_terms.data.get_unchecked(9 + 1),
                    b02_out[1] * stage.outer_terms.data.get_unchecked(9 + 2),
                ];
                let b12 = [
                    b00_out[2],
                    b01_out[2] * stage.outer_terms.data.get_unchecked(18 + 1),
                    b02_out[2] * stage.outer_terms.data.get_unchecked(18 + 2),
                ];

                let r00 = radix_3_dft(&b10, &stage.outer_terms);
                let r01 = radix_3_dft(&b11, &stage.outer_terms);
                let r02 = radix_3_dft(&b12, &stage.outer_terms);

                let results = [
                    r00[0], r01[0], r02[0], r00[1], r01[1], r02[1], r00[2], r01[2], r02[2],
                ];
                for row in 0..N {
                    *output_ptr.add(start_idx + (k + row * big_n / N) * stride) = results[row];
                }
            }
        }
    }
}

fn fft_butterfly_radix_9_s0_unsafe<T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
    outer_terms: &Matrix<Complex<T>>,
) {
    let input_ptr = input.as_ptr();
    let output_ptr = output.as_mut_ptr();
    const N: usize = 9;
    let stride = input.len() / N;
    let big_n = N;
    for start_idx in 0..stride {
        for k in 0..big_n / N {
            unsafe {
                let inputs: [Complex<T>; N] =
                    std::array::from_fn(|idx| *input_ptr.add(start_idx + (N * k + idx) * stride));

                let b00 = [inputs[0], inputs[3], inputs[6]];
                let b01 = [inputs[1], inputs[4], inputs[7]];
                let b02 = [inputs[2], inputs[5], inputs[8]];

                let b00_out = radix_3_dft(&b00, &outer_terms);
                let b01_out = radix_3_dft(&b01, &outer_terms);
                let b02_out = radix_3_dft(&b02, &outer_terms);

                // Recall the general way that outer terms work:
                // First row is all 1s.
                // Second row rotates 1/r clockwise.
                // Third row rotates 2/r clockwise.
                let b10 = [b00_out[0], b01_out[0], b02_out[0]];
                let b11 = [
                    b00_out[1],
                    b01_out[1] * outer_terms.data.get_unchecked(9 + 1),
                    b02_out[1] * outer_terms.data.get_unchecked(9 + 2),
                ];
                let b12 = [
                    b00_out[2],
                    b01_out[2] * outer_terms.data.get_unchecked(18 + 1),
                    b02_out[2] * outer_terms.data.get_unchecked(18 + 2),
                ];

                let r00 = radix_3_dft(&b10, &outer_terms);
                let r01 = radix_3_dft(&b11, &outer_terms);
                let r02 = radix_3_dft(&b12, &outer_terms);

                let results = [
                    r00[0], r01[0], r02[0], r00[1], r01[1], r02[1], r00[2], r01[2], r02[2],
                ];
                for row in 0..N {
                    *output_ptr.add(start_idx + (k + row * big_n / N) * stride) = results[row];
                }
            }
        }
    }
}

fn odd_radix_dft<const N: usize, T: Float + FloatConst>(
    input: &[Complex<T>; N],
    outer_terms: &Matrix<Complex<T>>,
) -> [Complex<T>; N] {
    let zero = Complex::new(T::zero(), T::zero());
    let mut result = [zero; N];
    // Precompute sums and differences.
    let mut sums = [zero; N];
    let mut diffs = [zero; N];
    for col in 1..N / 2 + 1 {
        sums[col] = input[col] + input[N - col];
        diffs[col] = input[col] - input[N - col];
    }

    // 0,0 term is never indexed, so compute it out of the loop.
    result[0] = input[0];

    // 0th row is never indexed, so compute it out of the loop.
    for col in 1..N / 2 + 1 {
        // outer_terms is always 1 along the first row and column, so we can skip the multiplication.
        result[0] = result[0] + sums[col];
    }

    for row in 1..N / 2 + 1 {
        // Column zero is not indexed below, so precompute here. Recall the
        // outer_terms is 1 along the first column.
        //
        // Per our identity, we split this sum into a "real" coefficient and an
        // "imaginary" coefficient. Note that they are in fact both complex-valued.
        let mut re_part = input[0];
        let mut im_part = zero;

        for col in 1..N / 2 + 1 {
            unsafe {
                let w = outer_terms.data.get_unchecked(row * N + col);
                re_part = re_part + sums[col].scale(w.re);
                im_part = im_part + diffs[col].scale(w.im);
            }
        }

        // Then apply the identity to get both rows.
        let i_im_part = mul_i(im_part);
        result[row] = re_part + i_im_part;
        result[N - row] = re_part - i_im_part;
    }
    result
}

fn fft_butterfly_radix_p_unsafe<const N: usize, T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
    stride: usize,
    big_n: usize,
    twiddles: &[Complex<T>],
    stage: &ExecutionStage<T>,
) {
    let input_ptr = input.as_ptr();
    let output_ptr = output.as_mut_ptr();
    for start_idx in 0..stride {
        for k in 0..big_n / N {
            unsafe {
                // Collect inputs.
                let inputs: [Complex<T>; N] =
                    std::array::from_fn(|idx| *input_ptr.add(start_idx + (N * k + idx) * stride));
                // Collect relevant twiddles.
                let twiddles: [&Complex<T>; N] =
                    std::array::from_fn(|idx| twiddles.get_unchecked(idx * k * stride));
                // Precompute a term used in the inner loop.
                let it: [Complex<T>; N] = std::array::from_fn(|idx| {
                    if idx == 0 {
                        inputs[idx]
                    } else {
                        inputs[idx] * twiddles[idx]
                    }
                });
                let result = odd_radix_dft(&it, &stage.outer_terms);
                for row in 0..N {
                    *output_ptr.add(start_idx + (k + row * big_n / N) * stride) = result[row];
                }
            }
        }
    }
}

fn fft_butterfly_radix_p_s0_unsafe<const N: usize, T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
    outer_terms: &Matrix<Complex<T>>,
) {
    let input_ptr = input.as_ptr();
    let output_ptr = output.as_mut_ptr();
    let stride = input.len() / N;
    let big_n = N;
    for start_idx in 0..stride {
        for k in 0..big_n / N {
            unsafe {
                // Collect inputs.
                let inputs: [Complex<T>; N] =
                    std::array::from_fn(|idx| *input_ptr.add(start_idx + (N * k + idx) * stride));
                // Precompute a term used in the inner loop.
                let result = odd_radix_dft(&inputs, outer_terms);
                for row in 0..N {
                    *output_ptr.add(start_idx + (k + row * big_n / N) * stride) = result[row];
                }
            }
        }
    }
}

#[inline(always)]
fn mul_i<T: Float + FloatConst>(x: Complex<T>) -> Complex<T> {
    Complex::new(-x.im, x.re)
}

#[inline(always)]
fn mul_ni<T: Float + FloatConst>(x: Complex<T>) -> Complex<T> {
    Complex::new(x.im, -x.re)
}

fn fft_butterfly_radix_4_unsafe<T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
    stride: usize,
    big_n: usize,
    twiddles: &[Complex<T>],
    _: &ExecutionStage<T>,
) {
    let input_ptr = input.as_ptr();
    let output_ptr = output.as_mut_ptr();
    for start_idx in 0..stride {
        for k in 0..big_n / 4 {
            unsafe {
                // Collect inputs.
                let i0 = *input_ptr.add(start_idx + 4 * k * stride);
                let i1 = *input_ptr.add(start_idx + (4 * k + 1) * stride);
                let i2 = *input_ptr.add(start_idx + (4 * k + 2) * stride);
                let i3 = *input_ptr.add(start_idx + (4 * k + 3) * stride);
                // Collect relevant twiddles.
                let ot1 = twiddles.get_unchecked(1 * k * stride);
                let ot2 = twiddles.get_unchecked(2 * k * stride);
                let ot3 = twiddles.get_unchecked(3 * k * stride);

                let a = i0;
                let b = ot1 * i1;
                let c = ot2 * i2;
                let d = ot3 * i3;

                // To derive this, write the output assignments in terms of
                // a/b/c/d, then factor out!
                let ac_sum = a + c;
                let ac_diff = a - c;
                let bd_sum = b + d;
                let bd_diff_ni = mul_ni(b - d);

                *output_ptr.add(start_idx + k * stride) = ac_sum + bd_sum;
                *output_ptr.add(start_idx + (k + big_n / 4) * stride) = ac_diff + bd_diff_ni;
                *output_ptr.add(start_idx + (k + big_n / 2) * stride) = ac_sum - bd_sum;
                *output_ptr.add(start_idx + (k + 3 * big_n / 4) * stride) = ac_diff - bd_diff_ni;
            }
        }
    }
}

fn fft_butterfly_radix_4_s0_unsafe<T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
    _: &Matrix<Complex<T>>,
) {
    let stride = input.len() / 4;
    let big_n = 4;
    let input_ptr = input.as_ptr();
    let output_ptr = output.as_mut_ptr();
    for start_idx in 0..stride {
        for k in 0..big_n / 4 {
            unsafe {
                // Collect inputs.
                let i0 = *input_ptr.add(start_idx + 4 * k * stride);
                let i1 = *input_ptr.add(start_idx + (4 * k + 1) * stride);
                let i2 = *input_ptr.add(start_idx + (4 * k + 2) * stride);
                let i3 = *input_ptr.add(start_idx + (4 * k + 3) * stride);

                let a = i0;
                let b = i1;
                let c = i2;
                let d = i3;

                // To derive this, write the output assignments in terms of
                // a/b/c/d, then factor out!
                let ac_sum = a + c;
                let ac_diff = a - c;
                let bd_sum = b + d;
                let bd_diff_ni = mul_ni(b - d);

                *output_ptr.add(start_idx + k * stride) = ac_sum + bd_sum;
                *output_ptr.add(start_idx + (k + big_n / 4) * stride) = ac_diff + bd_diff_ni;
                *output_ptr.add(start_idx + (k + big_n / 2) * stride) = ac_sum - bd_sum;
                *output_ptr.add(start_idx + (k + 3 * big_n / 4) * stride) = ac_diff - bd_diff_ni;
            }
        }
    }
}

#[inline(always)]
fn rot_45<T: Float + FloatConst>(c: Complex<T>) -> Complex<T> {
    let s = T::FRAC_1_SQRT_2();
    // The standard 2D rotation matrix gives:
    //  [ cos(pi/4) -sin(pi/4)]   [ s -s ]
    //  [ sin(pi/4)  cos(pi/4)] = [ s  s ]
    Complex::<T>::new(c.re - c.im, c.re + c.im) * s
}

#[inline(always)]
fn rot_90<T: Float + FloatConst>(c: Complex<T>) -> Complex<T> {
    // The standard 2D rotation matrix gives:
    //  [ cos(pi/2) -sin(pi/2)]   [ 0 -1 ]
    //  [ sin(pi/2)  cos(pi/2)] = [ 1  0 ]
    Complex::<T>::new(-c.im, c.re)
}

#[inline(always)]
fn rot_180<T: Float + FloatConst>(c: Complex<T>) -> Complex<T> {
    // The standard 2D rotation matrix gives:
    //  [ cos(pi) -sin(pi)]   [ -1  0 ]
    //  [ sin(pi)  cos(pi)] = [  0 -1 ]
    Complex::<T>::new(-c.re, -c.im)
}

#[inline(always)]
fn rot_270<T: Float + FloatConst>(c: Complex<T>) -> Complex<T> {
    // The standard 2D rotation matrix gives:
    //  [ cos(3pi/2) -sin(3pi/2)]   [  0 1 ]
    //  [ sin(3pi/2)  cos(3pi/2)] = [ -1 0 ]
    Complex::<T>::new(c.im, -c.re)
}

fn fft_butterfly_radix_8_unsafe<T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
    stride: usize,
    big_n: usize,
    twiddles: &[Complex<T>],
    _: &ExecutionStage<T>,
) {
    let input_ptr = input.as_ptr();
    let output_ptr = output.as_mut_ptr();
    for start_idx in 0..stride {
        for k in 0..big_n / 8 {
            unsafe {
                // Collect inputs.
                let i0 = *input_ptr.add(start_idx + 8 * k * stride);
                let i1 = *input_ptr.add(start_idx + (8 * k + 1) * stride);
                let i2 = *input_ptr.add(start_idx + (8 * k + 2) * stride);
                let i3 = *input_ptr.add(start_idx + (8 * k + 3) * stride);
                let i4 = *input_ptr.add(start_idx + (8 * k + 4) * stride);
                let i5 = *input_ptr.add(start_idx + (8 * k + 5) * stride);
                let i6 = *input_ptr.add(start_idx + (8 * k + 6) * stride);
                let i7 = *input_ptr.add(start_idx + (8 * k + 7) * stride);

                // Collect relevant twiddles.
                let ot1 = twiddles.get_unchecked(1 * k * stride);
                let ot2 = twiddles.get_unchecked(2 * k * stride);
                let ot3 = twiddles.get_unchecked(3 * k * stride);
                let ot4 = twiddles.get_unchecked(4 * k * stride);
                let ot5 = twiddles.get_unchecked(5 * k * stride);
                let ot6 = twiddles.get_unchecked(6 * k * stride);
                let ot7 = twiddles.get_unchecked(7 * k * stride);

                let a = i0;
                let b = ot1 * i1;
                let c = ot2 * i2;
                let d = ot3 * i3;
                let e = ot4 * i4;
                let f = ot5 * i5;
                let g = ot6 * i6;
                let h = ot7 * i7;

                let ae_sum = a + e;
                let ae_diff = a - e;
                let bf_sum = b + f;
                let bf_diff = b - f;
                let cg_sum = c + g;
                let cg_diff = c - g;
                let dh_sum = d + h;
                let dh_diff = d - h;

                let w00 = ae_sum + cg_sum;
                let w01 = ae_sum - cg_sum;
                let w10 = ae_diff + rot_270(cg_diff);
                let w11 = ae_diff - rot_270(cg_diff);
                let x00 = bf_sum + dh_sum;
                let x01 = rot_270(bf_sum) + rot_90(dh_sum);
                let x10 = rot_45(rot_270(bf_diff) + rot_180(dh_diff));
                let x11 = rot_45(rot_180(bf_diff) + rot_270(dh_diff));

                *output_ptr.add(start_idx + k * stride) = w00 + x00;
                *output_ptr.add(start_idx + (k + big_n / 8) * stride) = w10 + x10;
                *output_ptr.add(start_idx + (k + big_n / 4) * stride) = w01 + x01;
                *output_ptr.add(start_idx + (k + 3 * big_n / 8) * stride) = w11 + x11;
                *output_ptr.add(start_idx + (k + big_n / 2) * stride) = w00 - x00;
                *output_ptr.add(start_idx + (k + 5 * big_n / 8) * stride) = w10 - x10;
                *output_ptr.add(start_idx + (k + 3 * big_n / 4) * stride) = w01 - x01;
                *output_ptr.add(start_idx + (k + 7 * big_n / 8) * stride) = w11 - x11;
            }
        }
    }
}

fn fft_butterfly_radix_8_s0_unsafe<T: Float + FloatConst>(
    input: &mut [Complex<T>],
    output: &mut [Complex<T>],
    _: &Matrix<Complex<T>>,
) {
    let stride = input.len() / 8;
    let big_n = 8;
    let input_ptr = input.as_ptr();
    let output_ptr = output.as_mut_ptr();
    for start_idx in 0..stride {
        for k in 0..big_n / 8 {
            unsafe {
                // Collect inputs.
                let i0 = *input_ptr.add(start_idx + 8 * k * stride);
                let i1 = *input_ptr.add(start_idx + (8 * k + 1) * stride);
                let i2 = *input_ptr.add(start_idx + (8 * k + 2) * stride);
                let i3 = *input_ptr.add(start_idx + (8 * k + 3) * stride);
                let i4 = *input_ptr.add(start_idx + (8 * k + 4) * stride);
                let i5 = *input_ptr.add(start_idx + (8 * k + 5) * stride);
                let i6 = *input_ptr.add(start_idx + (8 * k + 6) * stride);
                let i7 = *input_ptr.add(start_idx + (8 * k + 7) * stride);

                let a = i0;
                let b = i1;
                let c = i2;
                let d = i3;
                let e = i4;
                let f = i5;
                let g = i6;
                let h = i7;

                let ae_sum = a + e;
                let ae_diff = a - e;
                let bf_sum = b + f;
                let bf_diff = b - f;
                let cg_sum = c + g;
                let cg_diff = c - g;
                let dh_sum = d + h;
                let dh_diff = d - h;

                let w00 = ae_sum + cg_sum;
                let w01 = ae_sum - cg_sum;
                let w10 = ae_diff + rot_270(cg_diff);
                let w11 = ae_diff - rot_270(cg_diff);
                let x00 = bf_sum + dh_sum;
                let x01 = rot_270(bf_sum) + rot_90(dh_sum);
                let x10 = rot_45(rot_270(bf_diff) + rot_180(dh_diff));
                let x11 = rot_45(rot_180(bf_diff) + rot_270(dh_diff));

                *output_ptr.add(start_idx + k * stride) = w00 + x00;
                *output_ptr.add(start_idx + (k + big_n / 8) * stride) = w10 + x10;
                *output_ptr.add(start_idx + (k + big_n / 4) * stride) = w01 + x01;
                *output_ptr.add(start_idx + (k + 3 * big_n / 8) * stride) = w11 + x11;
                *output_ptr.add(start_idx + (k + big_n / 2) * stride) = w00 - x00;
                *output_ptr.add(start_idx + (k + 5 * big_n / 8) * stride) = w10 - x10;
                *output_ptr.add(start_idx + (k + 3 * big_n / 4) * stride) = w01 - x01;
                *output_ptr.add(start_idx + (k + 7 * big_n / 8) * stride) = w11 - x11;
            }
        }
    }
}

struct Matrix<T> {
    data: Vec<T>,
    rows: usize,
    cols: usize,
}

impl<T: Clone + Default> Matrix<T> {
    pub fn new(rows: usize, cols: usize) -> Self {
        let data = vec![T::default(); rows * cols];
        Self { data, rows, cols }
    }
}

// Reference: https://doc.rust-lang.org/std/ops/trait.Index.html
impl<T> std::ops::Index<(usize, usize)> for Matrix<T> {
    type Output = T;
    fn index(&self, (r, c): (usize, usize)) -> &T {
        &self.data[r * self.cols + c]
    }
}

impl<T> std::ops::IndexMut<(usize, usize)> for Matrix<T> {
    fn index_mut(&mut self, (r, c): (usize, usize)) -> &mut T {
        &mut self.data[r * self.cols + c]
    }
}

pub struct ExecutionPlan<T: Float + FloatConst> {
    pub twiddles: Vec<Complex<T>>,
    pub stages: Vec<ExecutionStage<T>>,
    pub scratch: Vec<Complex<T>>,
}

pub struct ExecutionStage<T: Float + FloatConst> {
    radix: usize,
    butterfly:
        fn(&mut [Complex<T>], &mut [Complex<T>], usize, usize, &[Complex<T>], &ExecutionStage<T>),
    butterfly_s0: fn(&mut [Complex<T>], &mut [Complex<T>], &Matrix<Complex<T>>),
    // TODO explain wtf this is. Link to article?
    outer_terms: Matrix<Complex<T>>,
}

impl<T: Clone + Default + Float + FloatConst> ExecutionPlan<T> {
    pub fn new(data: &[Complex<T>]) -> anyhow::Result<Self> {
        let data_len = data.len();

        let available_radices = [9, 7, 8, 5, 4, 3, 2];
        let mut stages = Vec::new();
        let mut len = data_len;
        while len > 1 {
            let mut got_match = false;
            for radix in &available_radices {
                if len % radix == 0 {
                    let butterflies = Self::radix_to_butterfly(*radix)?;
                    stages.push(ExecutionStage {
                        radix: *radix,
                        butterfly: butterflies.0,
                        butterfly_s0: butterflies.1,
                        outer_terms: Self::generate_radix_outer_terms(*radix),
                    });
                    len /= radix;
                    got_match = true;
                    break;
                }
            }
            if !got_match {
                anyhow::bail!("Could not factor input length {}", len);
            }
        }

        let twiddles = precompute_twiddles::<T>(data_len);
        let scratch: Vec<Complex<T>> = vec![Complex::new(T::zero(), T::zero()); data.len()];
        Ok(Self {
            twiddles,
            stages,
            scratch,
        })
    }

    fn generate_radix_outer_terms(radix: usize) -> Matrix<Complex<T>> {
        assert!(radix > 1, "radix must be greater than 1");
        let mut data = Matrix::<Complex<T>>::new(radix, radix);
        for row in 0..radix {
            for col in 0..radix {
                let k = (row * col) % radix;
                let entry = Complex::<f64>::cis(-f64::TAU() * (k as f64) / (radix as f64));
                data.data[row * radix + col] =
                    Complex::<T>::new(T::from(entry.re).unwrap(), T::from(entry.im).unwrap());
            }
        }
        data
    }

    fn radix_to_butterfly(
        radix: usize,
    ) -> anyhow::Result<(
        fn(&mut [Complex<T>], &mut [Complex<T>], usize, usize, &[Complex<T>], &ExecutionStage<T>),
        fn(&mut [Complex<T>], &mut [Complex<T>], &Matrix<Complex<T>>),
    )> {
        match radix {
            8 => Ok((
                fft_butterfly_radix_8_unsafe,
                fft_butterfly_radix_8_s0_unsafe,
            )),
            4 => Ok((
                fft_butterfly_radix_4_unsafe,
                fft_butterfly_radix_4_s0_unsafe,
            )),
            2 => Ok((
                fft_butterfly_radix_2_unsafe,
                fft_butterfly_radix_2_s0_unsafe,
            )),
            9 => Ok((
                fft_butterfly_radix_9_unsafe,
                fft_butterfly_radix_9_s0_unsafe,
            )),
            3 => Ok((
                fft_butterfly_radix_3_unsafe,
                fft_butterfly_radix_3_s0_unsafe,
            )),
            5 => Ok((
                fft_butterfly_radix_p_unsafe::<5, T>,
                fft_butterfly_radix_p_s0_unsafe::<5, T>,
            )),
            7 => Ok((
                fft_butterfly_radix_p_unsafe::<7, T>,
                fft_butterfly_radix_p_s0_unsafe::<7, T>,
            )),
            _ => anyhow::bail!("Unsupported radix {}", radix),
        }
    }
}

pub fn fft<T: Float + FloatConst>(src: &mut [Complex<T>], plan: &mut ExecutionPlan<T>) {
    let scratch = plan.scratch.as_mut_slice();
    assert_eq!(src.len(), scratch.len());
    assert_eq!(plan.twiddles.len(), src.len());

    let n_iter = plan.stages.len();
    if n_iter % 2 == 1 {
        scratch.copy_from_slice(src);
    }

    let (mut input, mut output) = if n_iter % 2 == 0 {
        (scratch, src)
    } else {
        (src, scratch)
    };

    let big_n = input.len();
    let mut stride = big_n;
    let mut big_n = 1;
    for (stage_idx, stage) in plan.stages.iter().enumerate() {
        stride /= stage.radix;
        big_n *= stage.radix;
        std::mem::swap(&mut input, &mut output);

        if stage_idx == 0 {
            (stage.butterfly_s0)(input, output, &stage.outer_terms);
        } else {
            (stage.butterfly)(input, output, stride, big_n, &plan.twiddles, &stage);
        }
    }
}

#[cfg(test)]
mod tests {
    use rand::RngExt;
    use rand::SeedableRng;
    use rand::rngs::StdRng;

    use super::*;
    use num::complex::Complex;
    use std::time::Instant;

    // Cast one float type to another, truncating if needed.
    fn t0_to_t1<T0: Float, T1: Float>(val: T0) -> T1 {
        return num::cast(val).unwrap();
    }

    // Casts one Complex<T> to another.
    fn complex_to_t<T0: Float, T1: Float>(data: &[Complex<T0>]) -> Vec<Complex<T1>> {
        data.iter()
            .map(|c| Complex::new(t0_to_t1(c.re), t0_to_t1(c.im)))
            .collect()
    }

    fn evaluate_results(
        result_ref_64: &Vec<Complex<f64>>,
        result_cur_64: &Vec<Complex<f64>>,
        duration_dft: std::time::Duration,
        duration_cur: std::time::Duration,
        algo_name: &str,
    ) {
        assert_eq!(result_ref_64.len(), result_cur_64.len());

        println!("  Algorithm:  {}", algo_name);
        println!("    Duration:  {:?}", duration_cur);
        println!(
            "    Speedup:   {:?}",
            duration_dft.as_nanos() as f32 / duration_cur.as_nanos() as f32
        );

        let mut max_err: f64 = 0.0;
        let mut sum_err: f64 = 0.0;
        // TODO: median, 90p, 99p

        let result_64: Vec<Complex<f64>> = complex_to_t(&result_cur_64);

        for i in 0..result_ref_64.len() {
            max_err = max_err.max((result_ref_64[i] - result_64[i]).norm());
            sum_err += (result_ref_64[i] - result_64[i]).norm();
        }
        println!("    Max err:   {}", max_err);
        println!("    Avg err:   {}", sum_err / (result_ref_64.len() as f64));
    }

    fn evaluate_size(size: u32) {
        println!("size={}", size);

        let mut data: Vec<Complex<f64>> = Vec::new();
        data.resize(size as usize, Complex::ZERO);
        let mut rng = StdRng::seed_from_u64(42);
        for i in 0..data.len() {
            data[i] = Complex::new(rng.random(), rng.random());
        }

        let mut result_ref: Vec<rustfft::num_complex::Complex<f64>> = complex_to_t(&data);
        let mut ref_planner = rustfft::FftPlanner::<f64>::new();
        let ref_fft = ref_planner.plan_fft_forward(result_ref.len());
        let begin_instant = Instant::now();
        ref_fft.process(&mut result_ref);
        let duration_ref = begin_instant.elapsed();
        let result_ref_64: Vec<Complex<f64>> = complex_to_t(&result_ref);
        evaluate_results(
            &result_ref_64,
            &result_ref_64,
            std::time::Duration::from_secs(0),
            duration_ref,
            "REF",
        );

        let mut result_rfft: Vec<rustfft::num_complex::Complex<f32>> = complex_to_t(&data);
        let mut rfftner = rustfft::FftPlanner::<f32>::new();
        let rfft_fft = rfftner.plan_fft_forward(result_rfft.len());
        let mut begin_instant = Instant::now();
        rfft_fft.process(&mut result_rfft);
        let duration_rfft = begin_instant.elapsed();
        let result_rfft_64: Vec<Complex<f64>> = complex_to_t(&result_rfft);
        evaluate_results(
            &result_ref_64,
            &result_rfft_64,
            std::time::Duration::from_secs(0),
            duration_rfft,
            "RFFT",
        );

        let mut result_fft = complex_to_t(&data);
        let mut plan = ExecutionPlan::<f32>::new(&result_fft).unwrap();
        begin_instant = Instant::now();
        fft(&mut result_fft, &mut plan);
        let duration_fft = begin_instant.elapsed();
        let result_fft_64: Vec<Complex<f64>> = complex_to_t(&result_fft);
        evaluate_results(
            &result_ref_64,
            &result_fft_64,
            std::time::Duration::from_secs(0),
            duration_fft,
            "FFT v8 (plan)",
        );
    }

    #[test]
    fn naive_fft_matches_ref() {
        for exponent in 4..17 {
            evaluate_size(2u32.pow(exponent));
        }
        evaluate_size(3u32.pow(3)); // 27
        evaluate_size(3u32.pow(4)); // 81
        evaluate_size(3u32.pow(5)); // 243
        evaluate_size(3u32.pow(6)); // 729
        evaluate_size(3u32.pow(7)); // 2187
        evaluate_size(3u32.pow(8)); // 6561

        evaluate_size(5u32.pow(2)); // 25
        evaluate_size(5u32.pow(3)); // 125
        evaluate_size(5u32.pow(4)); // 625
        evaluate_size(5u32.pow(5)); // 3125

        evaluate_size(7u32.pow(2)); // 49
        evaluate_size(7u32.pow(3)); // 343
        evaluate_size(7u32.pow(4)); // 2401
    }
}