jpegli-rs 0.12.0

//! Forward Discrete Cosine Transform (DCT) implementation.
//!
//! This module provides 8x8 DCT-II transform used for JPEG encoding.
//!
//! Two algorithms are available:
//! - **Recursive (jpegli)**: Uses recursive splitting for exact C++ jpegli compatibility
//! - **AAN (libjpeg)**: Arai-Agui-Nakajima algorithm with only 5 multiplies per 1D DCT
//!
//! SIMD optimization via the `wide` crate is always enabled.
//!
//! # Safe/Unsafe Architecture
//!
//! - **Safe public APIs**: `forward_dct_8x8`, `transpose_8x8_simd`
//! - **Unsafe internal functions**: `forward_dct_8x8_fma`, `transpose_8x8_avx`
//!
//! The unsafe AVX/FMA functions use raw intrinsics for 8x8 matrix transpose
//! (`_mm256_unpacklo/hi_ps`, `_mm256_permute2f128_ps`) which have no `wide` equivalent.

use crate::foundation::consts::DCT_BLOCK_SIZE;
use crate::foundation::simd_types::Block8x8f;
use multiversed::multiversed;
use wide::f32x8;

// ============================================================================
// Recursive DCT (jpegli-compatible)
// ============================================================================

/// WcMultipliers constants for DCT.
/// Generated by: 1.0 / (2 * cos((i + 0.5) * pi / N))
/// (Used by scalar reference implementation)
#[allow(dead_code)]
const WC4: [f32; 2] = [0.541196100146197, 1.3065629648763764];

#[allow(dead_code)]
const WC8: [f32; 4] = [
    0.5097955791041592,
    0.6013448869350453,
    0.8999762231364156,
    2.5629154477415055,
];

#[allow(dead_code)]
const SQRT2: f32 = 1.41421356237;

// ============================================================================
// AAN DCT (Arai-Agui-Nakajima) - 5 multiplies + 29 adds per 1D DCT
// ============================================================================

/// AAN DCT constants from libjpeg jfdctflt.c
/// Based on: Y. Arai, T. Agui, M. Nakajima, "A fast DCT-SQ scheme for images", 1988
mod aan {
    /// cos(π/4) = √2/2
    pub const C4: f32 = 0.707106781;
    /// sin(π/8)
    pub const C6: f32 = 0.382683433;
    /// cos(π/8) - cos(3π/8)
    pub const C2_M_C6: f32 = 0.541196100;
    /// cos(π/8) + cos(3π/8)
    pub const C2_P_C6: f32 = 1.306562965;

    /// AAN scale factors (aanscales from libjpeg)
    /// The AAN algorithm produces scaled output; divide by these to get standard DCT.
    /// aanscales[k] = cos(k * π/16) * √2 for k > 0, and 1.0 for k = 0
    /// (Kept for reference; INV_SCALES is used in production for multiply instead of divide)
    #[allow(dead_code)]
    pub const SCALES: [f32; 8] = [
        1.0,
        1.387039845,
        1.306562965,
        1.175875602,
        1.0,
        0.785694958,
        0.541196100,
        0.275899379,
    ];

    /// Reciprocal of scale factors (for multiplication instead of division)
    pub const INV_SCALES: [f32; 8] = [
        1.0,
        1.0 / 1.387039845,
        1.0 / 1.306562965,
        1.0 / 1.175875602,
        1.0,
        1.0 / 0.785694958,
        1.0 / 0.541196100,
        1.0 / 0.275899379,
    ];
}

/// Transpose an 8x8 block (scalar version, used in tests).
#[cfg(test)]
#[inline]
fn transpose_8x8(input: &[f32; 64], output: &mut [f32; 64]) {
    for row in 0..8 {
        for col in 0..8 {
            output[col * 8 + row] = input[row * 8 + col];
        }
    }
}

/// AddReverse: out[i] = in1[i] + in2[N-1-i]
/// (Scalar reference implementation)
#[allow(dead_code)]
#[inline]
fn add_reverse<const N: usize>(in1: &[f32], in2: &[f32], out: &mut [f32]) {
    for i in 0..N {
        out[i] = in1[i] + in2[N - 1 - i];
    }
}

/// SubReverse: out[i] = in1[i] - in2[N-1-i]
/// (Scalar reference implementation)
#[allow(dead_code)]
#[inline]
fn sub_reverse<const N: usize>(in1: &[f32], in2: &[f32], out: &mut [f32]) {
    for i in 0..N {
        out[i] = in1[i] - in2[N - 1 - i];
    }
}

/// B transform: coeff[0] = coeff[0] * sqrt2 + coeff[1]; then cumulative sum
/// (Scalar reference implementation)
#[allow(dead_code)]
#[inline]
fn b_transform<const N: usize>(coeff: &mut [f32]) {
    coeff[0] = coeff[0] * SQRT2 + coeff[1];
    for i in 1..(N - 1) {
        coeff[i] += coeff[i + 1];
    }
}

/// Multiply odd part by Wc multipliers
/// (Scalar reference implementation)
#[allow(dead_code)]
#[inline]
fn multiply_wc8(coeff: &mut [f32]) {
    for i in 0..4 {
        coeff[4 + i] *= WC8[i];
    }
}

/// (Scalar reference implementation)
#[allow(dead_code)]
#[inline]
fn multiply_wc4(coeff: &mut [f32]) {
    for i in 0..2 {
        coeff[2 + i] *= WC4[i];
    }
}

/// InverseEvenOdd: Interleave even/odd parts
/// out[2*i] = in[i] for i in 0..N/2
/// out[2*i+1] = in[N/2+i] for i in 0..N/2
/// (Scalar reference implementation)
#[allow(dead_code)]
#[inline]
fn inverse_even_odd<const N: usize>(input: &[f32], output: &mut [f32]) {
    let half = N / 2;
    for i in 0..half {
        output[2 * i] = input[i];
        output[2 * i + 1] = input[half + i];
    }
}

/// DCT base case for N=2
/// (Scalar reference implementation)
#[allow(dead_code)]
#[inline]
fn dct1d_2(mem: &mut [f32]) {
    let in1 = mem[0];
    let in2 = mem[1];
    mem[0] = in1 + in2;
    mem[1] = in1 - in2;
}

/// DCT for N=4 (recursive)
/// (Scalar reference implementation)
#[allow(dead_code)]
#[inline]
fn dct1d_4(mem: &mut [f32]) {
    let mut tmp = [0.0f32; 4];

    // AddReverse<2>: tmp[0:2] = mem[0:2] + reverse(mem[2:4])
    add_reverse::<2>(&mem[0..2], &mem[2..4], &mut tmp[0..2]);

    // DCT1D<2> on first half
    dct1d_2(&mut tmp[0..2]);

    // SubReverse<2>: tmp[2:4] = mem[0:2] - reverse(mem[2:4])
    sub_reverse::<2>(&mem[0..2], &mem[2..4], &mut tmp[2..4]);

    // Multiply<4>: multiply tmp[2:4] by WC4
    multiply_wc4(&mut tmp);

    // DCT1D<2> on second half
    dct1d_2(&mut tmp[2..4]);

    // B<2>: just coeff[0] *= sqrt2 + coeff[1] (no loop since N-1=1)
    tmp[2] = tmp[2] * SQRT2 + tmp[3];

    // InverseEvenOdd<4>
    inverse_even_odd::<4>(&tmp, mem);
}

/// DCT for N=8 (recursive)
/// (Scalar reference implementation)
#[allow(dead_code)]
#[inline]
fn dct1d_8(mem: &mut [f32]) {
    let mut tmp = [0.0f32; 8];

    // AddReverse<4>: tmp[0:4] = mem[0:4] + reverse(mem[4:8])
    add_reverse::<4>(&mem[0..4], &mem[4..8], &mut tmp[0..4]);

    // DCT1D<4> on first half
    dct1d_4(&mut tmp[0..4]);

    // SubReverse<4>: tmp[4:8] = mem[0:4] - reverse(mem[4:8])
    sub_reverse::<4>(&mem[0..4], &mem[4..8], &mut tmp[4..8]);

    // Multiply<8>: multiply tmp[4:8] by WC8
    multiply_wc8(&mut tmp);

    // DCT1D<4> on second half
    dct1d_4(&mut tmp[4..8]);

    // B<4>: coeff[0] = coeff[0] * sqrt2 + coeff[1]; then cumulative sum
    b_transform::<4>(&mut tmp[4..8]);

    // InverseEvenOdd<8>
    inverse_even_odd::<8>(&tmp, mem);
}

// SIMD-optimized implementations (always available via `wide` crate)
pub(crate) mod simd {
    use super::*;

    // Raw AVX2 intrinsics - only available with `unsafe_simd` feature
    #[cfg(all(feature = "unsafe_simd", target_arch = "x86_64"))]
    use core::arch::x86_64::{
        __m256, _mm256_add_ps, _mm256_fmadd_ps, _mm256_loadu_ps, _mm256_mul_ps,
        _mm256_permute2f128_ps, _mm256_set1_ps, _mm256_storeu_ps, _mm256_sub_ps,
        _mm256_unpackhi_ps, _mm256_unpacklo_ps,
    };

    // SIMD versions of WC constants for parallel DCT
    // Note: The parallel DCT approach has poor cache locality and doesn't
    // improve performance compared to row-by-row processing. Kept for reference.
    #[allow(dead_code)]
    const WC4_0: f32x8 = f32x8::new([0.541196100146197; 8]);
    #[allow(dead_code)]
    const WC4_1: f32x8 = f32x8::new([1.3065629648763764; 8]);

    #[allow(dead_code)]
    const WC8_0: f32x8 = f32x8::new([0.5097955791041592; 8]);
    #[allow(dead_code)]
    const WC8_1: f32x8 = f32x8::new([0.6013448869350453; 8]);
    #[allow(dead_code)]
    const WC8_2: f32x8 = f32x8::new([0.8999762231364156; 8]);
    #[allow(dead_code)]
    const WC8_3: f32x8 = f32x8::new([2.5629154477415055; 8]);

    #[allow(dead_code)]
    const SQRT2_VEC: f32x8 = f32x8::new([1.41421356237; 8]);

    /// DCT base case for N=2 on SIMD vectors
    /// Processes 8 independent 2-point DCTs in parallel
    #[allow(dead_code)]
    #[inline]
    fn dct1d_2_simd(m0: &mut f32x8, m1: &mut f32x8) {
        let in0 = *m0;
        let in1 = *m1;
        *m0 = in0 + in1;
        *m1 = in0 - in1;
    }

    /// DCT for N=4 on SIMD vectors
    /// Processes 8 independent 4-point DCTs in parallel
    #[allow(dead_code)]
    #[inline]
    fn dct1d_4_simd(m: &mut [f32x8; 4]) {
        // AddReverse<2>: tmp[0:2] = m[0:2] + reverse(m[2:4])
        let t0 = m[0] + m[3]; // m[0] + m[4-1-0]
        let t1 = m[1] + m[2]; // m[1] + m[4-1-1]

        // SubReverse<2>: tmp[2:4] = m[0:2] - reverse(m[2:4])
        let t2 = m[0] - m[3];
        let t3 = m[1] - m[2];

        // DCT1D<2> on first half (t0, t1)
        let r0 = t0 + t1;
        let r1 = t0 - t1;

        // Multiply by WC4
        let t2_scaled = t2 * WC4_0;
        let t3_scaled = t3 * WC4_1;

        // DCT1D<2> on second half
        let r2 = t2_scaled + t3_scaled;
        let r3 = t2_scaled - t3_scaled;

        // B<2>: r2 = r2 * sqrt2 + r3
        let r2_final = r2 * SQRT2_VEC + r3;

        // InverseEvenOdd<4>: interleave even/odd
        m[0] = r0; // even[0]
        m[1] = r2_final; // odd[0]
        m[2] = r1; // even[1]
        m[3] = r3; // odd[1]
    }

    /// DCT for N=8 on SIMD vectors
    /// Processes 8 independent 8-point DCTs in parallel
    #[allow(dead_code)]
    #[inline]
    fn dct1d_8_simd(m: &mut [f32x8; 8]) {
        // AddReverse<4>: tmp[0:4] = m[0:4] + reverse(m[4:8])
        let t0 = m[0] + m[7];
        let t1 = m[1] + m[6];
        let t2 = m[2] + m[5];
        let t3 = m[3] + m[4];

        // SubReverse<4>: tmp[4:8] = m[0:4] - reverse(m[4:8])
        let t4 = m[0] - m[7];
        let t5 = m[1] - m[6];
        let t6 = m[2] - m[5];
        let t7 = m[3] - m[4];

        // DCT1D<4> on first half (t0-t3)
        let mut first = [t0, t1, t2, t3];
        dct1d_4_simd(&mut first);

        // Multiply by WC8
        let t4_scaled = t4 * WC8_0;
        let t5_scaled = t5 * WC8_1;
        let t6_scaled = t6 * WC8_2;
        let t7_scaled = t7 * WC8_3;

        // DCT1D<4> on second half
        let mut second = [t4_scaled, t5_scaled, t6_scaled, t7_scaled];
        dct1d_4_simd(&mut second);

        // B<4>: coeff[0] = coeff[0] * sqrt2 + coeff[1]; then cumulative sum
        second[0] = second[0] * SQRT2_VEC + second[1];
        second[1] += second[2];
        second[2] += second[3];
        // second[3] stays the same

        // InverseEvenOdd<8>: interleave
        m[0] = first[0];
        m[1] = second[0];
        m[2] = first[1];
        m[3] = second[1];
        m[4] = first[2];
        m[5] = second[2];
        m[6] = first[3];
        m[7] = second[3];
    }

    /// Process 8 rows simultaneously using SIMD with AVX2 transpose.
    /// Uses cache-friendly row loads + fast transpose instead of element-by-element gather.
    #[multiversed]
    #[allow(dead_code)]
    pub fn dct_8rows_parallel(input: &[f32; 64], output: &mut [f32; 64]) {
        // Step 1: Load all 8 rows (cache-friendly sequential access, zero-cost)
        let mut rows: [f32x8; 8] = [f32x8::ZERO; 8];
        for row in 0..8 {
            let k = row * 8;
            let row_slice: [f32; 8] = input[k..k + 8].try_into().unwrap();
            rows[row] = f32x8::from(row_slice);
        }

        // Step 2: Transpose to column-major using AVX2 if available
        // After transpose: mem[col] = [row0[col], row1[col], ..., row7[col]]
        let mut mem: [f32x8; 8];

        #[cfg(all(feature = "unsafe_simd", target_arch = "x86_64"))]
        {
            if is_x86_feature_detected!("avx2") {
                // Use AVX2 transpose on the raw arrays
                let mut input_flat = [0.0f32; 64];
                let mut output_flat = [0.0f32; 64];
                for (i, r) in rows.iter().enumerate() {
                    let arr = r.to_array();
                    input_flat[i * 8..i * 8 + 8].copy_from_slice(&arr);
                }
                unsafe {
                    transpose_8x8_avx(&input_flat, &mut output_flat);
                }
                mem = [f32x8::ZERO; 8];
                for col in 0..8 {
                    let k = col * 8;
                    let col_slice: [f32; 8] = output_flat[k..k + 8].try_into().unwrap();
                    mem[col] = f32x8::from(col_slice);
                }
            } else {
                // Scalar transpose fallback
                mem = transpose_f32x8_array(&rows);
            }
        }

        #[cfg(not(all(feature = "unsafe_simd", target_arch = "x86_64")))]
        {
            mem = transpose_f32x8_array(&rows);
        }

        // Step 3: Execute 8 DCT-1D operations in parallel
        dct1d_8_simd(&mut mem);

        // Step 4: Transpose back to row-major
        #[cfg(all(feature = "unsafe_simd", target_arch = "x86_64"))]
        {
            if is_x86_feature_detected!("avx2") {
                let mut input_flat = [0.0f32; 64];
                for (i, m) in mem.iter().enumerate() {
                    let arr = m.to_array();
                    input_flat[i * 8..i * 8 + 8].copy_from_slice(&arr);
                }
                unsafe {
                    transpose_8x8_avx(&input_flat, output);
                }
                return;
            }
        }

        // Scalar store fallback
        let result = transpose_f32x8_array(&mem);
        for (i, r) in result.iter().enumerate() {
            let arr = r.to_array();
            output[i * 8..i * 8 + 8].copy_from_slice(&arr);
        }
    }

    /// Transpose an 8x8 matrix of f32x8 vectors using wide's built-in transpose.
    /// AVX-accelerated when available.
    #[allow(dead_code)]
    #[inline]
    fn transpose_f32x8_array(input: &[f32x8; 8]) -> [f32x8; 8] {
        f32x8::transpose(*input)
    }

    /// AVX2-optimized 8x8 transpose using Highway's algorithm.
    ///
    /// Uses InterleaveLower/Upper (vunpcklps/vunpckhps) and ConcatLowerLower/UpperUpper
    /// (vperm2f128) to perform the transpose entirely in registers.
    /// Used by dct_8rows_parallel; transpose_8x8_simd has its own inline implementation.
    ///
    /// This is a low-level function. Production code should use the safe `transpose_8x8_simd` wrapper.
    ///
    /// # Safety
    /// Requires AVX2 support.
    #[cfg(all(feature = "unsafe_simd", target_arch = "x86_64"))]
    #[target_feature(enable = "avx2")]
    #[allow(dead_code)]
    #[inline]
    pub(crate) unsafe fn transpose_8x8_avx(input: &[f32; 64], output: &mut [f32; 64]) {
        // Load 8 rows into YMM registers
        let i0 = _mm256_loadu_ps(input.as_ptr());
        let i1 = _mm256_loadu_ps(input.as_ptr().add(8));
        let i2 = _mm256_loadu_ps(input.as_ptr().add(16));
        let i3 = _mm256_loadu_ps(input.as_ptr().add(24));
        let i4 = _mm256_loadu_ps(input.as_ptr().add(32));
        let i5 = _mm256_loadu_ps(input.as_ptr().add(40));
        let i6 = _mm256_loadu_ps(input.as_ptr().add(48));
        let i7 = _mm256_loadu_ps(input.as_ptr().add(56));

        // Phase 1: InterleaveLower/Upper pairs
        // For f32: unpacklo interleaves elements [0,1] from each lane
        //          unpackhi interleaves elements [2,3] from each lane
        let q0 = _mm256_unpacklo_ps(i0, i2); // [i0[0],i2[0],i0[1],i2[1] | i0[4],i2[4],i0[5],i2[5]]
        let q1 = _mm256_unpacklo_ps(i1, i3);
        let q2 = _mm256_unpackhi_ps(i0, i2);
        let q3 = _mm256_unpackhi_ps(i1, i3);
        let q4 = _mm256_unpacklo_ps(i4, i6);
        let q5 = _mm256_unpacklo_ps(i5, i7);
        let q6 = _mm256_unpackhi_ps(i4, i6);
        let q7 = _mm256_unpackhi_ps(i5, i7);

        // Phase 2: Another round of InterleaveLower/Upper
        let r0 = _mm256_unpacklo_ps(q0, q1);
        let r1 = _mm256_unpackhi_ps(q0, q1);
        let r2 = _mm256_unpacklo_ps(q2, q3);
        let r3 = _mm256_unpackhi_ps(q2, q3);
        let r4 = _mm256_unpacklo_ps(q4, q5);
        let r5 = _mm256_unpackhi_ps(q4, q5);
        let r6 = _mm256_unpacklo_ps(q6, q7);
        let r7 = _mm256_unpackhi_ps(q6, q7);

        // Phase 3: ConcatLowerLower/UpperUpper to swap 128-bit lanes
        // vperm2f128 with imm8=0x20: concat low lanes (b_lo, a_lo)
        // vperm2f128 with imm8=0x31: concat high lanes (b_hi, a_hi)
        let o0 = _mm256_permute2f128_ps(r0, r4, 0x20);
        let o1 = _mm256_permute2f128_ps(r1, r5, 0x20);
        let o2 = _mm256_permute2f128_ps(r2, r6, 0x20);
        let o3 = _mm256_permute2f128_ps(r3, r7, 0x20);
        let o4 = _mm256_permute2f128_ps(r0, r4, 0x31);
        let o5 = _mm256_permute2f128_ps(r1, r5, 0x31);
        let o6 = _mm256_permute2f128_ps(r2, r6, 0x31);
        let o7 = _mm256_permute2f128_ps(r3, r7, 0x31);

        // Store transposed rows
        _mm256_storeu_ps(output.as_mut_ptr(), o0);
        _mm256_storeu_ps(output.as_mut_ptr().add(8), o1);
        _mm256_storeu_ps(output.as_mut_ptr().add(16), o2);
        _mm256_storeu_ps(output.as_mut_ptr().add(24), o3);
        _mm256_storeu_ps(output.as_mut_ptr().add(32), o4);
        _mm256_storeu_ps(output.as_mut_ptr().add(40), o5);
        _mm256_storeu_ps(output.as_mut_ptr().add(48), o6);
        _mm256_storeu_ps(output.as_mut_ptr().add(56), o7);
    }

    /// SIMD-optimized 8x8 transpose with compile-time dispatch via multiversion.
    /// Uses AVX2 intrinsics when available, falls back to scalar.
    ///
    /// The AVX2 version uses vunpcklps/vunpckhps and vperm2f128 for efficient
    /// register-to-register transpose (Highway algorithm).
    #[multiversed]
    pub fn transpose_8x8_simd(input: &[f32; 64], output: &mut [f32; 64]) {
        // When compiling for AVX2 target with unsafe_simd feature,
        // use raw intrinsics for maximum performance.
        #[cfg(all(feature = "unsafe_simd", target_arch = "x86_64"))]
        {
            if is_x86_feature_detected!("avx2") {
                // SAFETY: AVX2 is available (checked above). The multiversion macro
                // ensures this branch compiles with target_feature(enable = "avx2").
                unsafe {
                    // Load 8 rows into YMM registers
                    let i0 = _mm256_loadu_ps(input.as_ptr());
                    let i1 = _mm256_loadu_ps(input.as_ptr().add(8));
                    let i2 = _mm256_loadu_ps(input.as_ptr().add(16));
                    let i3 = _mm256_loadu_ps(input.as_ptr().add(24));
                    let i4 = _mm256_loadu_ps(input.as_ptr().add(32));
                    let i5 = _mm256_loadu_ps(input.as_ptr().add(40));
                    let i6 = _mm256_loadu_ps(input.as_ptr().add(48));
                    let i7 = _mm256_loadu_ps(input.as_ptr().add(56));

                    // Phase 1: InterleaveLower/Upper pairs
                    let q0 = _mm256_unpacklo_ps(i0, i2);
                    let q1 = _mm256_unpacklo_ps(i1, i3);
                    let q2 = _mm256_unpackhi_ps(i0, i2);
                    let q3 = _mm256_unpackhi_ps(i1, i3);
                    let q4 = _mm256_unpacklo_ps(i4, i6);
                    let q5 = _mm256_unpacklo_ps(i5, i7);
                    let q6 = _mm256_unpackhi_ps(i4, i6);
                    let q7 = _mm256_unpackhi_ps(i5, i7);

                    // Phase 2: Another round of InterleaveLower/Upper
                    let r0 = _mm256_unpacklo_ps(q0, q1);
                    let r1 = _mm256_unpackhi_ps(q0, q1);
                    let r2 = _mm256_unpacklo_ps(q2, q3);
                    let r3 = _mm256_unpackhi_ps(q2, q3);
                    let r4 = _mm256_unpacklo_ps(q4, q5);
                    let r5 = _mm256_unpackhi_ps(q4, q5);
                    let r6 = _mm256_unpacklo_ps(q6, q7);
                    let r7 = _mm256_unpackhi_ps(q6, q7);

                    // Phase 3: ConcatLowerLower/UpperUpper to swap 128-bit lanes
                    let o0 = _mm256_permute2f128_ps(r0, r4, 0x20);
                    let o1 = _mm256_permute2f128_ps(r1, r5, 0x20);
                    let o2 = _mm256_permute2f128_ps(r2, r6, 0x20);
                    let o3 = _mm256_permute2f128_ps(r3, r7, 0x20);
                    let o4 = _mm256_permute2f128_ps(r0, r4, 0x31);
                    let o5 = _mm256_permute2f128_ps(r1, r5, 0x31);
                    let o6 = _mm256_permute2f128_ps(r2, r6, 0x31);
                    let o7 = _mm256_permute2f128_ps(r3, r7, 0x31);

                    // Store transposed rows
                    _mm256_storeu_ps(output.as_mut_ptr(), o0);
                    _mm256_storeu_ps(output.as_mut_ptr().add(8), o1);
                    _mm256_storeu_ps(output.as_mut_ptr().add(16), o2);
                    _mm256_storeu_ps(output.as_mut_ptr().add(24), o3);
                    _mm256_storeu_ps(output.as_mut_ptr().add(32), o4);
                    _mm256_storeu_ps(output.as_mut_ptr().add(40), o5);
                    _mm256_storeu_ps(output.as_mut_ptr().add(48), o6);
                    _mm256_storeu_ps(output.as_mut_ptr().add(56), o7);
                }
                return;
            }
        }

        // Use wide's built-in transpose (AVX-accelerated when available)
        transpose_8x8_wide(input, output);
    }

    /// Transpose using wide's built-in f32x8::transpose (AVX-accelerated).
    /// (Safe fallback when unsafe_simd feature is disabled)
    #[allow(dead_code)]
    #[inline]
    fn transpose_8x8_wide(input: &[f32; 64], output: &mut [f32; 64]) {
        let rows = [
            f32x8::from(<[f32; 8]>::try_from(&input[0..8]).unwrap()),
            f32x8::from(<[f32; 8]>::try_from(&input[8..16]).unwrap()),
            f32x8::from(<[f32; 8]>::try_from(&input[16..24]).unwrap()),
            f32x8::from(<[f32; 8]>::try_from(&input[24..32]).unwrap()),
            f32x8::from(<[f32; 8]>::try_from(&input[32..40]).unwrap()),
            f32x8::from(<[f32; 8]>::try_from(&input[40..48]).unwrap()),
            f32x8::from(<[f32; 8]>::try_from(&input[48..56]).unwrap()),
            f32x8::from(<[f32; 8]>::try_from(&input[56..64]).unwrap()),
        ];
        let transposed = f32x8::transpose(rows);
        output[0..8].copy_from_slice(&transposed[0].to_array());
        output[8..16].copy_from_slice(&transposed[1].to_array());
        output[16..24].copy_from_slice(&transposed[2].to_array());
        output[24..32].copy_from_slice(&transposed[3].to_array());
        output[32..40].copy_from_slice(&transposed[4].to_array());
        output[40..48].copy_from_slice(&transposed[5].to_array());
        output[48..56].copy_from_slice(&transposed[6].to_array());
        output[56..64].copy_from_slice(&transposed[7].to_array());
    }

    /// Transpose 8 f32x8 vectors in-place using wide's built-in transpose.
    #[inline(always)]
    fn transpose_vec(rows: [f32x8; 8]) -> [f32x8; 8] {
        f32x8::transpose(rows)
    }

    /// 1/8 scaling factor for DCT normalization.
    /// Applied after each 1D DCT pass, giving 1/64 total scaling (matching C++ jpegli).
    const DCT_SCALE: f32x8 = f32x8::new([1.0 / 8.0; 8]);

    /// Apply 1/8 scaling to 8 f32x8 vectors.
    #[inline(always)]
    fn scale_vec(v: [f32x8; 8]) -> [f32x8; 8] {
        [
            v[0] * DCT_SCALE,
            v[1] * DCT_SCALE,
            v[2] * DCT_SCALE,
            v[3] * DCT_SCALE,
            v[4] * DCT_SCALE,
            v[5] * DCT_SCALE,
            v[6] * DCT_SCALE,
            v[7] * DCT_SCALE,
        ]
    }

    /// Vectorized 1D DCT on 8 rows in parallel.
    /// Input: 8 f32x8 vectors where each represents corresponding values from 8 rows.
    /// (i.e., m[0] contains element 0 from each of 8 rows)
    /// Returns: transformed values in the same layout.
    #[inline(always)]
    fn dct_1d_vec(m: [f32x8; 8]) -> [f32x8; 8] {
        // AddReverse<4>: tmp[0:4] = m[0:4] + reverse(m[4:8])
        let t0 = m[0] + m[7];
        let t1 = m[1] + m[6];
        let t2 = m[2] + m[5];
        let t3 = m[3] + m[4];

        // SubReverse<4>: tmp[4:8] = m[0:4] - reverse(m[4:8])
        let t4 = m[0] - m[7];
        let t5 = m[1] - m[6];
        let t6 = m[2] - m[5];
        let t7 = m[3] - m[4];

        // DCT1D<4> on first half (t0-t3)
        let first = dct_1d_4_vec([t0, t1, t2, t3]);

        // Multiply by WC8
        let t4_scaled = t4 * WC8_0;
        let t5_scaled = t5 * WC8_1;
        let t6_scaled = t6 * WC8_2;
        let t7_scaled = t7 * WC8_3;

        // DCT1D<4> on second half
        let mut second = dct_1d_4_vec([t4_scaled, t5_scaled, t6_scaled, t7_scaled]);

        // B<4>: coeff[0] = coeff[0] * sqrt2 + coeff[1]; then cumulative sum
        second[0] = second[0] * SQRT2_VEC + second[1];
        second[1] += second[2];
        second[2] += second[3];
        // second[3] stays the same

        // InverseEvenOdd<8>: interleave
        [
            first[0], second[0], first[1], second[1], first[2], second[2], first[3], second[3],
        ]
    }

    /// DCT for N=4 on SIMD vectors (helper for dct_1d_vec)
    #[inline(always)]
    fn dct_1d_4_vec(m: [f32x8; 4]) -> [f32x8; 4] {
        // AddReverse<2>: tmp[0:2] = m[0:2] + reverse(m[2:4])
        let t0 = m[0] + m[3];
        let t1 = m[1] + m[2];

        // SubReverse<2>: tmp[2:4] = m[0:2] - reverse(m[2:4])
        let t2 = m[0] - m[3];
        let t3 = m[1] - m[2];

        // DCT1D<2> on first half (t0, t1)
        let r0 = t0 + t1;
        let r1 = t0 - t1;

        // Multiply by WC4
        let t2_scaled = t2 * WC4_0;
        let t3_scaled = t3 * WC4_1;

        // DCT1D<2> on second half
        let r2 = t2_scaled + t3_scaled;
        let r3 = t2_scaled - t3_scaled;

        // B<2>: r2 = r2 * sqrt2 + r3
        let r2_final = r2 * SQRT2_VEC + r3;

        // InverseEvenOdd<4>: interleave even/odd
        [r0, r2_final, r1, r3]
    }

    /// Full SIMD-optimized 2D forward DCT with f32x8 chaining.
    /// Keeps data in SIMD registers throughout the pipeline.
    /// Uses only 2 transposes instead of multiple gather/scatter cycles.
    #[multiversed]
    pub fn forward_dct_8x8_simd_chained(input: &[f32; 64]) -> [f32; 64] {
        // Load input as rows (each f32x8 is one row of the 8x8 block)
        let rows = [
            f32x8::from(<[f32; 8]>::try_from(&input[0..8]).unwrap()),
            f32x8::from(<[f32; 8]>::try_from(&input[8..16]).unwrap()),
            f32x8::from(<[f32; 8]>::try_from(&input[16..24]).unwrap()),
            f32x8::from(<[f32; 8]>::try_from(&input[24..32]).unwrap()),
            f32x8::from(<[f32; 8]>::try_from(&input[32..40]).unwrap()),
            f32x8::from(<[f32; 8]>::try_from(&input[40..48]).unwrap()),
            f32x8::from(<[f32; 8]>::try_from(&input[48..56]).unwrap()),
            f32x8::from(<[f32; 8]>::try_from(&input[56..64]).unwrap()),
        ];

        // Transpose: rows[j] -> cols[i] where cols[i] = [row0[i], row1[i], ..., row7[i]]
        // This layout allows dct_1d_vec to process 8 rows in parallel
        let cols = transpose_vec(rows);

        // Row DCT with 1/8 scaling (first pass)
        let cols_after_row = scale_vec(dct_1d_vec(cols));

        // Transpose back to row layout (also correct layout for column DCT)
        let rows_for_col = transpose_vec(cols_after_row);

        // Column DCT with 1/8 scaling (second pass)
        // Total scaling: 1/8 * 1/8 = 1/64, matching C++ jpegli
        let final_rows = scale_vec(dct_1d_vec(rows_for_col));

        // Store results (scaling already applied)
        let mut output = [0.0f32; 64];

        output[0..8].copy_from_slice(&final_rows[0].to_array());
        output[8..16].copy_from_slice(&final_rows[1].to_array());
        output[16..24].copy_from_slice(&final_rows[2].to_array());
        output[24..32].copy_from_slice(&final_rows[3].to_array());
        output[32..40].copy_from_slice(&final_rows[4].to_array());
        output[40..48].copy_from_slice(&final_rows[5].to_array());
        output[48..56].copy_from_slice(&final_rows[6].to_array());
        output[56..64].copy_from_slice(&final_rows[7].to_array());

        output
    }

    /// Wide-native 2D forward DCT: takes Block8x8f, returns Block8x8f.
    ///
    /// This eliminates all conversion overhead when data is already in wide format.
    /// Use this in hot paths where blocks are stored as Block8x8f.
    #[multiversed]
    #[inline]
    pub fn forward_dct_8x8_wide(input: &Block8x8f) -> Block8x8f {
        // Transpose: rows[j] -> cols[i] where cols[i] = [row0[i], row1[i], ..., row7[i]]
        let cols = transpose_vec(input.rows);

        // Row DCT with 1/8 scaling (first pass)
        let cols_after_row = scale_vec(dct_1d_vec(cols));

        // Transpose back to row layout
        let rows_for_col = transpose_vec(cols_after_row);

        // Column DCT with 1/8 scaling (second pass)
        // Total scaling: 1/8 * 1/8 = 1/64, matching C++ jpegli
        let final_rows = scale_vec(dct_1d_vec(rows_for_col));

        Block8x8f { rows: final_rows }
    }

    // ========================================================================
    // Raw AVX2+FMA DCT Implementation
    // No wide crate, pure intrinsics for maximum performance
    // ========================================================================

    /// In-place 8x8 transpose on 8 __m256 registers.
    /// After transpose, r[i] contains column i from all 8 original rows.
    #[cfg(all(feature = "unsafe_simd", target_arch = "x86_64"))]
    #[target_feature(enable = "avx2")]
    #[inline]
    unsafe fn transpose_8x8_avx_inplace(r: &mut [__m256; 8]) {
        // Phase 1: InterleaveLower/Upper pairs
        let q0 = _mm256_unpacklo_ps(r[0], r[2]);
        let q1 = _mm256_unpacklo_ps(r[1], r[3]);
        let q2 = _mm256_unpackhi_ps(r[0], r[2]);
        let q3 = _mm256_unpackhi_ps(r[1], r[3]);
        let q4 = _mm256_unpacklo_ps(r[4], r[6]);
        let q5 = _mm256_unpacklo_ps(r[5], r[7]);
        let q6 = _mm256_unpackhi_ps(r[4], r[6]);
        let q7 = _mm256_unpackhi_ps(r[5], r[7]);

        // Phase 2: Another round
        let s0 = _mm256_unpacklo_ps(q0, q1);
        let s1 = _mm256_unpackhi_ps(q0, q1);
        let s2 = _mm256_unpacklo_ps(q2, q3);
        let s3 = _mm256_unpackhi_ps(q2, q3);
        let s4 = _mm256_unpacklo_ps(q4, q5);
        let s5 = _mm256_unpackhi_ps(q4, q5);
        let s6 = _mm256_unpacklo_ps(q6, q7);
        let s7 = _mm256_unpackhi_ps(q6, q7);

        // Phase 3: ConcatLowerLower/UpperUpper
        r[0] = _mm256_permute2f128_ps(s0, s4, 0x20);
        r[1] = _mm256_permute2f128_ps(s1, s5, 0x20);
        r[2] = _mm256_permute2f128_ps(s2, s6, 0x20);
        r[3] = _mm256_permute2f128_ps(s3, s7, 0x20);
        r[4] = _mm256_permute2f128_ps(s0, s4, 0x31);
        r[5] = _mm256_permute2f128_ps(s1, s5, 0x31);
        r[6] = _mm256_permute2f128_ps(s2, s6, 0x31);
        r[7] = _mm256_permute2f128_ps(s3, s7, 0x31);
    }

    /// DCT base case for N=2 on raw AVX2 vectors.
    /// Computes: out0 = in0 + in1, out1 = in0 - in1
    #[cfg(all(feature = "unsafe_simd", target_arch = "x86_64"))]
    #[target_feature(enable = "avx2")]
    #[inline]
    unsafe fn dct1d_2_fma(m0: &mut __m256, m1: &mut __m256) {
        let in0 = *m0;
        let in1 = *m1;
        *m0 = _mm256_add_ps(in0, in1);
        *m1 = _mm256_sub_ps(in0, in1);
    }

    /// DCT for N=4 on raw AVX2 vectors with FMA.
    #[cfg(all(feature = "unsafe_simd", target_arch = "x86_64"))]
    #[target_feature(enable = "avx2", enable = "fma")]
    #[inline]
    unsafe fn dct1d_4_fma(m: &mut [__m256; 4]) {
        let wc4_0 = _mm256_set1_ps(0.541196100146197);
        let wc4_1 = _mm256_set1_ps(1.3065629648763764);
        let sqrt2 = _mm256_set1_ps(1.41421356237);

        // AddReverse<2>: tmp[0:2] = m[0:2] + reverse(m[2:4])
        let t0 = _mm256_add_ps(m[0], m[3]);
        let t1 = _mm256_add_ps(m[1], m[2]);

        // SubReverse<2>
        let t2 = _mm256_sub_ps(m[0], m[3]);
        let t3 = _mm256_sub_ps(m[1], m[2]);

        // DCT1D<2> on first half
        let r0 = _mm256_add_ps(t0, t1);
        let r1 = _mm256_sub_ps(t0, t1);

        // Multiply by WC4
        let t2_scaled = _mm256_mul_ps(t2, wc4_0);
        let t3_scaled = _mm256_mul_ps(t3, wc4_1);

        // DCT1D<2> on second half
        let r2 = _mm256_add_ps(t2_scaled, t3_scaled);
        let r3 = _mm256_sub_ps(t2_scaled, t3_scaled);

        // B<2>: r2 = r2 * sqrt2 + r3 (use FMA)
        let r2_final = _mm256_fmadd_ps(r2, sqrt2, r3);

        // InverseEvenOdd<4>: interleave
        m[0] = r0;
        m[1] = r2_final;
        m[2] = r1;
        m[3] = r3;
    }

    /// DCT for N=8 on raw AVX2 vectors with FMA.
    /// Processes 8 independent 8-point DCTs in parallel.
    #[cfg(all(feature = "unsafe_simd", target_arch = "x86_64"))]
    #[target_feature(enable = "avx2", enable = "fma")]
    #[inline]
    unsafe fn dct1d_8_fma(m: &mut [__m256; 8]) {
        let wc8_0 = _mm256_set1_ps(0.5097955791041592);
        let wc8_1 = _mm256_set1_ps(0.6013448869350453);
        let wc8_2 = _mm256_set1_ps(0.8999762231364156);
        let wc8_3 = _mm256_set1_ps(2.5629154477415055);
        let sqrt2 = _mm256_set1_ps(1.41421356237);

        // AddReverse<4>: tmp[0:4] = m[0:4] + reverse(m[4:8])
        let t0 = _mm256_add_ps(m[0], m[7]);
        let t1 = _mm256_add_ps(m[1], m[6]);
        let t2 = _mm256_add_ps(m[2], m[5]);
        let t3 = _mm256_add_ps(m[3], m[4]);

        // SubReverse<4>
        let t4 = _mm256_sub_ps(m[0], m[7]);
        let t5 = _mm256_sub_ps(m[1], m[6]);
        let t6 = _mm256_sub_ps(m[2], m[5]);
        let t7 = _mm256_sub_ps(m[3], m[4]);

        // DCT1D<4> on first half
        let mut first = [t0, t1, t2, t3];
        dct1d_4_fma(&mut first);

        // Multiply by WC8
        let t4_scaled = _mm256_mul_ps(t4, wc8_0);
        let t5_scaled = _mm256_mul_ps(t5, wc8_1);
        let t6_scaled = _mm256_mul_ps(t6, wc8_2);
        let t7_scaled = _mm256_mul_ps(t7, wc8_3);

        // DCT1D<4> on second half
        let mut second = [t4_scaled, t5_scaled, t6_scaled, t7_scaled];
        dct1d_4_fma(&mut second);

        // B<4>: cumulative sum with FMA
        // second[0] = second[0] * sqrt2 + second[1]
        second[0] = _mm256_fmadd_ps(second[0], sqrt2, second[1]);
        // second[1] += second[2]
        second[1] = _mm256_add_ps(second[1], second[2]);
        // second[2] += second[3]
        second[2] = _mm256_add_ps(second[2], second[3]);
        // second[3] stays the same

        // InverseEvenOdd<8>: interleave
        m[0] = first[0];
        m[1] = second[0];
        m[2] = first[1];
        m[3] = second[1];
        m[4] = first[2];
        m[5] = second[2];
        m[6] = first[3];
        m[7] = second[3];
    }

    /// Full 8x8 forward DCT using raw AVX2+FMA intrinsics.
    /// This is the fastest implementation, avoiding all wide crate overhead.
    ///
    /// Algorithm:
    /// 1. Load rows into registers: reg[i] = row i
    /// 2. Transpose: reg[i] = column i (position i of all rows)
    /// 3. Row DCT: dct1d_8_fma processes (reg[0]..reg[7]) as 8 positions
    ///    Each lane processes a separate row. Result: reg[i] = coef i of all rows
    /// 4. Transpose: reg[i] = row i of coefficient matrix
    /// 5. Column DCT: dct1d_8_fma processes columns (reg[k][j] = coef[k,j])
    ///    Result: reg[i] = final row i
    /// 6. Scale and store
    ///
    /// This is a low-level function. Production code should use the safe `forward_dct_8x8` wrapper.
    ///
    /// # Safety
    /// Caller must ensure AVX2+FMA are available.
    #[cfg(all(feature = "unsafe_simd", target_arch = "x86_64"))]
    #[target_feature(enable = "avx2", enable = "fma")]
    #[inline] // Critical: inline into multiversioned caller
    pub(crate) unsafe fn forward_dct_8x8_fma(input: &[f32; 64], output: &mut [f32; 64]) {
        let scale = _mm256_set1_ps(1.0 / 8.0);

        // Load 8 rows: reg[i] = row i of input
        let mut reg = [
            _mm256_loadu_ps(input.as_ptr()),
            _mm256_loadu_ps(input.as_ptr().add(8)),
            _mm256_loadu_ps(input.as_ptr().add(16)),
            _mm256_loadu_ps(input.as_ptr().add(24)),
            _mm256_loadu_ps(input.as_ptr().add(32)),
            _mm256_loadu_ps(input.as_ptr().add(40)),
            _mm256_loadu_ps(input.as_ptr().add(48)),
            _mm256_loadu_ps(input.as_ptr().add(56)),
        ];

        // Transpose: reg[i] = column i = [row0[i], row1[i], ..., row7[i]]
        // dct1d_8_fma treats reg[0..7] as positions 0-7 of the sequence
        // Each lane (element) within a register is a separate row being processed
        transpose_8x8_avx_inplace(&mut reg);

        // Row DCT: all 8 rows processed in parallel
        // Lane k gets: DCT of (reg[0][k], reg[1][k], ..., reg[7][k]) = DCT of row k
        // Result: reg[i][k] = coefficient i of row k's DCT
        dct1d_8_fma(&mut reg);

        // Apply 1/8 scaling after row DCT (first pass)
        for r in &mut reg {
            *r = _mm256_mul_ps(*r, scale);
        }

        // Transpose: reg[i][j] = coef[i, j] (row-major coefficient matrix)
        // This sets up for column DCT: lane j processes column j
        transpose_8x8_avx_inplace(&mut reg);

        // Column DCT: all 8 columns processed in parallel
        // Lane j gets: DCT of (reg[0][j], reg[1][j], ..., reg[7][j]) = DCT of column j
        // Result: reg[i][j] = final 2D DCT coefficient at (i, j)
        dct1d_8_fma(&mut reg);

        // Apply 1/8 scaling after column DCT (second pass)
        // Total scaling: 1/8 * 1/8 = 1/64, matching C++ jpegli
        for r in &mut reg {
            *r = _mm256_mul_ps(*r, scale);
        }

        // reg[i] is now row i of the final output - store directly
        _mm256_storeu_ps(output.as_mut_ptr(), reg[0]);
        _mm256_storeu_ps(output.as_mut_ptr().add(8), reg[1]);
        _mm256_storeu_ps(output.as_mut_ptr().add(16), reg[2]);
        _mm256_storeu_ps(output.as_mut_ptr().add(24), reg[3]);
        _mm256_storeu_ps(output.as_mut_ptr().add(32), reg[4]);
        _mm256_storeu_ps(output.as_mut_ptr().add(40), reg[5]);
        _mm256_storeu_ps(output.as_mut_ptr().add(48), reg[6]);
        _mm256_storeu_ps(output.as_mut_ptr().add(56), reg[7]);
    }
}

/// 1D DCT on all 8 rows (scalar, but faster than SIMD due to cache locality)
/// Kept as reference; forward_dct_8x8 uses SIMD chained approach instead.
#[allow(dead_code)]
#[inline]
fn dct_rows(input: &[f32; 64], output: &mut [f32; 64]) {
    // Process rows sequentially - good cache access pattern
    // The parallel SIMD approach has overhead from vector construction
    for row in 0..8 {
        let mut tmp = [0.0f32; 8];
        for i in 0..8 {
            tmp[i] = input[row * 8 + i];
        }
        dct1d_8(&mut tmp);
        for i in 0..8 {
            output[row * 8 + i] = tmp[i];
        }
    }
}

/// Performs an 8x8 forward DCT on the input block.
///
/// The input should be pixel values (typically shifted by -128 for level shift).
/// The output is DCT coefficients ready for quantization.
/// Includes 1/8 scaling for JPEG decoder compatibility.
///
/// Uses raw AVX2+FMA intrinsics when available for maximum performance.
///
/// # Arguments
/// * `input` - 8x8 block of input values in row-major order
///
/// # Returns
/// 8x8 block of DCT coefficients
/// Performs an 8x8 forward DCT with automatic CPU dispatch.
///
/// Uses multiversion for one-time dispatch at load (not per-call).
/// Inside each version, `cfg(target_feature)` provides zero-cost branching.
#[multiversed]
#[must_use]
pub fn forward_dct_8x8(input: &[f32; DCT_BLOCK_SIZE]) -> [f32; DCT_BLOCK_SIZE] {
    // Use raw AVX2 intrinsics when unsafe_simd feature is enabled
    #[cfg(all(
        feature = "unsafe_simd",
        target_arch = "x86_64",
        target_feature = "avx2",
        target_feature = "fma"
    ))]
    {
        let mut output = [0.0f32; 64];
        // SAFETY: This code path only compiles when target_feature includes avx2+fma
        unsafe {
            simd::forward_dct_8x8_fma(input, &mut output);
        }
        return output;
    }

    // Safe fallback using wide crate SIMD (portable to all platforms)
    #[cfg(not(all(
        feature = "unsafe_simd",
        target_arch = "x86_64",
        target_feature = "avx2",
        target_feature = "fma"
    )))]
    {
        forward_dct_8x8_scalar(input)
    }
}

/// Fallback for forward DCT using wide crate SIMD (portable to all platforms).
/// Public so multiversion blocks can call it directly via cfg(target_feature).
#[inline]
pub fn forward_dct_8x8_scalar(input: &[f32; DCT_BLOCK_SIZE]) -> [f32; DCT_BLOCK_SIZE] {
    simd::forward_dct_8x8_simd_chained(input)
}

/// Performs forward DCT on a block with level shift.
///
/// # Arguments
/// * `input` - 8x8 block of 8-bit pixel values
///
/// # Returns
/// 8x8 block of DCT coefficients
#[must_use]
pub fn forward_dct_8x8_u8(input: &[u8; DCT_BLOCK_SIZE]) -> [f32; DCT_BLOCK_SIZE] {
    // Convert to f32 with level shift (-128)
    let mut shifted = [0.0f32; DCT_BLOCK_SIZE];
    let level_shift = f32x8::splat(128.0);

    // Process 8 values at a time
    for chunk in 0..8 {
        let k = chunk * 8;
        let v = f32x8::from([
            input[k] as f32,
            input[k + 1] as f32,
            input[k + 2] as f32,
            input[k + 3] as f32,
            input[k + 4] as f32,
            input[k + 5] as f32,
            input[k + 6] as f32,
            input[k + 7] as f32,
        ]);
        let result = v - level_shift;
        let arr: [f32; 8] = result.into();
        shifted[k..k + 8].copy_from_slice(&arr);
    }

    forward_dct_8x8(&shifted)
}

/// Performs forward DCT on multiple blocks.
pub fn forward_dct_blocks(blocks: &[[f32; DCT_BLOCK_SIZE]]) -> Vec<[f32; DCT_BLOCK_SIZE]> {
    blocks.iter().map(forward_dct_8x8).collect()
}

// ============================================================================
// AAN DCT Implementation
// ============================================================================

/// AAN 1D DCT: 5 multiplies + 29 adds (+ 8 descaling multiplies)
///
/// This is the core of the Arai-Agui-Nakajima algorithm.
/// Processes one row/column of 8 elements in-place.
/// Includes descaling to match standard DCT normalization.
#[inline(always)]
fn aan_dct_1d(d: &mut [f32; 8]) {
    // Stage 1: Butterfly - split into even/odd symmetric pairs
    let tmp0 = d[0] + d[7];
    let tmp7 = d[0] - d[7];
    let tmp1 = d[1] + d[6];
    let tmp6 = d[1] - d[6];
    let tmp2 = d[2] + d[5];
    let tmp5 = d[2] - d[5];
    let tmp3 = d[3] + d[4];
    let tmp4 = d[3] - d[4];

    // Stage 2: Even part (produces coefficients at positions 0, 2, 4, 6)
    let tmp10 = tmp0 + tmp3;
    let tmp13 = tmp0 - tmp3;
    let tmp11 = tmp1 + tmp2;
    let tmp12 = tmp1 - tmp2;

    d[0] = tmp10 + tmp11; // DC coefficient
    d[4] = tmp10 - tmp11;

    // MULTIPLY #1: z1 = (tmp12 + tmp13) * cos(π/4)
    let z1 = (tmp12 + tmp13) * aan::C4;
    d[2] = tmp13 + z1;
    d[6] = tmp13 - z1;

    // Stage 3: Odd part (produces coefficients at positions 1, 3, 5, 7)
    let tmp10 = tmp4 + tmp5;
    let tmp11 = tmp5 + tmp6;
    let tmp12 = tmp6 + tmp7;

    // MULTIPLY #2: z5 = (tmp10 - tmp12) * sin(π/8)
    let z5 = (tmp10 - tmp12) * aan::C6;
    // MULTIPLY #3: z2 = (cos(π/8) - cos(3π/8)) * tmp10 + z5
    let z2 = aan::C2_M_C6 * tmp10 + z5;
    // MULTIPLY #4: z4 = (cos(π/8) + cos(3π/8)) * tmp12 + z5
    let z4 = aan::C2_P_C6 * tmp12 + z5;
    // MULTIPLY #5: z3 = tmp11 * cos(π/4)
    let z3 = tmp11 * aan::C4;

    let z11 = tmp7 + z3;
    let z13 = tmp7 - z3;

    d[5] = z13 + z2;
    d[3] = z13 - z2;
    d[1] = z11 + z4;
    d[7] = z11 - z4;

    // Descale to match standard DCT normalization
    // (multiply by reciprocal of aanscales)
    d[0] *= aan::INV_SCALES[0];
    d[1] *= aan::INV_SCALES[1];
    d[2] *= aan::INV_SCALES[2];
    d[3] *= aan::INV_SCALES[3];
    d[4] *= aan::INV_SCALES[4];
    d[5] *= aan::INV_SCALES[5];
    d[6] *= aan::INV_SCALES[6];
    d[7] *= aan::INV_SCALES[7];
}

/// AAN forward DCT for 8x8 block.
///
/// Uses the Arai-Agui-Nakajima algorithm which requires only 5 multiplies
/// and 29 adds per 1D transform (vs ~11 multiplies in naive approach).
///
/// Based on libjpeg's jfdctflt.c implementation.
#[inline]
#[must_use]
pub fn aan_forward_dct_8x8(input: &[f32; DCT_BLOCK_SIZE]) -> [f32; DCT_BLOCK_SIZE] {
    let mut data = *input;

    // Pass 1: Process rows
    for row in 0..8 {
        let offset = row * 8;
        let mut row_data: [f32; 8] = [
            data[offset],
            data[offset + 1],
            data[offset + 2],
            data[offset + 3],
            data[offset + 4],
            data[offset + 5],
            data[offset + 6],
            data[offset + 7],
        ];
        aan_dct_1d(&mut row_data);
        data[offset] = row_data[0];
        data[offset + 1] = row_data[1];
        data[offset + 2] = row_data[2];
        data[offset + 3] = row_data[3];
        data[offset + 4] = row_data[4];
        data[offset + 5] = row_data[5];
        data[offset + 6] = row_data[6];
        data[offset + 7] = row_data[7];
    }

    // Apply 1/8 scaling after row pass (matching C++ jpegli)
    for v in &mut data {
        *v *= 0.125;
    }

    // Pass 2: Process columns (strided access)
    for col in 0..8 {
        let mut col_data: [f32; 8] = [
            data[col],
            data[col + 8],
            data[col + 16],
            data[col + 24],
            data[col + 32],
            data[col + 40],
            data[col + 48],
            data[col + 56],
        ];
        aan_dct_1d(&mut col_data);
        data[col] = col_data[0];
        data[col + 8] = col_data[1];
        data[col + 16] = col_data[2];
        data[col + 24] = col_data[3];
        data[col + 32] = col_data[4];
        data[col + 40] = col_data[5];
        data[col + 48] = col_data[6];
        data[col + 56] = col_data[7];
    }

    // Apply 1/8 scaling after column pass (total 1/64, matching C++ jpegli)
    for v in &mut data {
        *v *= 0.125;
    }

    data
}

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

    #[test]
    fn test_dct_dc_only() {
        // Constant block should have only DC coefficient
        let input = [128.0f32; DCT_BLOCK_SIZE];
        let output = forward_dct_8x8(&input);

        // DC coefficient should be non-zero
        assert!(output[0].abs() > 1.0);

        // AC coefficients should be near zero
        for i in 1..DCT_BLOCK_SIZE {
            assert!(
                output[i].abs() < 0.001,
                "AC[{}] = {} should be ~0",
                i,
                output[i]
            );
        }
    }

    #[test]
    fn test_dct_zero_block() {
        let input = [0.0f32; DCT_BLOCK_SIZE];
        let output = forward_dct_8x8(&input);

        for (i, &v) in output.iter().enumerate() {
            assert!(v.abs() < 0.001, "coeff[{}] = {} should be 0", i, v);
        }
    }

    #[test]
    fn test_dct_u8_level_shift() {
        // All 128s should produce near-zero output (after -128 shift)
        let input = [128u8; DCT_BLOCK_SIZE];
        let output = forward_dct_8x8_u8(&input);

        for (i, &v) in output.iter().enumerate() {
            assert!(v.abs() < 0.001, "coeff[{}] = {} should be 0", i, v);
        }
    }

    #[test]
    fn test_transpose() {
        let mut input = [0.0f32; 64];
        for i in 0..64 {
            input[i] = i as f32;
        }

        let mut output = [0.0f32; 64];
        transpose_8x8(&input, &mut output);

        // Check a few values
        assert_eq!(output[0], 0.0); // (0,0)
        assert_eq!(output[1], 8.0); // (0,1) <- (1,0) in original
        assert_eq!(output[8], 1.0); // (1,0) <- (0,1) in original
        assert_eq!(output[9], 9.0); // (1,1)
    }

    #[test]
    fn test_simd_transpose_matches_scalar() {
        let mut input = [0.0f32; 64];
        for i in 0..64 {
            input[i] = (i as f32 * 1.7).sin() * 100.0;
        }

        let mut scalar_output = [0.0f32; 64];
        let mut simd_output = [0.0f32; 64];

        transpose_8x8(&input, &mut scalar_output);
        simd::transpose_8x8_simd(&input, &mut simd_output);

        for i in 0..64 {
            assert!(
                (scalar_output[i] - simd_output[i]).abs() < 1e-6,
                "Mismatch at {}: scalar={} simd={}",
                i,
                scalar_output[i],
                simd_output[i]
            );
        }
    }

    #[test]
    fn test_simd_dct_rows_matches_scalar() {
        // Test with various input patterns
        let patterns: Vec<[f32; 64]> = vec![
            // Constant block
            [64.0; 64],
            // Gradient
            {
                let mut arr = [0.0f32; 64];
                for i in 0..64 {
                    arr[i] = i as f32;
                }
                arr
            },
            // Sinusoidal pattern
            {
                let mut arr = [0.0f32; 64];
                for i in 0..64 {
                    arr[i] = (i as f32 * 0.3).sin() * 100.0;
                }
                arr
            },
            // Checkerboard
            {
                let mut arr = [0.0f32; 64];
                for row in 0..8 {
                    for col in 0..8 {
                        arr[row * 8 + col] = if (row + col) % 2 == 0 { 100.0 } else { -100.0 };
                    }
                }
                arr
            },
        ];

        for (pattern_idx, input) in patterns.iter().enumerate() {
            // Scalar implementation
            let mut scalar_output = [0.0f32; 64];
            for row in 0..8 {
                let mut tmp = [0.0f32; 8];
                for i in 0..8 {
                    tmp[i] = input[row * 8 + i];
                }
                dct1d_8(&mut tmp);
                for i in 0..8 {
                    scalar_output[row * 8 + i] = tmp[i];
                }
            }

            // SIMD implementation
            let mut simd_output = [0.0f32; 64];
            simd::dct_8rows_parallel(input, &mut simd_output);

            // Compare
            let mut max_error = 0.0f32;
            for i in 0..64 {
                let error = (scalar_output[i] - simd_output[i]).abs();
                max_error = max_error.max(error);
            }

            assert!(
                max_error < 1e-4,
                "Pattern {}: SIMD DCT rows max error {} exceeds threshold",
                pattern_idx,
                max_error
            );
        }
    }

    #[test]
    fn test_aan_dct_matches_recursive() {
        // Test AAN DCT against recursive implementation with various patterns
        let patterns: Vec<[f32; 64]> = vec![
            // Constant block
            [64.0; 64],
            // Gradient
            {
                let mut arr = [0.0f32; 64];
                for i in 0..64 {
                    arr[i] = i as f32;
                }
                arr
            },
            // Sinusoidal pattern
            {
                let mut arr = [0.0f32; 64];
                for i in 0..64 {
                    arr[i] = (i as f32 * 0.3).sin() * 100.0;
                }
                arr
            },
            // Checkerboard
            {
                let mut arr = [0.0f32; 64];
                for row in 0..8 {
                    for col in 0..8 {
                        arr[row * 8 + col] = if (row + col) % 2 == 0 { 100.0 } else { -100.0 };
                    }
                }
                arr
            },
            // Random-ish pattern
            {
                let mut arr = [0.0f32; 64];
                for i in 0..64 {
                    arr[i] = ((i * 17 + 31) % 256) as f32 - 128.0;
                }
                arr
            },
        ];

        for (pattern_idx, input) in patterns.iter().enumerate() {
            let recursive_output = forward_dct_8x8(input);
            let aan_output = aan_forward_dct_8x8(input);

            let mut max_error = 0.0f32;
            for i in 0..64 {
                let error = (recursive_output[i] - aan_output[i]).abs();
                max_error = max_error.max(error);
            }

            // AAN and recursive should produce nearly identical results
            // Allow small floating-point differences
            assert!(
                max_error < 1e-4,
                "Pattern {}: AAN vs recursive max error {} exceeds threshold.\n\
                 Recursive[0..8]: {:?}\n\
                 AAN[0..8]: {:?}",
                pattern_idx,
                max_error,
                &recursive_output[0..8],
                &aan_output[0..8]
            );
        }
    }

    #[cfg(all(feature = "unsafe_simd", target_arch = "x86_64"))]
    #[test]
    fn test_fma_dct_matches_scalar() {
        if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
            return; // Skip on CPUs without AVX2+FMA
        }

        let patterns: Vec<[f32; 64]> = vec![
            [64.0; 64],
            {
                let mut arr = [0.0f32; 64];
                for i in 0..64 {
                    arr[i] = i as f32;
                }
                arr
            },
            {
                let mut arr = [0.0f32; 64];
                for i in 0..64 {
                    arr[i] = ((i as f32) * 0.5).sin() * 100.0 + 50.0;
                }
                arr
            },
        ];

        for (pattern_idx, input) in patterns.iter().enumerate() {
            let scalar_output = forward_dct_8x8_scalar(input);

            let mut avx2_output = [0.0f32; 64];
            unsafe {
                simd::forward_dct_8x8_fma(input, &mut avx2_output);
            }

            let mut max_error = 0.0f32;
            for i in 0..64 {
                let error = (scalar_output[i] - avx2_output[i]).abs();
                max_error = max_error.max(error);
            }

            assert!(
                max_error < 1e-4,
                "Pattern {}: AVX2 vs scalar max error {} exceeds threshold.\n\
                 Scalar[0..8]: {:?}\n\
                 AVX2[0..8]: {:?}",
                pattern_idx,
                max_error,
                &scalar_output[0..8],
                &avx2_output[0..8]
            );
        }
    }

    #[cfg(all(feature = "unsafe_simd", target_arch = "x86_64"))]
    #[test]
    #[ignore] // Run with: cargo test --release --features unsafe_simd bench_fma_dct -- --ignored --nocapture
    fn bench_fma_dct_vs_scalar() {
        use std::hint::black_box;
        use std::time::Instant;

        if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
            println!("AVX2+FMA not available, skipping benchmark");
            return;
        }

        let input: [f32; 64] = [
            52., 55., 61., 66., 70., 61., 64., 73., 63., 59., 55., 90., 109., 85., 69., 72., 62.,
            59., 68., 113., 144., 104., 66., 73., 63., 58., 71., 122., 154., 106., 70., 69., 67.,
            61., 68., 104., 126., 88., 68., 70., 79., 65., 60., 70., 77., 68., 58., 75., 85., 71.,
            64., 59., 55., 61., 65., 83., 87., 79., 69., 68., 65., 76., 78., 94.,
        ];

        let iterations = 10_000_000;
        let mut output = [0.0f32; 64];

        // Warmup scalar
        for _ in 0..10000 {
            let _ = forward_dct_8x8_scalar(black_box(&input));
        }

        let start = Instant::now();
        for _ in 0..iterations {
            let result = forward_dct_8x8_scalar(black_box(&input));
            black_box(&result);
        }
        let scalar_time = start.elapsed();
        let scalar_ns = scalar_time.as_nanos() as f64 / iterations as f64;

        // Warmup AVX2
        for _ in 0..10000 {
            unsafe {
                simd::forward_dct_8x8_fma(black_box(&input), &mut output);
            }
        }

        let start = Instant::now();
        for _ in 0..iterations {
            unsafe {
                simd::forward_dct_8x8_fma(black_box(&input), black_box(&mut output));
            }
        }
        let avx2_time = start.elapsed();
        let avx2_ns = avx2_time.as_nanos() as f64 / iterations as f64;

        println!("DCT Performance ({} iterations):", iterations);
        println!("  Scalar: {:.2} ns/DCT", scalar_ns);
        println!("  AVX2:   {:.2} ns/DCT", avx2_ns);
        println!(
            "  Speedup: {:.2}x ({})",
            scalar_ns / avx2_ns,
            if avx2_ns < scalar_ns {
                "AVX2 faster"
            } else {
                "Scalar faster"
            }
        );
    }

    #[test]
    fn test_wide_dct_matches_array_dct() {
        use crate::foundation::simd_types::Block8x8f;

        // Test with various patterns
        let patterns: [[f32; 64]; 4] = [
            [128.0; 64], // Constant
            {
                let mut arr = [0.0f32; 64];
                for i in 0..64 {
                    arr[i] = i as f32;
                }
                arr
            },
            {
                let mut arr = [0.0f32; 64];
                for i in 0..64 {
                    arr[i] = (i as f32 * 0.3).sin() * 100.0;
                }
                arr
            },
            {
                let mut arr = [0.0f32; 64];
                for row in 0..8 {
                    for col in 0..8 {
                        arr[row * 8 + col] = if (row + col) % 2 == 0 { 100.0 } else { -100.0 };
                    }
                }
                arr
            },
        ];

        for (idx, input) in patterns.iter().enumerate() {
            // Array-based DCT
            let array_output = forward_dct_8x8(input);

            // Wide-native DCT
            let block_input = Block8x8f::from_array(input);
            let block_output = simd::forward_dct_8x8_wide(&block_input);
            let wide_output = block_output.to_array();

            // Compare
            let mut max_error = 0.0f32;
            for i in 0..64 {
                let error = (array_output[i] - wide_output[i]).abs();
                max_error = max_error.max(error);
            }

            assert!(
                max_error < 1e-4,
                "Pattern {}: wide DCT differs from array DCT by {}",
                idx,
                max_error
            );
        }
    }
}