/*
 * // Copyright (c) Radzivon Bartoshyk 6/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.
 */

#[rustfmt::skip]
pub(crate) static ZIGZAG: [(usize, usize); 64] = [
    (0,0),(1,0),(0,1),(2,0),(1,1),(0,2),(3,0),(2,1),
    (1,2),(0,3),(3,1),(2,2),(1,3),(3,2),(2,3),(3,3),
    (4,0),(5,0),(4,1),(6,0),(5,1),(4,2),(7,0),(6,1),
    (5,2),(4,3),(7,1),(6,2),(5,3),(7,2),(6,3),(7,3),
    (0,4),(1,4),(0,5),(2,4),(1,5),(0,6),(3,4),(2,5),
    (1,6),(0,7),(3,5),(2,6),(1,7),(3,6),(2,7),(3,7),
    (4,4),(5,4),(4,5),(6,4),(5,5),(4,6),(7,4),(6,5),
    (5,6),(4,7),(7,5),(6,6),(5,7),(7,6),(6,7),(7,7),
];

/// HEVC up-right diagonal scan for a single 4×4 block, `(row, col)`.
#[rustfmt::skip]
pub(crate) static DIAG_SCAN_4X4: [(usize, usize); 16] = [
    (0,0),(1,0),(0,1),(2,0),(1,1),(0,2),(3,0),(2,1),
    (1,2),(0,3),(3,1),(2,2),(1,3),(3,2),(2,3),(3,3),
];

/// HEVC sub-block-major up-right diagonal scan for a 16×16 TU, `(row, col)`.
/// Layout matches the 8×8 `ZIGZAG`: a 4×4 grid of 4×4 sub-blocks, each sub-block
/// scanned by `DIAG_SCAN_4X4`, the sub-blocks themselves visited in up-right
/// diagonal order. Generated by the same construction that reproduces `ZIGZAG`
/// and `DIAG_SCAN_4X4` exactly, so `coeffs[sb*16 + k]` indexing is consistent
/// across all TU sizes.
#[rustfmt::skip]
pub(crate) static ZIGZAG_16X16: [(usize, usize); 256] = [
    (0,0),(1,0),(0,1),(2,0),(1,1),(0,2),(3,0),(2,1),
    (1,2),(0,3),(3,1),(2,2),(1,3),(3,2),(2,3),(3,3),
    (4,0),(5,0),(4,1),(6,0),(5,1),(4,2),(7,0),(6,1),
    (5,2),(4,3),(7,1),(6,2),(5,3),(7,2),(6,3),(7,3),
    (0,4),(1,4),(0,5),(2,4),(1,5),(0,6),(3,4),(2,5),
    (1,6),(0,7),(3,5),(2,6),(1,7),(3,6),(2,7),(3,7),
    (8,0),(9,0),(8,1),(10,0),(9,1),(8,2),(11,0),(10,1),
    (9,2),(8,3),(11,1),(10,2),(9,3),(11,2),(10,3),(11,3),
    (4,4),(5,4),(4,5),(6,4),(5,5),(4,6),(7,4),(6,5),
    (5,6),(4,7),(7,5),(6,6),(5,7),(7,6),(6,7),(7,7),
    (0,8),(1,8),(0,9),(2,8),(1,9),(0,10),(3,8),(2,9),
    (1,10),(0,11),(3,9),(2,10),(1,11),(3,10),(2,11),(3,11),
    (12,0),(13,0),(12,1),(14,0),(13,1),(12,2),(15,0),(14,1),
    (13,2),(12,3),(15,1),(14,2),(13,3),(15,2),(14,3),(15,3),
    (8,4),(9,4),(8,5),(10,4),(9,5),(8,6),(11,4),(10,5),
    (9,6),(8,7),(11,5),(10,6),(9,7),(11,6),(10,7),(11,7),
    (4,8),(5,8),(4,9),(6,8),(5,9),(4,10),(7,8),(6,9),
    (5,10),(4,11),(7,9),(6,10),(5,11),(7,10),(6,11),(7,11),
    (0,12),(1,12),(0,13),(2,12),(1,13),(0,14),(3,12),(2,13),
    (1,14),(0,15),(3,13),(2,14),(1,15),(3,14),(2,15),(3,15),
    (12,4),(13,4),(12,5),(14,4),(13,5),(12,6),(15,4),(14,5),
    (13,6),(12,7),(15,5),(14,6),(13,7),(15,6),(14,7),(15,7),
    (8,8),(9,8),(8,9),(10,8),(9,9),(8,10),(11,8),(10,9),
    (9,10),(8,11),(11,9),(10,10),(9,11),(11,10),(10,11),(11,11),
    (4,12),(5,12),(4,13),(6,12),(5,13),(4,14),(7,12),(6,13),
    (5,14),(4,15),(7,13),(6,14),(5,15),(7,14),(6,15),(7,15),
    (12,8),(13,8),(12,9),(14,8),(13,9),(12,10),(15,8),(14,9),
    (13,10),(12,11),(15,9),(14,10),(13,11),(15,10),(14,11),(15,11),
    (8,12),(9,12),(8,13),(10,12),(9,13),(8,14),(11,12),(10,13),
    (9,14),(8,15),(11,13),(10,14),(9,15),(11,14),(10,15),(11,15),
    (12,12),(13,12),(12,13),(14,12),(13,13),(12,14),(15,12),(14,13),
    (13,14),(12,15),(15,13),(14,14),(13,15),(15,14),(14,15),(15,15),
];

// ── Mode-dependent coefficient scans (HEVC §6.5.3/6.5.4/6.5.5) ──────────────
// For 4×4 and 8×8-luma intra TBs the scan order depends on the prediction mode:
// modes 6..=14 use the vertical scan, 22..=30 the horizontal scan, all others
// (and every larger TB) the up-right diagonal scan. Encoder and decoder must
// agree, so this selection is conformance-critical.

/// scanIdx (0 = diagonal, 1 = horizontal, 2 = vertical) for an intra TB.
///
/// Mode-dependent for 4×4 (any component), 8×8 luma, and — in 4:4:4 — 8×8 chroma
/// (HEVC `get_intra_scan_idx`: `log2==2 || (log2==3 && (cIdx==0 || 4:4:4))`).
pub(crate) fn scan_idx_for(mode: u8, log2_ts: u32, is_luma: bool, is_444: bool) -> u8 {
    let mode_dependent = log2_ts == 2 || (log2_ts == 3 && (is_luma || is_444));
    if !mode_dependent {
        return 0;
    }
    if (6..=14).contains(&mode) {
        2 // vertical
    } else if (22..=30).contains(&mode) {
        1 // horizontal
    } else {
        0 // diagonal
    }
}

/// Raw scan order over a `blk`×`blk` array, returning `(x, y)` = `(col, row)`.
/// Runtime reference generator, kept only to cross-check the const tables in tests.
#[cfg(test)]
fn raw_scan(blk: usize, scan_idx: u8) -> Vec<(usize, usize)> {
    let mut v = Vec::with_capacity(blk * blk);
    match scan_idx {
        1 => {
            // horizontal: row-major
            for y in 0..blk {
                for x in 0..blk {
                    v.push((x, y));
                }
            }
        }
        2 => {
            // vertical: column-major
            for x in 0..blk {
                for y in 0..blk {
                    v.push((x, y));
                }
            }
        }
        _ => {
            // up-right diagonal (HEVC 6.5.3)
            let b = blk as i32;
            let (mut x, mut y) = (0i32, 0i32);
            loop {
                while y >= 0 {
                    if x < b && y < b {
                        v.push((x as usize, y as usize));
                    }
                    y -= 1;
                    x += 1;
                }
                y = x;
                x = 0;
                if v.len() >= blk * blk {
                    break;
                }
            }
        }
    }
    v
}

/// Coefficient scan for a TB as `(row, col)`, organised as a grid of 4×4
/// sub-blocks (each sub-block and the sub-block grid scanned in `scan_idx`
/// order) so that `coeffs[sb*16 + k]` indexing stays consistent. The diagonal
/// case returns the canonical static tables directly; horizontal/vertical are
/// generated once and cached.
pub(crate) fn coeff_scan(log2_ts: u32, scan_idx: u8) -> &'static [(usize, usize)] {
    if scan_idx == 0 {
        return match log2_ts {
            2 => &DIAG_SCAN_4X4,
            3 => &ZIGZAG,
            4 => &ZIGZAG_16X16,
            _ => panic!("unsupported transform size log2={log2_ts}"),
        };
    }
    match (log2_ts, scan_idx) {
        (2, 1) => &COEFF_SCAN_4_H,
        (2, 2) => &COEFF_SCAN_4_V,
        (3, 1) => &COEFF_SCAN_8_H,
        (3, 2) => &COEFF_SCAN_8_V,
        (4, 1) => &COEFF_SCAN_16_H,
        (4, 2) => &COEFF_SCAN_16_V,
        _ => panic!("unsupported scan log2={log2_ts} idx={scan_idx}"),
    }
}

/// Sub-block grid scan `(sbx, sby)` for a TB.
pub(crate) fn sb_scan_for(log2_ts: u32, scan_idx: u8) -> &'static [(usize, usize)] {
    match (log2_ts, scan_idx) {
        // sb_side == 1 → single sub-block, scan order is trivial for every idx.
        (2, _) => &SB_SCAN_1,
        (3, 0) => &SB_SCAN_2_D,
        (3, 1) => &SB_SCAN_2_H,
        (3, 2) => &SB_SCAN_2_V,
        (4, 0) => &SB_SCAN_4_D,
        (4, 1) => &SB_SCAN_4_H,
        (4, 2) => &SB_SCAN_4_V,
        _ => panic!("unsupported sb scan log2={log2_ts} idx={scan_idx}"),
    }
}

// ── Compile-time scan tables ────────────────────────────────────────────────
// Built by `const fn` at compile time so no table is generated or heap-allocated
// at run time. `const_scan_equiv` (tests) asserts they equal the runtime
// generator below, which the diagonal tests in turn pin to the canonical tables.

const fn raw_scan_const(blk: usize, scan_idx: u8) -> ([(usize, usize); 16], usize) {
    let mut v = [(0usize, 0usize); 16];
    let mut i = 0;
    match scan_idx {
        1 => {
            let mut y = 0;
            while y < blk {
                let mut x = 0;
                while x < blk {
                    v[i] = (x, y);
                    i += 1;
                    x += 1;
                }
                y += 1;
            }
        }
        2 => {
            let mut x = 0;
            while x < blk {
                let mut y = 0;
                while y < blk {
                    v[i] = (x, y);
                    i += 1;
                    y += 1;
                }
                x += 1;
            }
        }
        _ => {
            let b = blk as i32;
            let mut x = 0i32;
            let mut y = 0i32;
            loop {
                while y >= 0 {
                    if x < b && y < b {
                        v[i] = (x as usize, y as usize);
                        i += 1;
                    }
                    y -= 1;
                    x += 1;
                }
                y = x;
                x = 0;
                if i >= blk * blk {
                    break;
                }
            }
        }
    }
    (v, blk * blk)
}

const fn build_coeff_scan_const<const NN: usize>(
    log2_ts: u32,
    scan_idx: u8,
) -> [(usize, usize); NN] {
    let n = 1usize << log2_ts; // NN == n*n
    let mut out = [(0usize, 0usize); NN];
    if n <= 4 {
        let (raw, len) = raw_scan_const(n, scan_idx);
        let mut i = 0;
        while i < len {
            out[i] = (raw[i].1, raw[i].0);
            i += 1;
        }
        return out;
    }
    let sb_side = n / 4;
    let (sb, sb_len) = raw_scan_const(sb_side, scan_idx);
    let (inner, inner_len) = raw_scan_const(4, scan_idx);
    let mut o = 0;
    let mut s = 0;
    while s < sb_len {
        let (sbx, sby) = (sb[s].0, sb[s].1);
        let mut k = 0;
        while k < inner_len {
            let (ix, iy) = (inner[k].0, inner[k].1);
            out[o] = (sby * 4 + iy, sbx * 4 + ix);
            o += 1;
            k += 1;
        }
        s += 1;
    }
    out
}

const fn build_sb_scan_const<const NN: usize>(log2_ts: u32, scan_idx: u8) -> [(usize, usize); NN] {
    let sb_side = (1usize << log2_ts) / 4; // NN == sb_side*sb_side (>=1)
    let (raw, _len) = raw_scan_const(sb_side, scan_idx); // sb_side >= 1 for log2_ts >= 2
    let mut out = [(0usize, 0usize); NN];
    let mut i = 0;
    while i < NN {
        out[i] = raw[i];
        i += 1;
    }
    out
}

const COEFF_SCAN_4_H: [(usize, usize); 16] = build_coeff_scan_const::<16>(2, 1);
const COEFF_SCAN_4_V: [(usize, usize); 16] = build_coeff_scan_const::<16>(2, 2);
const COEFF_SCAN_8_H: [(usize, usize); 64] = build_coeff_scan_const::<64>(3, 1);
const COEFF_SCAN_8_V: [(usize, usize); 64] = build_coeff_scan_const::<64>(3, 2);
const COEFF_SCAN_16_H: [(usize, usize); 256] = build_coeff_scan_const::<256>(4, 1);
const COEFF_SCAN_16_V: [(usize, usize); 256] = build_coeff_scan_const::<256>(4, 2);

const SB_SCAN_1: [(usize, usize); 1] = build_sb_scan_const::<1>(2, 0);
const SB_SCAN_2_D: [(usize, usize); 4] = build_sb_scan_const::<4>(3, 0);
const SB_SCAN_2_H: [(usize, usize); 4] = build_sb_scan_const::<4>(3, 1);
const SB_SCAN_2_V: [(usize, usize); 4] = build_sb_scan_const::<4>(3, 2);
const SB_SCAN_4_D: [(usize, usize); 16] = build_sb_scan_const::<16>(4, 0);
const SB_SCAN_4_H: [(usize, usize); 16] = build_sb_scan_const::<16>(4, 1);
const SB_SCAN_4_V: [(usize, usize); 16] = build_sb_scan_const::<16>(4, 2);

#[cfg(test)]
fn build_coeff_scan(log2_ts: u32, scan_idx: u8) -> Vec<(usize, usize)> {
    let n = 1usize << log2_ts;
    if n <= 4 {
        return raw_scan(n, scan_idx)
            .into_iter()
            .map(|(x, y)| (y, x))
            .collect();
    }
    let sb_side = n / 4;
    let sb = raw_scan(sb_side, scan_idx);
    let inner = raw_scan(4, scan_idx);
    let mut out = Vec::with_capacity(n * n);
    for &(sbx, sby) in &sb {
        for &(ix, iy) in &inner {
            out.push((sby * 4 + iy, sbx * 4 + ix));
        }
    }
    out
}

#[cfg(test)]
fn build_sb_scan(log2_ts: u32, scan_idx: u8) -> Vec<(usize, usize)> {
    let sb_side = (1usize << log2_ts) / 4;
    if sb_side <= 1 {
        return vec![(0, 0)];
    }
    raw_scan(sb_side, scan_idx)
}

#[cfg(test)]
mod scan_tests {
    use super::*;
    #[test]
    fn diagonal_matches_static_tables() {
        // The generator (used for the horizontal/vertical scans) must reproduce
        // the canonical diagonal tables exactly when given scan_idx 0.
        assert_eq!(&build_coeff_scan(2, 0)[..], &DIAG_SCAN_4X4[..], "4x4 diag");
        assert_eq!(&build_coeff_scan(3, 0)[..], &ZIGZAG[..], "8x8 diag");
        assert_eq!(&build_coeff_scan(4, 0)[..], &ZIGZAG_16X16[..], "16x16 diag");
    }
    #[test]
    fn const_tables_match_runtime_generator() {
        // Production uses the const tables; this pins them to the runtime generator
        // (which the diagonal test above pins to the canonical HEVC tables).
        assert_eq!(&COEFF_SCAN_4_H[..], &build_coeff_scan(2, 1)[..]);
        assert_eq!(&COEFF_SCAN_4_V[..], &build_coeff_scan(2, 2)[..]);
        assert_eq!(&COEFF_SCAN_8_H[..], &build_coeff_scan(3, 1)[..]);
        assert_eq!(&COEFF_SCAN_8_V[..], &build_coeff_scan(3, 2)[..]);
        assert_eq!(&COEFF_SCAN_16_H[..], &build_coeff_scan(4, 1)[..]);
        assert_eq!(&COEFF_SCAN_16_V[..], &build_coeff_scan(4, 2)[..]);
        assert_eq!(&SB_SCAN_1[..], &build_sb_scan(2, 0)[..]);
        assert_eq!(&SB_SCAN_2_D[..], &build_sb_scan(3, 0)[..]);
        assert_eq!(&SB_SCAN_2_H[..], &build_sb_scan(3, 1)[..]);
        assert_eq!(&SB_SCAN_2_V[..], &build_sb_scan(3, 2)[..]);
        assert_eq!(&SB_SCAN_4_D[..], &build_sb_scan(4, 0)[..]);
        assert_eq!(&SB_SCAN_4_H[..], &build_sb_scan(4, 1)[..]);
        assert_eq!(&SB_SCAN_4_V[..], &build_sb_scan(4, 2)[..]);
    }
    #[test]
    fn scan_idx_selection() {
        assert_eq!(scan_idx_for(0, 3, true, false), 0); // PLANAR diag
        assert_eq!(scan_idx_for(6, 3, true, false), 2); // mode 6 vertical (8x8 luma)
        assert_eq!(scan_idx_for(26, 3, true, false), 1); // mode 26 horizontal
        assert_eq!(scan_idx_for(6, 4, true, false), 0); // 16x16 always diagonal
        assert_eq!(scan_idx_for(26, 2, false, false), 1); // 4x4 chroma mode-dependent
        assert_eq!(scan_idx_for(6, 3, false, false), 0); // 8x8 chroma non-444: diagonal
        assert_eq!(scan_idx_for(6, 3, false, true), 2); // 8x8 chroma 4:4:4: mode-dependent
    }
}
