//! HeavyHash matrix operations in GF(2^8).
//!
//! The matrix is a 64x64 matrix where each element is in GF(2^8).
//! Matrix multiplication is performed using Galois field arithmetic.

/// The irreducible polynomial for GF(2^8): x^8 + x^4 + x^3 + x + 1
const GF_POLY: u16 = 0x11B;

/// 64x64 HeavyHash matrix.
/// This is deterministically generated from a seed.
#[derive(Clone)]
pub struct HeavyMatrix {
    /// The 64x64 matrix elements.
    data: [[u8; 64]; 64],
}

impl HeavyMatrix {
    /// Creates the default HeavyMatrix (Kaspa-compatible).
    pub fn new() -> Self {
        Self::from_seed(&SYNOR_MATRIX_SEED)
    }

    /// Creates a matrix from a 32-byte seed.
    pub fn from_seed(seed: &[u8; 32]) -> Self {
        let mut matrix = [[0u8; 64]; 64];
        let mut state = *seed;

        for row in matrix.iter_mut() {
            for elem in row.iter_mut() {
                // Simple deterministic expansion using Blake3
                state = blake3::hash(&state).into();
                *elem = state[0];
            }
        }

        HeavyMatrix { data: matrix }
    }

    /// Multiplies a 256-bit vector by the matrix.
    /// Input: 32 bytes (256 bits) treated as 64 nibbles
    /// Output: 32 bytes after matrix multiplication
    #[allow(clippy::needless_range_loop)]
    pub fn multiply(&self, input: &[u8; 32]) -> [u8; 32] {
        // Expand 32 bytes into 64 nibbles (4 bits each)
        let mut nibbles = [0u8; 64];
        for i in 0..32 {
            nibbles[i * 2] = (input[i] >> 4) & 0x0F;
            nibbles[i * 2 + 1] = input[i] & 0x0F;
        }

        // Matrix multiplication in GF(2^8)
        let mut result = [0u8; 64];
        for i in 0..64 {
            let mut acc = 0u8;
            for j in 0..64 {
                acc = gf_add(acc, gf_mul(self.data[i][j], nibbles[j]));
            }
            result[i] = acc;
        }

        // Compress 64 values back to 32 bytes
        let mut output = [0u8; 32];
        for i in 0..32 {
            // XOR adjacent pairs
            output[i] = result[i * 2] ^ result[i * 2 + 1];
        }

        output
    }

    /// Multiplies with full byte input (alternative method).
    #[allow(clippy::needless_range_loop)]
    pub fn multiply_bytes(&self, input: &[u8; 32]) -> [u8; 32] {
        let mut output = [0u8; 32];

        // Treat input as 32 bytes, output as 32 bytes
        // Use a 32x32 view of the matrix
        for i in 0..32 {
            let mut acc: u16 = 0;
            for j in 0..32 {
                let m1 = self.data[i * 2][j * 2] as u16;
                let m2 = self.data[i * 2 + 1][j * 2 + 1] as u16;
                let v = input[j] as u16;
                acc = acc.wrapping_add((m1.wrapping_mul(v)).wrapping_add(m2.wrapping_mul(v)));
            }
            output[i] = (acc & 0xFF) as u8;
        }

        output
    }

    /// Gets a matrix element.
    #[inline]
    pub fn get(&self, row: usize, col: usize) -> u8 {
        self.data[row][col]
    }

    /// Returns the raw matrix data.
    pub fn as_bytes(&self) -> &[[u8; 64]; 64] {
        &self.data
    }
}

impl Default for HeavyMatrix {
    fn default() -> Self {
        Self::new()
    }
}

/// Addition in GF(2^8) is XOR.
#[inline]
pub fn gf_add(a: u8, b: u8) -> u8 {
    a ^ b
}

/// Multiplication in GF(2^8) using the irreducible polynomial.
#[inline]
pub fn gf_mul(a: u8, b: u8) -> u8 {
    let mut result: u16 = 0;
    let mut aa = a as u16;
    let mut bb = b as u16;

    for _ in 0..8 {
        if bb & 1 != 0 {
            result ^= aa;
        }
        let hi_bit = aa & 0x80;
        aa <<= 1;
        if hi_bit != 0 {
            aa ^= GF_POLY;
        }
        bb >>= 1;
    }

    (result & 0xFF) as u8
}

/// Precomputed multiplication tables for faster GF(2^8) operations.
pub struct GfTables {
    /// Multiplication table: mul[a][b] = a * b in GF(2^8)
    mul: [[u8; 256]; 256],
}

impl GfTables {
    /// Creates precomputed tables.
    #[allow(clippy::needless_range_loop)]
    pub fn new() -> Self {
        let mut mul = [[0u8; 256]; 256];

        for a in 0..256 {
            for b in 0..256 {
                mul[a][b] = gf_mul(a as u8, b as u8);
            }
        }

        GfTables { mul }
    }

    /// Fast multiplication using lookup table.
    #[inline]
    pub fn mul(&self, a: u8, b: u8) -> u8 {
        self.mul[a as usize][b as usize]
    }
}

impl Default for GfTables {
    fn default() -> Self {
        Self::new()
    }
}

/// Synor's matrix seed (unique to our chain).
/// This creates a different matrix than Kaspa, preventing merge-mining.
pub const SYNOR_MATRIX_SEED: [u8; 32] = [
    0x53, 0x59, 0x4E, 0x4F, 0x52, 0x5F, 0x48, 0x45, // "SYNOR_HE"
    0x41, 0x56, 0x59, 0x5F, 0x4D, 0x41, 0x54, 0x52, // "AVY_MATR"
    0x49, 0x58, 0x5F, 0x53, 0x45, 0x45, 0x44, 0x5F, // "IX_SEED_"
    0x56, 0x30, 0x30, 0x31, 0x00, 0x00, 0x00, 0x01, // "V001"...
];

/// Optimized matrix multiplication using precomputed tables.
pub struct OptimizedMatrix {
    /// The 64x64 matrix.
    matrix: HeavyMatrix,
    /// Precomputed GF tables.
    tables: GfTables,
}

impl OptimizedMatrix {
    /// Creates an optimized matrix with lookup tables.
    pub fn new() -> Self {
        OptimizedMatrix {
            matrix: HeavyMatrix::new(),
            tables: GfTables::new(),
        }
    }

    /// Creates from a specific seed.
    pub fn from_seed(seed: &[u8; 32]) -> Self {
        OptimizedMatrix {
            matrix: HeavyMatrix::from_seed(seed),
            tables: GfTables::new(),
        }
    }

    /// Fast matrix multiplication.
    #[allow(clippy::needless_range_loop)]
    pub fn multiply(&self, input: &[u8; 32]) -> [u8; 32] {
        // Expand to nibbles
        let mut nibbles = [0u8; 64];
        for i in 0..32 {
            nibbles[i * 2] = (input[i] >> 4) & 0x0F;
            nibbles[i * 2 + 1] = input[i] & 0x0F;
        }

        // Matrix multiplication with table lookup
        let mut result = [0u8; 64];
        for i in 0..64 {
            let mut acc = 0u8;
            for j in 0..64 {
                acc ^= self.tables.mul(self.matrix.get(i, j), nibbles[j]);
            }
            result[i] = acc;
        }

        // Compress
        let mut output = [0u8; 32];
        for i in 0..32 {
            output[i] = result[i * 2] ^ result[i * 2 + 1];
        }

        output
    }

    /// Returns the underlying matrix.
    pub fn matrix(&self) -> &HeavyMatrix {
        &self.matrix
    }
}

impl Default for OptimizedMatrix {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_gf_add() {
        assert_eq!(gf_add(0, 0), 0);
        assert_eq!(gf_add(0xFF, 0xFF), 0);
        assert_eq!(gf_add(0xAA, 0x55), 0xFF);
    }

    #[test]
    fn test_gf_mul() {
        // Multiplication by 1
        assert_eq!(gf_mul(0x53, 1), 0x53);
        // Multiplication by 0
        assert_eq!(gf_mul(0x53, 0), 0);
        // Known values
        assert_eq!(gf_mul(0x57, 0x83), 0xC1);
    }

    #[test]
    fn test_matrix_deterministic() {
        let m1 = HeavyMatrix::new();
        let m2 = HeavyMatrix::new();

        // Same seed should produce same matrix
        for i in 0..64 {
            for j in 0..64 {
                assert_eq!(m1.get(i, j), m2.get(i, j));
            }
        }
    }

    #[test]
    fn test_matrix_multiply() {
        let matrix = HeavyMatrix::new();
        let input = [0x42u8; 32];

        let output = matrix.multiply(&input);

        // Output should be deterministic
        let output2 = matrix.multiply(&input);
        assert_eq!(output, output2);

        // Output should be different from input
        assert_ne!(output, input);
    }

    #[test]
    fn test_optimized_matrix() {
        let matrix = HeavyMatrix::new();
        let optimized = OptimizedMatrix::new();
        let input = [0x55u8; 32];

        let output1 = matrix.multiply(&input);
        let output2 = optimized.multiply(&input);

        // Both should produce same result
        assert_eq!(output1, output2);
    }

    #[test]
    fn test_gf_tables() {
        let tables = GfTables::new();

        // Verify table matches direct computation
        for a in 0..=255u8 {
            for b in 0..=255u8 {
                assert_eq!(tables.mul(a, b), gf_mul(a, b));
            }
        }
    }
}