use crate::hash::hash_u32;

/// Performs a base-2 Owen scramble on an integer.
pub fn owen2(n: u32, seed: u32) -> u32 {
    let mut result = n;

    for i in 0..32 {
        let mask = (!0 << 1) << i;
        result ^= hash_u32(n & mask, seed) & (1 << i);
    }

    result
}

/// Performs a base-4 Owen scramble on an integer.
///
/// This is motivated by z-scrambling from the paper "Screen-Space
/// Blue-Noise Diffusion of Monte Carlo Sampling Error via Hierarchical
/// Ordering of Pixels" by Ahmed et al.  Their scrambling function is
/// actually just Owen scrambling in base 4, but their attempt at
/// implementing that function isn't great, both in terms of memory
/// usage and scramble quality.  (The paper itself is otherwise
/// fantastic, though.)
///
/// The implementation below is a full proper base-4 Owen scramble,
/// requiring only a 24-byte lookup table.
pub fn owen4(n: u32, seed: u32) -> u32 {
    const PERMUTATION_TABLE: [u8; 24] = [
        0 | (1 << 2) | (2 << 4) | (3 << 6), // [0, 1, 2, 3],
        0 | (1 << 2) | (3 << 4) | (2 << 6), // [0, 1, 3, 2],
        0 | (2 << 2) | (1 << 4) | (3 << 6), // [0, 2, 1, 3],
        0 | (2 << 2) | (3 << 4) | (1 << 6), // [0, 2, 3, 1],
        0 | (3 << 2) | (1 << 4) | (2 << 6), // [0, 3, 1, 2],
        0 | (3 << 2) | (2 << 4) | (1 << 6), // [0, 3, 2, 1],
        1 | (0 << 2) | (2 << 4) | (3 << 6), // [1, 0, 2, 3],
        1 | (0 << 2) | (3 << 4) | (2 << 6), // [1, 0, 3, 2],
        1 | (2 << 2) | (0 << 4) | (3 << 6), // [1, 2, 0, 3],
        1 | (2 << 2) | (3 << 4) | (0 << 6), // [1, 2, 3, 0],
        1 | (3 << 2) | (0 << 4) | (2 << 6), // [1, 3, 0, 2],
        1 | (3 << 2) | (2 << 4) | (0 << 6), // [1, 3, 2, 0],
        2 | (0 << 2) | (1 << 4) | (3 << 6), // [2, 0, 1, 3],
        2 | (0 << 2) | (3 << 4) | (1 << 6), // [2, 0, 3, 1],
        2 | (1 << 2) | (0 << 4) | (3 << 6), // [2, 1, 0, 3],
        2 | (1 << 2) | (3 << 4) | (0 << 6), // [2, 1, 3, 0],
        2 | (3 << 2) | (0 << 4) | (1 << 6), // [2, 3, 0, 1],
        2 | (3 << 2) | (1 << 4) | (0 << 6), // [2, 3, 1, 0],
        3 | (0 << 2) | (1 << 4) | (2 << 6), // [3, 0, 1, 2],
        3 | (0 << 2) | (2 << 4) | (1 << 6), // [3, 0, 2, 1],
        3 | (1 << 2) | (0 << 4) | (2 << 6), // [3, 1, 0, 2],
        3 | (1 << 2) | (2 << 4) | (0 << 6), // [3, 1, 2, 0],
        3 | (2 << 2) | (0 << 4) | (1 << 6), // [3, 2, 0, 1],
        3 | (2 << 2) | (1 << 4) | (0 << 6), // [3, 2, 1, 0],
    ];

    let mut result = 0;

    for i in 0..16 {
        let mask = (!0 << 2) << (i * 2);
        let perm_entry = PERMUTATION_TABLE[(hash_u32(n & mask, seed) % 24) as usize];
        let perm_cell_idx = ((n >> (i * 2)) & 0b11) as usize;

        result |= (((perm_entry >> (perm_cell_idx * 2)) & 0b11) as u32) << (i * 2);
    }

    result
}