#include <vector>
#include <cstdint>
#include <iostream>
#include <cstring>
#include <algorithm>
#include <omp.h>
#include <immintrin.h>

template <int BLOCK_SIZE>
void rsr_gemm_template(const int8_t* A, const int8_t* B, int32_t* C, int N, int K, int M) {
    int max_combinations = 1 << (2 * BLOCK_SIZE);
    std::vector<int16_t> lut(max_combinations * M);

    for (int k_start = 0; k_start < K; k_start += BLOCK_SIZE) {
        int k_end = std::min(k_start + BLOCK_SIZE, K);
        int current_block_width = k_end - k_start;
        int current_combinations = 1 << (2 * current_block_width);

        // Fill LUT
        #pragma omp parallel for
        for (int comb = 0; comb < current_combinations; ++comb) {
            int temp_comb = comb;
            for (int m = 0; m < M; ++m) {
                int16_t val = 0;
                int temp_c = temp_comb;
                for (int i = 0; i < current_block_width; ++i) {
                    int coeff = temp_c % 4;
                    temp_c /= 4;
                    if (k_start + i < K)
                        val += coeff * B[(k_start + i) * M + m];
                }
                lut[comb * M + m] = val;
            }
        }

        // Apply
        #pragma omp parallel for
        for (int n = 0; n < N; ++n) {
            int idx = 0;
            int power_of_4 = 1;
            for (int i = 0; i < current_block_width; ++i) {
                int8_t val = 0;
                if (k_start + i < K)
                    val = A[n * K + (k_start + i)];
                idx += (val & 3) * power_of_4;
                power_of_4 *= 4;
            }

            for (int m = 0; m < M; ++m) {
                C[n * M + m] += lut[idx * M + m];
            }
        }
    }
}

extern "C" {

void rsr_gemm_dispatch(const int8_t* A, const int8_t* B, int32_t* C, const int32_t* bias, int N, int K, int M, int block_size) {
    // Init C with Bias or 0
    #pragma omp parallel for
    for (int i = 0; i < N * M; ++i) {
        if (bias) C[i] = bias[i % M]; // Broadcast bias
        else C[i] = 0;
    }

    // Dispatch to Templated Kernels
    if (block_size == 4) {
        rsr_gemm_template<4>(A, B, C, N, K, M);
    } else if (block_size == 6) {
        rsr_gemm_template<6>(A, B, C, N, K, M);
    } else if (block_size == 8) {
        rsr_gemm_template<8>(A, B, C, N, K, M);
    } else {
        rsr_gemm_template<4>(A, B, C, N, K, M);
    }
}

}
