// SPDX-License-Identifier: MIT OR Apache-2.0

use j2k_transcode::dct97_2d::{
    dct8x8_blocks_then_dwt97_float, dct8x8_blocks_then_dwt97_float_with_scratch, Dct97GridScratch,
    Dwt97TwoDimensional,
};

#[test]
fn dct8x8_grid_to_2d_97_idct_scratch_path_matches_reference_for_structured_cases() {
    let blocks = structured_blocks(2, 2);
    let mut scratch = Dct97GridScratch::default();

    for (width, height) in [(8, 8), (13, 11), (16, 16)] {
        let scratch_path =
            dct8x8_blocks_then_dwt97_float_with_scratch(&blocks, 2, 2, width, height, &mut scratch)
                .expect("scratch 9/7 IDCT path accepts covered grid");
        let reference = dct8x8_blocks_then_dwt97_float(&blocks, 2, 2, width, height)
            .expect("reference 9/7 IDCT path accepts covered grid");

        assert!(
            max_abs_diff(&scratch_path, &reference) < 1.0e-9,
            "scratch 9/7 IDCT path diverged for {width}x{height}"
        );
    }
}

#[test]
fn dct8x8_grid_to_2d_97_idct_scratch_path_matches_reference_and_reuses_storage() {
    let large_blocks = structured_blocks(32, 32);
    let small_blocks = structured_blocks(2, 2);
    let mut scratch = Dct97GridScratch::default();

    let large =
        dct8x8_blocks_then_dwt97_float_with_scratch(&large_blocks, 32, 32, 255, 241, &mut scratch)
            .expect("scratch 9/7 IDCT path accepts covered large grid");
    let expected_large = dct8x8_blocks_then_dwt97_float(&large_blocks, 32, 32, 255, 241)
        .expect("reference 9/7 IDCT path accepts covered large grid");
    let capacity_after_large = scratch.spatial_sample_capacity();

    let small =
        dct8x8_blocks_then_dwt97_float_with_scratch(&small_blocks, 2, 2, 13, 11, &mut scratch)
            .expect("scratch 9/7 IDCT path accepts covered small grid");
    let expected_small = dct8x8_blocks_then_dwt97_float(&small_blocks, 2, 2, 13, 11)
        .expect("reference 9/7 IDCT path accepts covered small grid");

    assert!(
        max_abs_diff(&large, &expected_large) < 1.0e-9,
        "scratch 9/7 IDCT path diverged for large grid"
    );
    assert!(
        max_abs_diff(&small, &expected_small) < 1.0e-9,
        "scratch 9/7 IDCT path diverged for small grid"
    );
    assert_eq!(scratch.spatial_sample_capacity(), capacity_after_large);
}

fn max_abs_diff(actual: &Dwt97TwoDimensional<f64>, expected: &Dwt97TwoDimensional<f64>) -> f64 {
    assert_eq!(actual.low_width, expected.low_width);
    assert_eq!(actual.low_height, expected.low_height);
    assert_eq!(actual.high_width, expected.high_width);
    assert_eq!(actual.high_height, expected.high_height);

    actual
        .ll
        .iter()
        .zip(expected.ll.iter())
        .chain(actual.hl.iter().zip(expected.hl.iter()))
        .chain(actual.lh.iter().zip(expected.lh.iter()))
        .chain(actual.hh.iter().zip(expected.hh.iter()))
        .map(|(actual, expected)| (actual - expected).abs())
        .fold(0.0, f64::max)
}

fn structured_blocks(block_cols: usize, block_rows: usize) -> Vec<[[f64; 8]; 8]> {
    let mut blocks = Vec::with_capacity(block_cols * block_rows);
    for block_y in 0..block_rows {
        for block_x in 0..block_cols {
            let mut block = [[0.0; 8]; 8];
            block[0][0] = 384.0 + (block_x * 19 + block_y * 23) as f64;
            block[0][1] = -17.0 + block_x as f64;
            block[1][0] = 11.0 - block_y as f64;
            block[2][3] = 7.0;
            block[4][4] = -3.0;
            block[7][7] = 2.0;
            blocks.push(block);
        }
    }
    blocks
}
