#include <cstdint>
#include <cmath>
#include <algorithm>
#include <vector>

extern "C" {

/**
 * Quantize BFloat16 (simulated as float32 for convenience in this prototype) to INT8.
 *
 * Logic:
 * 1. Find max absolute value in block (Micro-scaling).
 * 2. Scale = 127 / max_val.
 * 3. Quantized = round(Input * Scale).
 *
 * Ideally input should be uint16_t* for BF16, but ctypes handles float* easier.
 */
void quantize_bf16_to_int8(const float* input, int8_t* output, float* scales, int N, int K, int block_size) {
    // Process row by row, block by block
    for (int n = 0; n < N; ++n) {
        for (int k = 0; k < K; k += block_size) {
            int k_end = std::min(k + block_size, K);
            int current_block_size = k_end - k;

            // 1. Find Max
            float max_val = 1e-9; // Avoid div by zero
            for (int i = 0; i < current_block_size; ++i) {
                float val = std::abs(input[n * K + k + i]);
                if (val > max_val) max_val = val;
            }

            // 2. Compute Scale
            float scale = 3.0f / max_val; // We want 2-bit quantization [0..3] for A, or [-X..X] for B.
            // The prompt says "A: 2-bit integers [0,3]".
            // Let's assume we are quantizing for the A matrix of RSR which expects [0,3].
            // Mapping: min -> 0, max -> 3? Or centered?
            // RSR code expects A in [0,3].
            // Let's map [-max, max] to [0, 3] roughly for demo.

            // Actually, RSR usually works on weights B (precomputed) and activations A (runtime).
            // Let's implement a standard symmetric quantization to [-127, 127] for generic INT8,
            // and a specific one for TRME A-matrix [0,3].

            // Mode: TRME A-Matrix (2-bit unsigned)
            scale = 3.0f / max_val;
            scales[(n * K + k) / block_size] = max_val; // Store inverse scale or max

            for (int i = 0; i < current_block_size; ++i) {
                float val = input[n * K + k + i];
                // Simple abs projection for 2-bit magnitude
                int8_t q = (int8_t)(std::abs(val) * scale);
                if (q > 3) q = 3;
                output[n * K + k + i] = q;
            }
        }
    }
}

}
