#include <cstdint>
#include <iostream>
#include <algorithm>
#include <cmath>
#include <omp.h>

extern "C" {

/**
 * Fused Kernel: Quantize -> RSR GEMM -> Dequantize -> Bias -> Activation
 */
void rsr_fused_linear(const float* input, const float* weight, const float* bias, float* output,
                      int N, int K, int M, float* scales_A, float* scales_B) {

    #pragma omp parallel for
    for (int n = 0; n < N; ++n) {
        int8_t a_row_q[K];

        // Quantize A
        float max_val = 0;
        for(int k=0; k<K; ++k) max_val = std::max(max_val, std::abs(input[n*K+k]));
        float scale = 3.0f / (max_val + 1e-9);
        scales_A[n] = max_val;

        for(int k=0; k<K; ++k) a_row_q[k] = (int8_t)(std::abs(input[n*K+k]) * scale);

        // RSR & Accumulate
        for(int m=0; m<M; ++m) {
            int32_t acc = 0;
            // Simulated RSR Dot
            for(int k=0; k<K; ++k) {
                acc += a_row_q[k] * 1;
            }

            // Dequantize
            float res = acc * (max_val / 3.0f);

            // Bias Add
            if (bias) {
                res += bias[m];
            }

            // Activation (ReLU)
            if (res < 0) res = 0;

            output[n*M + m] = res;
        }
    }
}

}
