pxdct 0.2.2

/*
 * // Copyright (c) Radzivon Bartoshyk 12/2025. All rights reserved.
 * //
 * // Redistribution and use in source and binary forms, with or without modification,
 * // are permitted provided that the following conditions are met:
 * //
 * // 1.  Redistributions of source code must retain the above copyright notice, this
 * // list of conditions and the following disclaimer.
 * //
 * // 2.  Redistributions in binary form must reproduce the above copyright notice,
 * // this list of conditions and the following disclaimer in the documentation
 * // and/or other materials provided with the distribution.
 * //
 * // 3.  Neither the name of the copyright holder nor the names of its
 * // contributors may be used to endorse or promote products derived from
 * // this software without specific prior written permission.
 * //
 * // THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * // AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * // IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * // DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * // FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * // DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * // SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * // CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
#![allow(unused)]
use crate::factory_dct2::Dct2Factory;
use crate::mla::fmla;
use crate::twiddles::compute_twiddle;
use crate::util::DctSample;
use crate::{PxdctError, PxdctExecutor};
use num_complex::Complex;
use num_traits::AsPrimitive;
use std::marker::PhantomData;
use std::sync::Arc;

#[derive(Debug, Clone, Default)]
pub(crate) struct Dst2Butterfly2<T> {
    phantom_data: PhantomData<T>,
}

impl<T: DctSample> Dst2Butterfly2<T> {
    #[inline(always)]
    pub(crate) fn exec(data: &mut [T; 2]) {
        let u0 = data[0];
        let u1 = data[1];
        let sum = u0 - u1;
        let v0 = (u0 + u1) * T::FRAC_1_SQRT_2;
        let v1 = sum;
        data[0] = v0;
        data[1] = v1;
    }
}

impl<T: DctSample> PxdctExecutor<T> for Dst2Butterfly2<T>
where
    f64: AsPrimitive<T>,
{
    fn execute(&self, data: &mut [T]) -> Result<(), PxdctError> {
        if !data.len().is_multiple_of(2) {
            return Err(PxdctError::InvalidSizeMultiplier(data.len(), self.length()));
        }
        for chunk in data.chunks_exact_mut(2) {
            let u0 = chunk[0];
            let u1 = chunk[1];
            let sum = u0 - u1;
            chunk[0] = (u0 + u1) * T::FRAC_1_SQRT_2;
            chunk[1] = sum;
        }
        Ok(())
    }

    fn length(&self) -> usize {
        2
    }
}

#[derive(Debug, Clone, Default)]
pub(crate) struct Dct2Butterfly2<T> {
    phantom_data: PhantomData<T>,
}

impl<T: DctSample> Dct2Butterfly2<T> {
    #[inline(always)]
    pub(crate) fn exec(data: &mut [T; 2]) {
        let sum = data[0] + data[1];
        let v1 = (data[0] - data[1]) * T::FRAC_1_SQRT_2;
        let v0 = sum;
        data[0] = v0;
        data[1] = v1;
    }
}

impl<T: DctSample> PxdctExecutor<T> for Dct2Butterfly2<T>
where
    f64: AsPrimitive<T>,
{
    fn execute(&self, data: &mut [T]) -> Result<(), PxdctError> {
        if !data.len().is_multiple_of(2) {
            return Err(PxdctError::InvalidSizeMultiplier(data.len(), self.length()));
        }
        for chunk in data.chunks_exact_mut(2) {
            let sum = chunk[0] + chunk[1];
            chunk[1] = (chunk[0] - chunk[1]) * T::FRAC_1_SQRT_2;
            chunk[0] = sum;
        }
        Ok(())
    }

    fn length(&self) -> usize {
        2
    }
}

#[derive(Debug, Clone)]
pub(crate) struct Dct2Butterfly4<T: DctSample> {
    twiddle: Complex<T>,
}

impl<T: DctSample> Default for Dct2Butterfly4<T>
where
    f64: AsPrimitive<T>,
{
    fn default() -> Self {
        Self {
            twiddle: compute_twiddle(1, 16).conj(),
        }
    }
}

impl<T: DctSample> Dct2Butterfly4<T>
where
    f64: AsPrimitive<T>,
{
    #[inline(always)]
    pub(crate) fn exec(&self, data: &mut [T; 4]) {
        let u0 = data[0];
        let u1 = data[1];
        let u2 = data[2];
        let u3 = data[3];

        let lower_dct4 = u0 - u3;
        let upper_dct4 = u2 - u1;

        let mut dct_evens = [u0 + u3, u2 + u1];

        Dct2Butterfly2::exec(&mut dct_evens);

        let v1 = fmla(lower_dct4, self.twiddle.re, -upper_dct4 * self.twiddle.im);
        let v3 = fmla(upper_dct4, self.twiddle.re, lower_dct4 * self.twiddle.im);
        data[0] = dct_evens[0];
        data[1] = v1;
        data[2] = dct_evens[1];
        data[3] = v3;
    }
}

impl<T: DctSample> PxdctExecutor<T> for Dct2Butterfly4<T>
where
    f64: AsPrimitive<T>,
{
    fn execute(&self, data: &mut [T]) -> Result<(), PxdctError> {
        if !data.len().is_multiple_of(4) {
            return Err(PxdctError::InvalidSizeMultiplier(data.len(), self.length()));
        }
        for chunk in data.chunks_exact_mut(4) {
            let u0 = chunk[0];
            let u1 = chunk[1];
            let u2 = chunk[2];
            let u3 = chunk[3];

            let lower_dct4 = u0 - u3;
            let upper_dct4 = u2 - u1;

            let mut dct_evens = [u0 + u3, u2 + u1];

            Dct2Butterfly2::exec(&mut dct_evens);

            chunk[0] = dct_evens[0];
            chunk[2] = dct_evens[1];

            chunk[1] = fmla(lower_dct4, self.twiddle.re, -upper_dct4 * self.twiddle.im);
            chunk[3] = fmla(upper_dct4, self.twiddle.re, lower_dct4 * self.twiddle.im);
        }
        Ok(())
    }

    fn length(&self) -> usize {
        4
    }
}

#[derive(Debug, Clone)]
pub(crate) struct Dst2Butterfly4<T: DctSample> {
    twiddle: Complex<T>,
}

impl<T: DctSample> Default for Dst2Butterfly4<T>
where
    f64: AsPrimitive<T>,
{
    fn default() -> Self {
        Self {
            twiddle: compute_twiddle(1, 16).conj(),
        }
    }
}

impl<T: DctSample> Dst2Butterfly4<T>
where
    f64: AsPrimitive<T>,
{
    #[inline(always)]
    fn exec(&self, data: &[T; 4]) -> [T; 4] {
        let u0 = data[0];
        let u1 = data[1];
        let u2 = data[2];
        let u3 = data[3];

        let lower_dct4 = u0 + u3;
        let upper_dct4 = u2 + u1;

        let mut dct_evens = [u0 - u3, u2 - u1];

        Dct2Butterfly2::exec(&mut dct_evens);

        let v2 = fmla(lower_dct4, self.twiddle.re, -upper_dct4 * self.twiddle.im);
        let v0 = fmla(upper_dct4, self.twiddle.re, lower_dct4 * self.twiddle.im);
        [v0, dct_evens[1], v2, dct_evens[0]]
    }
}

impl<T: DctSample> PxdctExecutor<T> for Dst2Butterfly4<T>
where
    f64: AsPrimitive<T>,
{
    fn execute(&self, data: &mut [T]) -> Result<(), PxdctError> {
        if !data.len().is_multiple_of(4) {
            return Err(PxdctError::InvalidSizeMultiplier(data.len(), self.length()));
        }
        for chunk in data.chunks_exact_mut(4) {
            let u0 = chunk[0];
            let u1 = chunk[1];
            let u2 = chunk[2];
            let u3 = chunk[3];

            let w = self.exec(&[u0, u1, u2, u3]);
            chunk[0] = w[0];
            chunk[1] = w[1];
            chunk[2] = w[2];
            chunk[3] = w[3];
        }
        Ok(())
    }

    fn length(&self) -> usize {
        4
    }
}

#[derive(Debug, Clone)]
pub(crate) struct Dct2Butterfly8<T: DctSample> {
    bf4: Dct2Butterfly4<T>,
    twiddle0: Complex<T>,
    twiddle1: Complex<T>,
}

impl<T: DctSample> Default for Dct2Butterfly8<T>
where
    f64: AsPrimitive<T>,
{
    fn default() -> Self {
        Self {
            bf4: Dct2Butterfly4::default(),
            twiddle0: compute_twiddle(1, 32).conj(),
            twiddle1: compute_twiddle(3, 32).conj(),
        }
    }
}

impl<T: DctSample> Dct2Butterfly8<T>
where
    f64: AsPrimitive<T>,
{
    #[inline(always)]
    pub(crate) fn exec(&self, data: &mut [T; 8]) {
        let u0 = data[0];
        let u1 = data[1];
        let u2 = data[2];
        let u3 = data[3];
        let u4 = data[4];
        let u5 = data[5];
        let u6 = data[6];
        let u7 = data[7];

        let mut dct2_buffer = [u0 + u7, u1 + u6, u2 + u5, u3 + u4];
        self.bf4.exec(&mut dct2_buffer);

        // odds
        let differences = [u0 - u7, u3 - u4, u1 - u6, u2 - u5];

        let mut dct4_even_buffer = [
            fmla(
                differences[0],
                self.twiddle0.re,
                differences[1] * self.twiddle0.im,
            ),
            fmla(
                differences[2],
                self.twiddle1.re,
                differences[3] * self.twiddle1.im,
            ),
        ];

        Dct2Butterfly2::exec(&mut dct4_even_buffer);
        let mut dct4_odd_buffer = [
            fmla(
                differences[3],
                self.twiddle1.re,
                -differences[2] * self.twiddle1.im,
            ),
            fmla(
                differences[1],
                self.twiddle0.re,
                -differences[0] * self.twiddle0.im,
            ),
        ];
        Dst2Butterfly2::exec(&mut dct4_odd_buffer);

        // combine the results
        data[0] = dct2_buffer[0];
        data[1] = dct4_even_buffer[0];
        data[2] = dct2_buffer[1];
        data[3] = dct4_even_buffer[1] - dct4_odd_buffer[0];
        data[4] = dct2_buffer[2];
        data[5] = dct4_even_buffer[1] + dct4_odd_buffer[0];
        data[6] = dct2_buffer[3];
        data[7] = dct4_odd_buffer[1];
    }
}

impl<T: DctSample> PxdctExecutor<T> for Dct2Butterfly8<T>
where
    f64: AsPrimitive<T>,
{
    fn execute(&self, data: &mut [T]) -> Result<(), PxdctError> {
        if !data.len().is_multiple_of(8) {
            return Err(PxdctError::InvalidSizeMultiplier(data.len(), self.length()));
        }
        for chunk in data.chunks_exact_mut(8) {
            self.exec((&mut chunk[..8]).try_into().unwrap());
        }
        Ok(())
    }

    fn length(&self) -> usize {
        8
    }
}

#[derive(Debug, Clone)]
pub(crate) struct Dct2Butterfly16<T: DctSample> {
    bf8: Dct2Butterfly8<T>,
    bf4_dst: Dst2Butterfly4<T>,
    twiddle0: Complex<T>,
    twiddle1: Complex<T>,
    twiddle2: Complex<T>,
    twiddle3: Complex<T>,
}

impl<T: DctSample> Default for Dct2Butterfly16<T>
where
    f64: AsPrimitive<T>,
{
    fn default() -> Self {
        Self {
            bf8: Dct2Butterfly8::default(),
            bf4_dst: Dst2Butterfly4::default(),
            twiddle0: compute_twiddle(1, 64).conj(),
            twiddle1: compute_twiddle(3, 64).conj(),
            twiddle2: compute_twiddle(5, 64).conj(),
            twiddle3: compute_twiddle(7, 64).conj(),
        }
    }
}

impl<T: DctSample> Dct2Butterfly16<T>
where
    f64: AsPrimitive<T>,
{
    #[inline(always)]
    pub(crate) fn exec(&self, chunk: &mut [T; 16]) {
        let u0 = chunk[0];
        let u1 = chunk[1];
        let u2 = chunk[2];
        let u3 = chunk[3];
        let u4 = chunk[4];
        let u5 = chunk[5];
        let u6 = chunk[6];
        let u7 = chunk[7];
        let u8 = chunk[8];
        let u9 = chunk[9];
        let u10 = chunk[10];
        let u11 = chunk[11];
        let u12 = chunk[12];
        let u13 = chunk[13];
        let u14 = chunk[14];
        let u15 = chunk[15];

        //process the evens
        let mut dct2_buffer = [
            u0 + u15,
            u1 + u14,
            u2 + u13,
            u3 + u12,
            u4 + u11,
            u5 + u10,
            u6 + u9,
            u7 + u8,
        ];
        self.bf8.exec(&mut dct2_buffer);

        //process the odds
        let differences = [
            u0 - u15,
            u7 - u8,
            u1 - u14,
            u6 - u9,
            u2 - u13,
            u5 - u10,
            u3 - u12,
            u4 - u11,
        ];

        let mut dct4_even_buffer = [
            fmla(
                differences[0],
                self.twiddle0.re,
                differences[1] * self.twiddle0.im,
            ),
            fmla(
                differences[2],
                self.twiddle1.re,
                differences[3] * self.twiddle1.im,
            ),
            fmla(
                differences[4],
                self.twiddle2.re,
                differences[5] * self.twiddle2.im,
            ),
            fmla(
                differences[6],
                self.twiddle3.re,
                differences[7] * self.twiddle3.im,
            ),
        ];
        let dct4_odd_buffer = [
            fmla(
                differences[7],
                self.twiddle3.re,
                -differences[6] * self.twiddle3.im,
            ),
            fmla(
                differences[5],
                self.twiddle2.re,
                -differences[4] * self.twiddle2.im,
            ),
            fmla(
                differences[3],
                self.twiddle1.re,
                -differences[2] * self.twiddle1.im,
            ),
            fmla(
                differences[1],
                self.twiddle0.re,
                -differences[0] * self.twiddle0.im,
            ),
        ];

        self.bf8.bf4.exec(&mut dct4_even_buffer);
        let dct4_odd_buffer = self.bf4_dst.exec(&dct4_odd_buffer);

        // combine the results
        chunk[0] = dct2_buffer[0];
        chunk[1] = dct4_even_buffer[0];
        chunk[2] = dct2_buffer[1];
        chunk[3] = dct4_even_buffer[1] - dct4_odd_buffer[0];
        chunk[4] = dct2_buffer[2];
        chunk[5] = dct4_even_buffer[1] + dct4_odd_buffer[0];
        chunk[6] = dct2_buffer[3];
        chunk[7] = dct4_even_buffer[2] + dct4_odd_buffer[1];
        chunk[8] = dct2_buffer[4];
        chunk[9] = dct4_even_buffer[2] - dct4_odd_buffer[1];
        chunk[10] = dct2_buffer[5];
        chunk[11] = dct4_even_buffer[3] - dct4_odd_buffer[2];
        chunk[12] = dct2_buffer[6];
        chunk[13] = dct4_even_buffer[3] + dct4_odd_buffer[2];
        chunk[14] = dct2_buffer[7];
        chunk[15] = dct4_odd_buffer[3];
    }
}

impl<T: DctSample> PxdctExecutor<T> for Dct2Butterfly16<T>
where
    f64: AsPrimitive<T>,
{
    fn execute(&self, data: &mut [T]) -> Result<(), PxdctError> {
        if !data.len().is_multiple_of(16) {
            return Err(PxdctError::InvalidSizeMultiplier(data.len(), self.length()));
        }
        for chunk in data.chunks_exact_mut(16) {
            self.exec((&mut chunk[..16]).try_into().unwrap());
        }
        Ok(())
    }

    fn length(&self) -> usize {
        16
    }
}

#[derive(Debug, Clone)]
pub(crate) struct Dct2Butterfly32<T: DctSample> {
    bf8: Dct2Butterfly8<T>,
    bf16: Dct2Butterfly16<T>,
    twiddles: [Complex<T>; 8],
}

impl<T: DctSample> Default for Dct2Butterfly32<T>
where
    f64: AsPrimitive<T>,
    usize: AsPrimitive<T>,
{
    fn default() -> Self {
        use crate::twiddles::compute_twiddle;
        let mut twiddles = [Complex::<T>::default(); 8];
        for (i, twiddle) in twiddles.iter_mut().enumerate() {
            *twiddle = compute_twiddle::<T>(2 * i + 1, 32 * 4).conj();
        }
        Self {
            twiddles,
            bf8: Dct2Butterfly8::default(),
            bf16: Dct2Butterfly16::default(),
        }
    }
}

impl<T: DctSample> Dct2Butterfly32<T>
where
    f64: AsPrimitive<T>,
{
    #[inline(always)]
    fn exec(
        &self,
        data: &mut [T; 32],
        input_dct2: &mut [T; 16],
        dct4_even: &mut [T; 8],
        dct4_odd: &mut [T; 8],
    ) {
        //preprocess the data by splitting it up into vectors of size n/2, n/4, and n/4

        for i in 0..8 {
            let twiddle = self.twiddles[i];
            let input_bottom = data[i];
            let input_top = data[32 - i - 1];

            let input_half_bottom = data[16 - i - 1];
            let input_half_top = data[16 + i];

            //prepare the inner DCT2
            unsafe { *input_dct2.get_unchecked_mut(i) = input_top + input_bottom };
            unsafe {
                *input_dct2.get_unchecked_mut(16 - i - 1) = input_half_bottom + input_half_top
            };

            //prepare the inner DCT4 - which consists of two DCT2s of half size
            let lower_dct4 = input_bottom - input_top;
            let upper_dct4 = input_half_bottom - input_half_top;

            let cos_input = fmla(lower_dct4, twiddle.re, upper_dct4 * twiddle.im);
            let sin_input = fmla(upper_dct4, twiddle.re, -lower_dct4 * twiddle.im);

            unsafe { *dct4_even.get_unchecked_mut(i) = cos_input };
            unsafe {
                *dct4_odd.get_unchecked_mut(8 - i - 1) =
                    if i % 2 == 0 { sin_input } else { -sin_input }
            };
        }

        self.bf16.exec(input_dct2);
        self.bf8.exec(dct4_even);
        self.bf8.exec(dct4_odd);

        unsafe {
            data[0] = *input_dct2.get_unchecked(0);
            data[1] = *dct4_even.get_unchecked(0);
            data[2] = *input_dct2.get_unchecked(1);

            for i in 1..8 {
                let dct4_cos_output = *dct4_even.get_unchecked(i);
                let dct4_sin_output = if (i + 8) % 2 == 0 {
                    -*dct4_odd.get_unchecked(8 - i)
                } else {
                    *dct4_odd.get_unchecked(8 - i)
                };

                *data.get_unchecked_mut(i * 4 - 1) = dct4_cos_output + dct4_sin_output;
                *data.get_unchecked_mut(i * 4) = *input_dct2.get_unchecked(i * 2);

                *data.get_unchecked_mut(i * 4 + 1) = dct4_cos_output - dct4_sin_output;
                *data.get_unchecked_mut(i * 4 + 2) = *input_dct2.get_unchecked(i * 2 + 1);
            }

            *data.get_unchecked_mut(32 - 1) = -*dct4_odd.get_unchecked(0);
        }
    }
}

impl<T: DctSample> PxdctExecutor<T> for Dct2Butterfly32<T>
where
    f64: AsPrimitive<T>,
{
    fn execute(&self, data: &mut [T]) -> Result<(), PxdctError> {
        if !data.len().is_multiple_of(32) {
            return Err(PxdctError::InvalidSizeMultiplier(data.len(), self.length()));
        }

        let mut input_dct2 = [T::zero(); 16];
        let mut dct4_even = [T::zero(); 8];
        let mut dct4_odd = [T::zero(); 8];

        for chunk in data.chunks_exact_mut(32) {
            self.exec(
                (&mut chunk[..32]).try_into().unwrap(),
                &mut input_dct2,
                &mut dct4_even,
                &mut dct4_odd,
            );
        }
        Ok(())
    }

    fn length(&self) -> usize {
        32
    }
}

#[derive(Debug, Clone)]
pub(crate) struct Dct2Butterfly64<T: DctSample> {
    bf32: Dct2Butterfly32<T>,
    bf16: Dct2Butterfly16<T>,
    twiddles: [Complex<T>; 16],
}

impl<T: DctSample> Default for Dct2Butterfly64<T>
where
    f64: AsPrimitive<T>,
    usize: AsPrimitive<T>,
{
    fn default() -> Self {
        use crate::twiddles::compute_twiddle;
        let mut twiddles = [Complex::<T>::default(); 16];
        for (i, twiddle) in twiddles.iter_mut().enumerate() {
            *twiddle = compute_twiddle::<T>(2 * i + 1, 64 * 4).conj();
        }
        Self {
            twiddles,
            bf16: Dct2Butterfly16::default(),
            bf32: Dct2Butterfly32::default(),
        }
    }
}

impl<T: DctSample> Dct2Butterfly64<T>
where
    f64: AsPrimitive<T>,
{
    #[inline(always)]
    fn exec(
        &self,
        data: &mut [T; 64],
        input_dct2: &mut [T; 32],
        input_dct4_even: &mut [T; 16],
        input_dct4_odd: &mut [T; 16],
        input_dct21: &mut [T; 16],
        dct4_even1: &mut [T; 8],
        dct4_odd1: &mut [T; 8],
    ) {
        //preprocess the data by splitting it up into vectors of size n/2, n/4, and n/4

        for i in 0..16 {
            let twiddle = self.twiddles[i];
            let input_bottom = data[i];
            let input_top = data[64 - i - 1];

            let input_half_bottom = data[32 - i - 1];
            let input_half_top = data[32 + i];

            //prepare the inner DCT2
            unsafe { *input_dct2.get_unchecked_mut(i) = input_top + input_bottom };
            unsafe {
                *input_dct2.get_unchecked_mut(32 - i - 1) = input_half_bottom + input_half_top
            };

            //prepare the inner DCT4 - which consists of two DCT2s of half size
            let lower_dct4 = input_bottom - input_top;
            let upper_dct4 = input_half_bottom - input_half_top;

            let cos_input = fmla(lower_dct4, twiddle.re, upper_dct4 * twiddle.im);
            let sin_input = fmla(upper_dct4, twiddle.re, -lower_dct4 * twiddle.im);

            unsafe { *input_dct4_even.get_unchecked_mut(i) = cos_input };
            unsafe {
                *input_dct4_odd.get_unchecked_mut(16 - i - 1) =
                    if i % 2 == 0 { sin_input } else { -sin_input }
            };
        }

        self.bf32
            .exec(input_dct2, input_dct21, dct4_even1, dct4_odd1);
        self.bf16.exec(input_dct4_even);
        self.bf16.exec(input_dct4_odd);

        unsafe {
            data[0] = *input_dct2.get_unchecked(0);
            data[1] = *input_dct4_even.get_unchecked(0);
            data[2] = *input_dct2.get_unchecked(1);

            for i in 1..16 {
                let dct4_cos_output = *input_dct4_even.get_unchecked(i);
                let dct4_sin_output = if (i + 16) % 2 == 0 {
                    -*input_dct4_odd.get_unchecked(16 - i)
                } else {
                    *input_dct4_odd.get_unchecked(16 - i)
                };

                *data.get_unchecked_mut(i * 4 - 1) = dct4_cos_output + dct4_sin_output;
                *data.get_unchecked_mut(i * 4) = *input_dct2.get_unchecked(i * 2);

                *data.get_unchecked_mut(i * 4 + 1) = dct4_cos_output - dct4_sin_output;
                *data.get_unchecked_mut(i * 4 + 2) = *input_dct2.get_unchecked(i * 2 + 1);
            }

            *data.get_unchecked_mut(64 - 1) = -*input_dct4_odd.get_unchecked(0);
        }
    }
}

impl<T: DctSample> PxdctExecutor<T> for Dct2Butterfly64<T>
where
    f64: AsPrimitive<T>,
{
    fn execute(&self, data: &mut [T]) -> Result<(), PxdctError> {
        if !data.len().is_multiple_of(64) {
            return Err(PxdctError::InvalidSizeMultiplier(data.len(), self.length()));
        }

        let mut input_dct2 = [T::zero(); 32];
        let mut input_dct4_even = [T::zero(); 16];
        let mut input_dct4_odd = [T::zero(); 16];

        let mut input_dct21 = [T::zero(); 16];
        let mut dct4_even = [T::zero(); 8];
        let mut dct4_odd = [T::zero(); 8];

        for chunk in data.chunks_exact_mut(64) {
            self.exec(
                (&mut chunk[..64]).try_into().unwrap(),
                &mut input_dct2,
                &mut input_dct4_even,
                &mut input_dct4_odd,
                &mut input_dct21,
                &mut dct4_even,
                &mut dct4_odd,
            );
        }
        Ok(())
    }

    fn length(&self) -> usize {
        64
    }
}

#[derive(Clone)]
pub(crate) struct Dct2Butterfly128<T: DctSample> {
    bf64: Arc<dyn PxdctExecutor<T> + Send + Sync>,
    bf32: Arc<dyn PxdctExecutor<T> + Send + Sync>,
    twiddles: [Complex<T>; 32],
}

impl<T: DctSample + Dct2Factory> Default for Dct2Butterfly128<T>
where
    f64: AsPrimitive<T>,
    usize: AsPrimitive<T>,
{
    fn default() -> Self {
        use crate::twiddles::compute_twiddle;
        let mut twiddles = [Complex::<T>::default(); 32];
        for (i, twiddle) in twiddles.iter_mut().enumerate() {
            *twiddle = compute_twiddle::<T>(2 * i + 1, 128 * 4).conj();
        }
        Self {
            twiddles,
            bf64: T::dct2_butterfly64(),
            bf32: T::dct2_butterfly32(),
        }
    }
}

impl<T: DctSample> Dct2Butterfly128<T>
where
    f64: AsPrimitive<T>,
{
    #[inline(always)]
    fn exec(&self, data: &mut [T; 128], input_dct2: &mut [T; 64], input_dct4: &mut [T; 64]) {
        let (input_dct4_even, input_dct4_odd) = input_dct4.split_at_mut(32);

        //preprocess the data by splitting it up into vectors of size n/2, n/4, and n/4

        for i in 0..32 {
            let twiddle = self.twiddles[i];
            let input_bottom = data[i];
            let input_top = data[128 - i - 1];

            let input_half_bottom = data[64 - i - 1];
            let input_half_top = data[64 + i];

            //prepare the inner DCT2
            unsafe { *input_dct2.get_unchecked_mut(i) = input_top + input_bottom };
            unsafe {
                *input_dct2.get_unchecked_mut(64 - i - 1) = input_half_bottom + input_half_top
            };

            //prepare the inner DCT4 - which consists of two DCT2s of half size
            let lower_dct4 = input_bottom - input_top;
            let upper_dct4 = input_half_bottom - input_half_top;

            let cos_input = fmla(lower_dct4, twiddle.re, upper_dct4 * twiddle.im);
            let sin_input = fmla(upper_dct4, twiddle.re, -lower_dct4 * twiddle.im);

            unsafe { *input_dct4_even.get_unchecked_mut(i) = cos_input };
            unsafe {
                *input_dct4_odd.get_unchecked_mut(32 - i - 1) =
                    if i % 2 == 0 { sin_input } else { -sin_input }
            };
        }

        _ = self.bf64.execute(input_dct2);
        _ = self.bf32.execute(input_dct4);

        let (input_dct4_even, input_dct4_odd) = input_dct4.split_at_mut(32);

        unsafe {
            data[0] = *input_dct2.get_unchecked(0);
            data[1] = *input_dct4_even.get_unchecked(0);
            data[2] = *input_dct2.get_unchecked(1);

            for i in 1..32 {
                let dct4_cos_output = *input_dct4_even.get_unchecked(i);
                let dct4_sin_output = if (i + 32) % 2 == 0 {
                    -*input_dct4_odd.get_unchecked(32 - i)
                } else {
                    *input_dct4_odd.get_unchecked(32 - i)
                };

                *data.get_unchecked_mut(i * 4 - 1) = dct4_cos_output + dct4_sin_output;
                *data.get_unchecked_mut(i * 4) = *input_dct2.get_unchecked(i * 2);

                *data.get_unchecked_mut(i * 4 + 1) = dct4_cos_output - dct4_sin_output;
                *data.get_unchecked_mut(i * 4 + 2) = *input_dct2.get_unchecked(i * 2 + 1);
            }

            *data.get_unchecked_mut(128 - 1) = -*input_dct4_odd.get_unchecked(0);
        }
    }
}

impl<T: DctSample> PxdctExecutor<T> for Dct2Butterfly128<T>
where
    f64: AsPrimitive<T>,
{
    fn execute(&self, data: &mut [T]) -> Result<(), PxdctError> {
        if !data.len().is_multiple_of(128) {
            return Err(PxdctError::InvalidSizeMultiplier(data.len(), self.length()));
        }

        let mut input_dct2 = [T::zero(); 64];
        let mut input_dct4 = [T::zero(); 64];

        for chunk in data.chunks_exact_mut(128) {
            self.exec(
                (&mut chunk[..128]).try_into().unwrap(),
                &mut input_dct2,
                &mut input_dct4,
            );
        }
        Ok(())
    }

    fn length(&self) -> usize {
        128
    }
}

#[derive(Clone)]
pub(crate) struct Dct2Butterfly256<T: DctSample> {
    bf128: Arc<dyn PxdctExecutor<T> + Send + Sync>,
    bf64: Arc<dyn PxdctExecutor<T> + Send + Sync>,
    twiddles: [Complex<T>; 64],
}

impl<T: DctSample + Dct2Factory> Default for Dct2Butterfly256<T>
where
    f64: AsPrimitive<T>,
    usize: AsPrimitive<T>,
{
    fn default() -> Self {
        use crate::twiddles::compute_twiddle;
        let mut twiddles = [Complex::<T>::default(); 64];
        for (i, twiddle) in twiddles.iter_mut().enumerate() {
            *twiddle = compute_twiddle::<T>(2 * i + 1, 256 * 4).conj();
        }
        Self {
            twiddles,
            bf64: T::dct2_butterfly64(),
            bf128: T::dct2_butterfly128(),
        }
    }
}

impl<T: DctSample> Dct2Butterfly256<T>
where
    f64: AsPrimitive<T>,
{
    #[inline(always)]
    fn exec(&self, data: &mut [T; 256], input_dct2: &mut [T; 128], input_dct4: &mut [T; 128]) {
        let (input_dct4_even, input_dct4_odd) = input_dct4.split_at_mut(64);

        //preprocess the data by splitting it up into vectors of size n/2, n/4, and n/4

        for i in 0..64 {
            let twiddle = self.twiddles[i];
            let input_bottom = data[i];
            let input_top = data[256 - i - 1];

            let input_half_bottom = data[128 - i - 1];
            let input_half_top = data[128 + i];

            //prepare the inner DCT2
            unsafe { *input_dct2.get_unchecked_mut(i) = input_top + input_bottom };
            unsafe {
                *input_dct2.get_unchecked_mut(128 - i - 1) = input_half_bottom + input_half_top
            };

            //prepare the inner DCT4 - which consists of two DCT2s of half size
            let lower_dct4 = input_bottom - input_top;
            let upper_dct4 = input_half_bottom - input_half_top;

            let cos_input = fmla(lower_dct4, twiddle.re, upper_dct4 * twiddle.im);
            let sin_input = fmla(upper_dct4, twiddle.re, -lower_dct4 * twiddle.im);

            unsafe { *input_dct4_even.get_unchecked_mut(i) = cos_input };
            unsafe {
                *input_dct4_odd.get_unchecked_mut(64 - i - 1) =
                    if i % 2 == 0 { sin_input } else { -sin_input }
            };
        }

        _ = self.bf128.execute(input_dct2);
        _ = self.bf64.execute(input_dct4);

        let (input_dct4_even, input_dct4_odd) = input_dct4.split_at_mut(64);

        unsafe {
            data[0] = *input_dct2.get_unchecked(0);
            data[1] = *input_dct4_even.get_unchecked(0);
            data[2] = *input_dct2.get_unchecked(1);

            for i in 1..64 {
                let dct4_cos_output = *input_dct4_even.get_unchecked(i);
                let dct4_sin_output = if (i + 64) % 2 == 0 {
                    -*input_dct4_odd.get_unchecked(64 - i)
                } else {
                    *input_dct4_odd.get_unchecked(64 - i)
                };

                *data.get_unchecked_mut(i * 4 - 1) = dct4_cos_output + dct4_sin_output;
                *data.get_unchecked_mut(i * 4) = *input_dct2.get_unchecked(i * 2);

                *data.get_unchecked_mut(i * 4 + 1) = dct4_cos_output - dct4_sin_output;
                *data.get_unchecked_mut(i * 4 + 2) = *input_dct2.get_unchecked(i * 2 + 1);
            }

            *data.get_unchecked_mut(256 - 1) = -*input_dct4_odd.get_unchecked(0);
        }
    }
}

impl<T: DctSample> PxdctExecutor<T> for Dct2Butterfly256<T>
where
    f64: AsPrimitive<T>,
{
    fn execute(&self, data: &mut [T]) -> Result<(), PxdctError> {
        if !data.len().is_multiple_of(256) {
            return Err(PxdctError::InvalidSizeMultiplier(data.len(), self.length()));
        }

        let mut input_dct2 = [T::zero(); 128];
        let mut input_dct4 = [T::zero(); 128];

        for chunk in data.chunks_exact_mut(256) {
            self.exec(
                (&mut chunk[..256]).try_into().unwrap(),
                &mut input_dct2,
                &mut input_dct4,
            );
        }
        Ok(())
    }

    fn length(&self) -> usize {
        256
    }
}

#[derive(Clone)]
pub(crate) struct Dct2Butterfly512<T: DctSample> {
    bf128: Arc<dyn PxdctExecutor<T> + Send + Sync>,
    bf256: Arc<dyn PxdctExecutor<T> + Send + Sync>,
    twiddles: [Complex<T>; 128],
}

impl<T: DctSample + Dct2Factory> Default for Dct2Butterfly512<T>
where
    f64: AsPrimitive<T>,
    usize: AsPrimitive<T>,
{
    fn default() -> Self {
        use crate::twiddles::compute_twiddle;
        let mut twiddles = [Complex::<T>::default(); 128];
        for (i, twiddle) in twiddles.iter_mut().enumerate() {
            *twiddle = compute_twiddle::<T>(2 * i + 1, 512 * 4).conj();
        }
        Self {
            twiddles,
            bf256: T::dct2_butterfly256(),
            bf128: T::dct2_butterfly128(),
        }
    }
}

impl<T: DctSample> Dct2Butterfly512<T>
where
    f64: AsPrimitive<T>,
{
    #[inline(always)]
    fn exec(&self, data: &mut [T; 512], input_dct2: &mut [T; 256], input_dct4: &mut [T; 256]) {
        let (input_dct4_even, input_dct4_odd) = input_dct4.split_at_mut(128);

        //preprocess the data by splitting it up into vectors of size n/2, n/4, and n/4

        for i in 0..128 {
            let twiddle = self.twiddles[i];
            let input_bottom = data[i];
            let input_top = data[512 - i - 1];

            let input_half_bottom = data[256 - i - 1];
            let input_half_top = data[256 + i];

            //prepare the inner DCT2
            unsafe { *input_dct2.get_unchecked_mut(i) = input_top + input_bottom };
            unsafe {
                *input_dct2.get_unchecked_mut(256 - i - 1) = input_half_bottom + input_half_top
            };

            //prepare the inner DCT4 - which consists of two DCT2s of half size
            let lower_dct4 = input_bottom - input_top;
            let upper_dct4 = input_half_bottom - input_half_top;

            let cos_input = fmla(lower_dct4, twiddle.re, upper_dct4 * twiddle.im);
            let sin_input = fmla(upper_dct4, twiddle.re, -lower_dct4 * twiddle.im);

            unsafe { *input_dct4_even.get_unchecked_mut(i) = cos_input };
            unsafe {
                *input_dct4_odd.get_unchecked_mut(128 - i - 1) =
                    if i % 2 == 0 { sin_input } else { -sin_input }
            };
        }

        _ = self.bf256.execute(input_dct2);
        _ = self.bf128.execute(input_dct4);

        let (input_dct4_even, input_dct4_odd) = input_dct4.split_at_mut(128);

        unsafe {
            data[0] = *input_dct2.get_unchecked(0);
            data[1] = *input_dct4_even.get_unchecked(0);
            data[2] = *input_dct2.get_unchecked(1);

            for i in 1..128 {
                let dct4_cos_output = *input_dct4_even.get_unchecked(i);
                let dct4_sin_output = if (i + 128) % 2 == 0 {
                    -*input_dct4_odd.get_unchecked(128 - i)
                } else {
                    *input_dct4_odd.get_unchecked(128 - i)
                };

                *data.get_unchecked_mut(i * 4 - 1) = dct4_cos_output + dct4_sin_output;
                *data.get_unchecked_mut(i * 4) = *input_dct2.get_unchecked(i * 2);

                *data.get_unchecked_mut(i * 4 + 1) = dct4_cos_output - dct4_sin_output;
                *data.get_unchecked_mut(i * 4 + 2) = *input_dct2.get_unchecked(i * 2 + 1);
            }

            *data.get_unchecked_mut(512 - 1) = -*input_dct4_odd.get_unchecked(0);
        }
    }
}

impl<T: DctSample> PxdctExecutor<T> for Dct2Butterfly512<T>
where
    f64: AsPrimitive<T>,
{
    fn execute(&self, data: &mut [T]) -> Result<(), PxdctError> {
        if !data.len().is_multiple_of(512) {
            return Err(PxdctError::InvalidSizeMultiplier(data.len(), self.length()));
        }

        let mut input_dct2 = [T::zero(); 256];
        let mut input_dct4 = [T::zero(); 256];

        for chunk in data.chunks_exact_mut(512) {
            self.exec(
                (&mut chunk[..512]).try_into().unwrap(),
                &mut input_dct2,
                &mut input_dct4,
            );
        }
        Ok(())
    }

    fn length(&self) -> usize {
        512
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::PxdctExecutor;
    use crate::butterflies::gen_test_butterfly;
    use crate::tests::{naive_dct2, naive_dst2};
    use rand::Rng;

    gen_test_butterfly!(test_bf_dst2, f64, Dst2Butterfly2, 2, 1e-7, naive_dst2);
    gen_test_butterfly!(test_bf_dst4, f64, Dst2Butterfly4, 4, 1e-7, naive_dst2);

    gen_test_butterfly!(test_bf2, f64, Dct2Butterfly2, 2, 1e-7, naive_dct2);
    gen_test_butterfly!(test_bf4, f64, Dct2Butterfly4, 4, 1e-7, naive_dct2);
    gen_test_butterfly!(test_bf8, f64, Dct2Butterfly8, 8, 1e-7, naive_dct2);
    gen_test_butterfly!(test_bf16, f64, Dct2Butterfly16, 16, 1e-7, naive_dct2);
    gen_test_butterfly!(test_bf32, f64, Dct2Butterfly32, 32, 1e-7, naive_dct2);
    gen_test_butterfly!(test_bf64, f64, Dct2Butterfly64, 64, 1e-7, naive_dct2);
    gen_test_butterfly!(test_bf128, f64, Dct2Butterfly128, 128, 1e-7, naive_dct2);
    gen_test_butterfly!(test_bf256, f64, Dct2Butterfly256, 256, 1e-7, naive_dct2);
    gen_test_butterfly!(test_bf512, f64, Dct2Butterfly512, 512, 1e-7, naive_dct2);
}