//! Forward Discrete Cosine Transform (DCT) implementation.
//!
//! This module provides 8x8 DCT-II transform used for JPEG encoding.
//! Uses the recursive splitting algorithm from jpegli for exact compatibility.
//!
//! SIMD optimization is available via the `simd` feature (enabled by default).

use crate::consts::DCT_BLOCK_SIZE;

#[cfg(feature = "simd")]
use wide::f32x8;

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

const WC8: [f32; 4] = [
    0.5097955791041592,
    0.6013448869350453,
    0.8999762231364156,
    2.5629154477415055,
];

const SQRT2: f32 = 1.41421356237;

/// Transpose an 8x8 block (scalar version).
#[cfg(any(not(feature = "simd"), 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]
#[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]
#[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
#[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
#[inline]
fn multiply_wc8(coeff: &mut [f32]) {
    for i in 0..4 {
        coeff[4 + i] *= WC8[i];
    }
}

#[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
#[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
#[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)
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)
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);
}

/// 1D DCT on all 8 rows (no per-row scaling, scaling handled in main function)
#[cfg(not(feature = "simd"))]
fn dct_rows(input: &[f32; 64], output: &mut [f32; 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 {
            output[row * 8 + i] = tmp[i];
        }
    }
}

// SIMD-optimized implementations
#[cfg(feature = "simd")]
mod simd {
    use super::*;

    /// SIMD-optimized 8x8 transpose.
    /// Uses proper 8x8 transpose algorithm with AVX-style shuffle patterns.
    pub fn transpose_8x8_simd(input: &[f32; 64], output: &mut [f32; 64]) {
        // Load all 8 rows
        let r0 = f32x8::from(&input[0..8]);
        let r1 = f32x8::from(&input[8..16]);
        let r2 = f32x8::from(&input[16..24]);
        let r3 = f32x8::from(&input[24..32]);
        let r4 = f32x8::from(&input[32..40]);
        let r5 = f32x8::from(&input[40..48]);
        let r6 = f32x8::from(&input[48..56]);
        let r7 = f32x8::from(&input[56..64]);

        // Convert to arrays for transpose operations
        let a0 = r0.to_array();
        let a1 = r1.to_array();
        let a2 = r2.to_array();
        let a3 = r3.to_array();
        let a4 = r4.to_array();
        let a5 = r5.to_array();
        let a6 = r6.to_array();
        let a7 = r7.to_array();

        // Transpose: output[col * 8 + row] = input[row * 8 + col]
        // Build each output row by gathering from columns
        let t0 = f32x8::from([a0[0], a1[0], a2[0], a3[0], a4[0], a5[0], a6[0], a7[0]]);
        let t1 = f32x8::from([a0[1], a1[1], a2[1], a3[1], a4[1], a5[1], a6[1], a7[1]]);
        let t2 = f32x8::from([a0[2], a1[2], a2[2], a3[2], a4[2], a5[2], a6[2], a7[2]]);
        let t3 = f32x8::from([a0[3], a1[3], a2[3], a3[3], a4[3], a5[3], a6[3], a7[3]]);
        let t4 = f32x8::from([a0[4], a1[4], a2[4], a3[4], a4[4], a5[4], a6[4], a7[4]]);
        let t5 = f32x8::from([a0[5], a1[5], a2[5], a3[5], a4[5], a5[5], a6[5], a7[5]]);
        let t6 = f32x8::from([a0[6], a1[6], a2[6], a3[6], a4[6], a5[6], a6[6], a7[6]]);
        let t7 = f32x8::from([a0[7], a1[7], a2[7], a3[7], a4[7], a5[7], a6[7], a7[7]]);

        // Store results
        output[0..8].copy_from_slice(&t0.to_array());
        output[8..16].copy_from_slice(&t1.to_array());
        output[16..24].copy_from_slice(&t2.to_array());
        output[24..32].copy_from_slice(&t3.to_array());
        output[32..40].copy_from_slice(&t4.to_array());
        output[40..48].copy_from_slice(&t5.to_array());
        output[48..56].copy_from_slice(&t6.to_array());
        output[56..64].copy_from_slice(&t7.to_array());
    }

    /// SIMD-optimized 1D DCT on a row using f32x8
    #[inline]
    pub fn dct1d_8_simd(input: f32x8) -> f32x8 {
        let arr = input.to_array();
        let mut mem = arr;

        // Use scalar algorithm (the structure doesn't vectorize well within a single row)
        super::dct1d_8(&mut mem);

        f32x8::from(mem)
    }
}

/// 1D DCT on all 8 rows (no per-row scaling, scaling handled in main function)
#[cfg(feature = "simd")]
fn dct_rows(input: &[f32; 64], output: &mut [f32; 64]) {
    for row in 0..8 {
        let row_vec = f32x8::from(&input[row * 8..row * 8 + 8]);
        let result = simd::dct1d_8_simd(row_vec);
        let arr = result.to_array();
        for i in 0..8 {
            output[row * 8 + i] = arr[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 SIMD optimization when the `simd` feature is enabled.
///
/// # Arguments
/// * `input` - 8x8 block of input values in row-major order
///
/// # Returns
/// 8x8 block of DCT coefficients
#[must_use]
pub fn forward_dct_8x8(input: &[f32; DCT_BLOCK_SIZE]) -> [f32; DCT_BLOCK_SIZE] {
    let mut scratch = [0.0f32; 64];
    let mut coefficients = [0.0f32; 64];

    // Row transform
    dct_rows(input, &mut scratch);

    // Transpose (use SIMD version when available)
    #[cfg(feature = "simd")]
    simd::transpose_8x8_simd(&scratch, &mut coefficients);
    #[cfg(not(feature = "simd"))]
    transpose_8x8(&scratch, &mut coefficients);

    // Column transform (on transposed data)
    dct_rows(&coefficients, &mut scratch);

    // Transpose back
    #[cfg(feature = "simd")]
    simd::transpose_8x8_simd(&scratch, &mut coefficients);
    #[cfg(not(feature = "simd"))]
    transpose_8x8(&scratch, &mut coefficients);

    // Apply 1/8 scaling for JPEG decoder compatibility
    // Note: C++ jpegli applies 1/8 per dimension (1/64 total), but standard JPEG
    // decoders like libjpeg/jpeg-decoder expect 1/8 total scaling.
    // Using 1/64 here breaks compatibility (decoded values are 8× too small).
    let scale = 1.0 / 8.0;
    for v in &mut coefficients {
        *v *= scale;
    }

    coefficients
}

/// 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];
    for (i, &v) in input.iter().enumerate() {
        shifted[i] = v as f32 - 128.0;
    }
    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()
}

#[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)
    }

    #[cfg(feature = "simd")]
    #[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]
            );
        }
    }
}
