pxdct 0.3.3

/*
 * // Copyright (c) Radzivon Bartoshyk 1/2026. 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.
 */
use crate::bidirectional::{BidirectionalStore, InPlaceStore};
use crate::factory_dct4::Dct4Factory;
use crate::mla::fmla;
use crate::type4::radixq_dct4_rotation_twiddle;
use crate::util::DctSample;
use crate::{PxdctError, PxdctExecutor};
use num_complex::Complex;
use num_traits::AsPrimitive;
use std::sync::Arc;

pub(crate) trait Dct4MixedRadix9Sample {
    const D4_R9_ROT_TWIDDLE_0: Self;
    const D4_R9_ROT_TWIDDLE_1: Self;
    const D4_R9_ROT_TWIDDLE_2: Self;
    const D4_R9_ROT_TWIDDLE_3: Self;
    const D4_R9_ROT_TWIDDLE_4: Self;
    const D4_R9_ROT_TWIDDLE_5: Self;
    const D4_R9_ROT_TWIDDLE_6: Self;
}

impl Dct4MixedRadix9Sample for f32 {
    const D4_R9_ROT_TWIDDLE_0: f32 = f32::from_bits(0x3f441b7d);
    const D4_R9_ROT_TWIDDLE_1: f32 = f32::from_bits(0x3f248dbb);
    const D4_R9_ROT_TWIDDLE_2: f32 = f32::from_bits(0x3f708fb2);
    const D4_R9_ROT_TWIDDLE_3: f32 = f32::from_bits(0x3eaf1d44);
    const D4_R9_ROT_TWIDDLE_4: f32 = f32::from_bits(0x3f5db3d7);
    const D4_R9_ROT_TWIDDLE_5: f32 = f32::from_bits(0xbe31d0d4);
    const D4_R9_ROT_TWIDDLE_6: f32 = f32::from_bits(0x3f7c1c5c);
}

impl Dct4MixedRadix9Sample for f64 {
    const D4_R9_ROT_TWIDDLE_0: f64 = f64::from_bits(0x3fe8836fa2cf5039);
    const D4_R9_ROT_TWIDDLE_1: f64 = f64::from_bits(0x3fe491b7523c161c);
    const D4_R9_ROT_TWIDDLE_2: f64 = f64::from_bits(0x3fee11f642522d1c);
    const D4_R9_ROT_TWIDDLE_3: f64 = f64::from_bits(0x3fd5e3a8748a0bf5);
    const D4_R9_ROT_TWIDDLE_4: f64 = f64::from_bits(0x3febb67ae8584cab);
    const D4_R9_ROT_TWIDDLE_5: f64 = f64::from_bits(0xbfc63a1a7e0b738c);
    const D4_R9_ROT_TWIDDLE_6: f64 = f64::from_bits(0x3fef838b8c811c17);
}

#[derive(Debug, Clone)]
pub(crate) struct Dct4Butterfly9<T> {
    _phantom: std::marker::PhantomData<T>,
}

impl<T: DctSample> Default for Dct4Butterfly9<T>
where
    f64: AsPrimitive<T>,
{
    fn default() -> Self {
        Self {
            _phantom: std::marker::PhantomData,
        }
    }
}

impl<T: DctSample + Dct4MixedRadix9Sample> Dct4Butterfly9<T> {
    #[inline(always)]
    pub(crate) fn exec<S: BidirectionalStore<T>>(&self, data: &mut S) {
        // This is autogenerated Radix-Q where Q = 9.
        let a_buffer = [data[4] * T::FRAC_1_SQRT_2];
        let c_buffer = [
            (data[0] + data[8]) * T::FRAC_1_SQRT_2,
            (data[1] + data[7]) * T::FRAC_1_SQRT_2,
            (data[2] + data[6]) * T::FRAC_1_SQRT_2,
            (data[3] + data[5]) * T::FRAC_1_SQRT_2,
        ];
        let s_buffer = [
            (data[0] - data[8]) * T::FRAC_1_SQRT_2,
            (data[1] - data[7]) * T::FRAC_1_SQRT_2,
            (data[2] - data[6]) * T::FRAC_1_SQRT_2,
            (data[3] - data[5]) * T::FRAC_1_SQRT_2,
        ];
        let c0 = c_buffer[0];
        let s0 = s_buffer[0];
        let s0_tw0 = s0 * T::D4_R9_ROT_TWIDDLE_1;
        let s0_tw1 = -s0 * T::D4_R9_ROT_TWIDDLE_0;
        let iq0 = fmla(c0, T::D4_R9_ROT_TWIDDLE_0, s0_tw0);
        let siq0 = fmla(c0, T::D4_R9_ROT_TWIDDLE_1, s0_tw1);
        let mut u0 = iq0;
        let mut u1 = u0;
        let mut v0 = siq0;
        u1 *= T::D4_R9_ROT_TWIDDLE_2;
        v0 *= T::D4_R9_ROT_TWIDDLE_3;
        let c1 = c_buffer[1];
        let s1 = s_buffer[1];
        let iq1 = fmla(c1, T::D4_R9_ROT_TWIDDLE_4, s1 * T::HALF);
        let siq1 = fmla(c1, T::HALF, -s1 * T::D4_R9_ROT_TWIDDLE_4);
        u0 += iq1;
        u1 = fmla(iq1, T::HALF, u1);
        v0 = fmla(siq1, T::D4_R9_ROT_TWIDDLE_4, v0);
        let c2 = c_buffer[2];
        let s2 = s_buffer[2];
        let iq2 = fmla(c2, T::D4_R9_ROT_TWIDDLE_2, s2 * T::D4_R9_ROT_TWIDDLE_3);
        let siq2 = fmla(c2, T::D4_R9_ROT_TWIDDLE_3, -s2 * T::D4_R9_ROT_TWIDDLE_2);
        u0 += iq2;
        u1 = fmla(iq2, T::D4_R9_ROT_TWIDDLE_5, u1);
        v0 = fmla(siq2, T::D4_R9_ROT_TWIDDLE_6, v0);
        let c3 = c_buffer[3];
        let s3 = s_buffer[3];
        let iq3 = fmla(c3, T::D4_R9_ROT_TWIDDLE_6, s3 * -T::D4_R9_ROT_TWIDDLE_5);
        let siq3 = fmla(c3, -T::D4_R9_ROT_TWIDDLE_5, -s3 * T::D4_R9_ROT_TWIDDLE_6);
        u0 += iq3;
        u1 = fmla(iq3, -T::D4_R9_ROT_TWIDDLE_0, u1);
        v0 = fmla(siq3, T::D4_R9_ROT_TWIDDLE_1, v0);
        let a0 = a_buffer[0];
        u0 += a0;
        u1 = u1 - a0;
        data[0] = u0;
        data[1] = u1 - v0;
        data[2] = u1 + v0;
        let mut u2 = iq0;
        let mut v2 = siq0;
        u2 *= T::D4_R9_ROT_TWIDDLE_0;
        v2 *= T::D4_R9_ROT_TWIDDLE_1;
        u2 = fmla(iq1, -T::HALF, u2);
        v2 = fmla(siq1, T::D4_R9_ROT_TWIDDLE_4, v2);
        u2 = fmla(iq2, -T::D4_R9_ROT_TWIDDLE_2, u2);
        v2 = fmla(siq2, -T::D4_R9_ROT_TWIDDLE_3, v2);
        u2 = fmla(iq3, -T::D4_R9_ROT_TWIDDLE_5, u2);
        v2 = fmla(siq3, -T::D4_R9_ROT_TWIDDLE_6, v2);
        u2 += a0;
        data[3] = u2 - v2;
        data[4] = u2 + v2;
        let mut u3 = iq0;
        let mut v3 = siq0;
        u3 *= T::HALF;
        v3 *= T::D4_R9_ROT_TWIDDLE_4;
        u3 = u3 - iq1;
        u3 = fmla(iq2, T::HALF, u3);
        v3 = fmla(siq2, -T::D4_R9_ROT_TWIDDLE_4, v3);
        u3 = fmla(iq3, T::HALF, u3);
        v3 = fmla(siq3, T::D4_R9_ROT_TWIDDLE_4, v3);
        u3 = u3 - a0;
        data[5] = u3 - v3;
        data[6] = u3 + v3;
        let mut u4 = iq0;
        let mut v4 = siq0;
        u4 *= -T::D4_R9_ROT_TWIDDLE_5;
        v4 *= T::D4_R9_ROT_TWIDDLE_6;
        u4 = fmla(iq1, -T::HALF, u4);
        v4 = fmla(siq1, -T::D4_R9_ROT_TWIDDLE_4, v4);
        u4 = fmla(iq2, T::D4_R9_ROT_TWIDDLE_0, u4);
        v4 = fmla(siq2, T::D4_R9_ROT_TWIDDLE_1, v4);
        u4 = fmla(iq3, -T::D4_R9_ROT_TWIDDLE_2, u4);
        v4 = fmla(siq3, -T::D4_R9_ROT_TWIDDLE_3, v4);
        u4 += a0;
        data[7] = u4 - v4;
        data[8] = u4 + v4;
    }
}

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

    fn execute_with_scratch(&self, data: &mut [T], _: &mut [T]) -> Result<(), PxdctError> {
        self.execute(data)
    }

    fn execute_into(&self, input: &[T], output: &mut [T]) -> Result<(), PxdctError> {
        self.execute_into_with_scratch(input, output, &mut [])
    }

    fn execute_into_with_scratch(
        &self,
        input: &[T],
        output: &mut [T],
        _: &mut [T],
    ) -> Result<(), PxdctError> {
        use crate::util::validate_oof_sizes;
        validate_oof_sizes!(input, output, 9);
        use crate::bidirectional::BiStore;
        for (src, dst) in input.chunks_exact(9).zip(output.chunks_exact_mut(9)) {
            self.exec(&mut BiStore::new(src, dst));
        }
        Ok(())
    }

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

    fn scratch_size(&self) -> usize {
        0
    }
}

#[derive(Debug, Clone)]
pub(crate) struct Dct4Butterfly27<T> {
    twiddles: [Complex<T>; 9],
    bf9: Dct4Butterfly9<T>,
}

impl<T: DctSample> Default for Dct4Butterfly27<T>
where
    f64: AsPrimitive<T>,
    usize: AsPrimitive<T>,
{
    fn default() -> Self {
        let mut twiddles = [Complex::<T>::default(); 9];
        for (k, dst) in twiddles.iter_mut().enumerate() {
            *dst = radixq_dct4_rotation_twiddle(3, 0, k, 27);
        }
        Self {
            twiddles,
            bf9: Dct4Butterfly9::default(),
        }
    }
}

impl<T: DctSample + Dct4MixedRadix9Sample> Dct4Butterfly27<T> {
    #[inline(always)]
    pub(crate) fn exec<S: BidirectionalStore<T>>(&self, data: &mut S) {
        let mut a_buffer = [
            data[1], data[4], data[7], data[10], data[13], data[16], data[19], data[22], data[25],
        ];
        let mut c_buffer = [
            data[0] + data[2],
            data[3] + data[5],
            data[6] + data[8],
            data[9] + data[11],
            data[12] + data[14],
            data[15] + data[17],
            data[18] + data[20],
            data[21] + data[23],
            data[24] + data[26],
        ];
        let mut s_buffer = [
            data[0] - data[2],
            data[5] - data[3],
            data[6] - data[8],
            data[11] - data[9],
            data[12] - data[14],
            data[17] - data[15],
            data[18] - data[20],
            data[23] - data[21],
            data[24] - data[26],
        ];

        self.bf9
            .exec(&mut InPlaceStore::new(a_buffer.as_mut_slice()));
        self.bf9
            .exec(&mut InPlaceStore::new(c_buffer.as_mut_slice()));
        self.bf9
            .exec(&mut InPlaceStore::new(s_buffer.as_mut_slice()));

        let q_modules = 27 / 3;
        let s = 2 * 27 / 3;

        for k in 0..q_modules {
            let c_v = c_buffer[k];
            let s_v = s_buffer[q_modules - 1 - k];
            let a_v = a_buffer[k];

            let twiddle = self.twiddles[k];

            let mut u0 = fmla(c_v, twiddle.re, s_v * twiddle.im);
            let mut u1 = u0;
            let mut v0 = fmla(c_v, twiddle.im, -s_v * twiddle.re);

            u0 += a_v;
            u1 *= T::HALF;
            v0 *= T::SQRT_3_OVER_2;
            u1 = u1 - a_v;

            let uc0 = u1 - v0;
            let uc1 = u1 + v0;

            data[k] = u0;
            data[s + k] = uc1;
            data[s - 1 - k] = uc0;
        }
    }
}

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

    fn execute_with_scratch(&self, data: &mut [T], _: &mut [T]) -> Result<(), PxdctError> {
        self.execute(data)
    }

    fn execute_into(&self, input: &[T], output: &mut [T]) -> Result<(), PxdctError> {
        self.execute_into_with_scratch(input, output, &mut [])
    }

    fn execute_into_with_scratch(
        &self,
        input: &[T],
        output: &mut [T],
        _: &mut [T],
    ) -> Result<(), PxdctError> {
        use crate::util::validate_oof_sizes;
        validate_oof_sizes!(input, output, 27);
        use crate::bidirectional::BiStore;
        for (src, dst) in input.chunks_exact(27).zip(output.chunks_exact_mut(27)) {
            self.exec(&mut BiStore::new(src, dst));
        }
        Ok(())
    }

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

    fn scratch_size(&self) -> usize {
        0
    }
}

#[derive(Clone)]
pub(crate) struct Dct4Butterfly81<T> {
    twiddles: [Complex<T>; 27],
    bf27: Arc<dyn PxdctExecutor<T> + Send + Sync>,
}

impl<T: DctSample + Dct4Factory> Default for Dct4Butterfly81<T>
where
    f64: AsPrimitive<T>,
    usize: AsPrimitive<T>,
{
    fn default() -> Self {
        let mut twiddles = [Complex::<T>::default(); 27];
        for (k, dst) in twiddles.iter_mut().enumerate() {
            *dst = radixq_dct4_rotation_twiddle(3, 0, k, 81);
        }
        Self {
            twiddles,
            bf27: T::dct4_butterfly27(),
        }
    }
}

impl<T: DctSample + Dct4MixedRadix9Sample> Dct4Butterfly81<T> {
    #[inline(always)]
    pub(crate) fn exec<S: BidirectionalStore<T>>(&self, data: &mut S) {
        let mut a_buffer = [T::zero(); 27];
        let mut c_buffer = [T::zero(); 27];
        let mut s_buffer = [T::zero(); 27];

        for n in 0..27 {
            a_buffer[n] = data[n * 3 + 1];
        }

        let mut sign = T::one();
        for n in 0..27 {
            let u0 = data[3 * n];
            let u1 = data[3 * n + 3 - 1];
            c_buffer[n] = u0 + u1;
            s_buffer[n] = (u0 - u1).mulsign(sign);

            sign = -sign;
        }

        _ = self.bf27.execute(&mut a_buffer);
        _ = self.bf27.execute(&mut c_buffer);
        _ = self.bf27.execute(&mut s_buffer);

        let q_modules = 81 / 3;
        let s = 2 * 81 / 3;

        for k in 0..q_modules {
            let c_v = c_buffer[k];
            let s_v = s_buffer[q_modules - 1 - k];
            let a_v = a_buffer[k];

            let twiddle = self.twiddles[k];

            let mut u0 = fmla(c_v, twiddle.re, s_v * twiddle.im);
            let mut u1 = u0;
            let mut v0 = fmla(c_v, twiddle.im, -s_v * twiddle.re);

            u0 += a_v;
            u1 *= T::HALF;
            v0 *= T::SQRT_3_OVER_2;
            u1 = u1 - a_v;

            let uc0 = u1 - v0;
            let uc1 = u1 + v0;

            data[k] = u0;
            data[s + k] = uc1;
            data[s - 1 - k] = uc0;
        }
    }
}

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

    fn execute_with_scratch(&self, data: &mut [T], _: &mut [T]) -> Result<(), PxdctError> {
        self.execute(data)
    }

    fn execute_into(&self, input: &[T], output: &mut [T]) -> Result<(), PxdctError> {
        self.execute_into_with_scratch(input, output, &mut [])
    }

    fn execute_into_with_scratch(
        &self,
        input: &[T],
        output: &mut [T],
        _: &mut [T],
    ) -> Result<(), PxdctError> {
        use crate::util::validate_oof_sizes;
        validate_oof_sizes!(input, output, 81);
        use crate::bidirectional::BiStore;
        for (src, dst) in input.chunks_exact(81).zip(output.chunks_exact_mut(81)) {
            self.exec(&mut BiStore::new(src, dst));
        }
        Ok(())
    }

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

    fn scratch_size(&self) -> usize {
        0
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::butterflies::gen_test_butterfly;

    gen_test_butterfly!(test_bf_dct4_9, f64, Dct4Butterfly9, 9, 1e-7, naive_dct4);
    gen_test_butterfly!(test_bf_dct4_27, f64, Dct4Butterfly27, 27, 1e-7, naive_dct4);
    gen_test_butterfly!(test_bf_dct4_81, f64, Dct4Butterfly81, 81, 1e-7, naive_dct4);
}