#include "hash_matrix.h"

// Generator matrices and identity
const hash_matrix_t MATRIX_I = {{1, 0, 0, 1}};
const hash_matrix_t MATRIX_A = {{1, 2, 0, 1}};
const hash_matrix_t MATRIX_B = {{1, 0, 2, 1}};

// 64-bit multiplication returning 128-bit result
uint128_t mul_u64(uint64_t x, uint64_t y) {
    return (uint128_t)x * y;
}

// 128-bit addition
uint128_t add_u128(uint128_t x, uint128_t y) {
    return x + y;
}

// Modular reduction helper functions (Dresdenboy's algorithm for Mersenne primes)
static inline uint128_t mod_p_round_1(uint128_t n) {
    uint64_t low_bits = (uint64_t)n & P;  // 63 bits
    uint64_t intermediate_bits = (uint64_t)(n >> 63);  // remaining bits
    return (uint128_t)low_bits + intermediate_bits;
}

static inline uint64_t mod_p_round_2(uint128_t n) {
    uint64_t low_bits = (uint64_t)n & P;
    uint64_t intermediate_bits = (uint64_t)(n >> 63);
    return low_bits + intermediate_bits;
}

// Fast modular reduction mod p = 2^63 - 1
uint64_t mod_p(uint128_t n) {
    // First reduction: n could be up to 128 bits
    if (n >> 64) {
        n = mod_p_round_1(n);
    }

    // Second reduction: n is at most 65 bits
    uint64_t n_small;
    if (n > P) {
        n_small = mod_p_round_2(n);
    } else {
        n_small = (uint64_t)n;
    }

    // Third reduction if needed
    if (n_small > P) {
        n_small = (n_small & P) + (n_small >> 63);
    }

    // Handle special case: p mod p = 0
    if (n_small == P) {
        return 0;
    }
    return n_small;
}

// Matrix multiplication mod p
hash_matrix_t matmul(hash_matrix_t a, hash_matrix_t b) {
    hash_matrix_t result;

    // result[0] = (a[0]*b[0] + a[1]*b[2]) mod p
    result.m[0] = mod_p(add_u128(mul_u64(a.m[0], b.m[0]), mul_u64(a.m[1], b.m[2])));

    // result[1] = (a[0]*b[1] + a[1]*b[3]) mod p
    result.m[1] = mod_p(add_u128(mul_u64(a.m[0], b.m[1]), mul_u64(a.m[1], b.m[3])));

    // result[2] = (a[2]*b[0] + a[3]*b[2]) mod p
    result.m[2] = mod_p(add_u128(mul_u64(a.m[2], b.m[0]), mul_u64(a.m[3], b.m[2])));

    // result[3] = (a[2]*b[1] + a[3]*b[3]) mod p
    result.m[3] = mod_p(add_u128(mul_u64(a.m[2], b.m[1]), mul_u64(a.m[3], b.m[3])));

    return result;
}

// Compare two matrices for equality
bool hash_matrix_eq(const hash_matrix_t* a, const hash_matrix_t* b) {
    return a->m[0] == b->m[0] && a->m[1] == b->m[1] &&
           a->m[2] == b->m[2] && a->m[3] == b->m[3];
}
