/* gemm_engine_test: M1.6 go/no-go — port gemm_gpu_v4 into the engine and
 * validate (a) correctness vs a CPU fp32 reference and (b) in-engine sustained
 * TFLOPs with a real activation producer (gpu_rmsnorm_quant_b) feeding B.
 *
 * Shapes from the real model (FFN gate/up/down, attn qkv):
 *   ffn_gate/up: M=17408, K=5120   attn_q: M=12288, K=5120
 *   ffn_down:    M=5120,  K=17408
 *
 * We can't load the real weights here, so we synthesize ternary Q2 weights
 * (codes {-1,0,1}, random f16 scale per 64) in the exact GGUF layout, run
 * rmsnorm_quant_b to produce the B operand, run the GEMM, and compare.
 *
 * Build: meson compile -C build gemm_engine_test && ./build/gemm_engine_test
 */
#include <chrono>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <random>
#include <vector>

#include <sycl/sycl.hpp>

#include "qxmx.h"
#include "qxmx_gpu.h"
#include "qxmx_gemm.h"

using namespace qx;

static std::mt19937_64 rng(0x6E6E6EULL);

/* f32 -> f16 bits (software; matches the GGUF/ggml_half layout). */
static uint16_t f32_to_f16_bits(float f) {
    sycl::half h = sycl::half(f);
    return sycl::bit_cast<uint16_t>(h);
}

/* f16 bits -> f32 (software; matches qxmx_gpu.cpp::f16_to_f32). */
static float f16_to_f32_local(uint16_t bits) {
    return float(sycl::bit_cast<sycl::half>(bits));
}

/* Synthesize a Q2_0_g64 weight matrix in GGUF layout (qx_block_q2 array):
 * M rows, K cols (K multiple of 64).  codes in {-1,0,1} -> q in {0,1,2}. */
static std::vector<qx_block_q2> synth_q2_weight(int M, int K) {
    int ngrp = K / QX_QK2;
    std::vector<qx_block_q2> W((size_t)M * ngrp);
    std::normal_distribution<float> g(0.0f, 1.0f);
    for (int m = 0; m < M; m++)
        for (int gg = 0; gg < ngrp; gg++) {
            float sc = 0.05f * std::fabs(g(rng)) + 0.001f;
            W[(size_t)m * ngrp + gg].d = f32_to_f16_bits(sc);
            for (int b = 0; b < QX_QK2 / 4; b++) {
                uint8_t pack = 0;
                for (int k = 0; k < 4; k++) {
                    int q = rng() % 3;   /* {0,1,2} -> value {−1,0,+1} */
                    pack |= (uint8_t)(q << (k * 2));
                }
                W[(size_t)m * ngrp + gg].qs[b] = pack;
            }
        }
    return W;
}

/* Repack Q2 weight -> engine A_codes layout [M/8][ngrp][2][64B] + A_sc_t [ngrp][M].
 * Identical to qxmx_gpu.cpp::repack_q2 (kept local so the test is standalone). */
static void repack_q2_local(const qx_block_q2* src, int M, int K,
                            uint8_t* codes, float* scales_t) {
    int ngrp = K / QX_QK2;
    for (int m8 = 0; m8 < M / 8; m8++)
        for (int gg = 0; gg < ngrp; gg++)
            for (int half = 0; half < 2; half++)
                for (int r = 0; r < 8; r++)
                    memcpy(codes + (((size_t)(m8 * ngrp + gg) * 2 + half) * 64 + r * 8),
                           src[(m8 * 8 + r) * ngrp + gg].qs + half * 8, 8);
    for (int m = 0; m < M; m++)
        for (int gg = 0; gg < ngrp; gg++)
            scales_t[gg * M + m] = f16_to_f32_local(src[m * ngrp + gg].d);
}

/* CPU fp32 reference: C = W @ rmsnorm(x_in), matching the GPU path which
 * feeds rmsnorm_quant_b(x_in) into the GEMM.  rmsnorm is fp32 here; the GPU
 * adds s8 quantization noise on top (~1e-3 mean|d| expected). */
static void rmsnorm_row(const float* x, const float* w, float* out, int K) {
    double ss = 0;
    for (int k = 0; k < K; k++) ss += (double)x[k] * x[k];
    float inv = 1.0f / sqrtf((float)(ss / K) + QX_RMS_EPS);
    for (int k = 0; k < K; k++) out[k] = x[k] * inv * w[k];
}

static void cpu_ref_row(const qx_block_q2* W, const float* b, float* acc,
                        int M, int K) {
    /* For one output row m of W (M weights), accumulate into acc[m] += W[m]·b. */
    int ngrp = K / QX_QK2;
    for (int m = 0; m < M; m++) {
        const qx_block_q2* row = W + (size_t)m * ngrp;
        float p = 0;
        for (int gg = 0; gg < ngrp; gg++) {
            float d = f16_to_f32_local(row[gg].d);
            float s = 0;
            for (int k = 0; k < QX_QK2; k++) {
                int q = (row[gg].qs[k / 4] >> ((k % 4) * 2)) & 3;
                s += (float)(q - 1) * b[gg * QX_QK2 + k];
            }
            p += d * s;
        }
        acc[m] += p;
    }
}

struct Shape { int M, K; const char* name; };

int main() {
    sycl::queue q(sycl::gpu_selector_v, sycl::property::queue::in_order{});
    printf("gemm_engine_test (M1.6 go/no-go): gemm_gpu_v4 in-engine, real producer feeding B\n");

    Shape shapes[] = {
        {17408, 5120, "ffn_gate/up"},
        {5120, 17408, "ffn_down"},
        {12288, 5120, "attn_q"},
    };
    int N = 256;   /* n_tok = chunk size */

    int bad_total = 0;
    for (auto sh : shapes) {
        int M = sh.M, K = sh.K;
        int ngrp = K / QX_QK2;
        printf("\n=== %s  M=%d K=%d N=%d ===\n", sh.name, M, K, N);

        /* ---- synthesize weight + activation ---- */
        auto W = synth_q2_weight(M, K);
        std::vector<float> x_in((size_t)N * K);
        std::normal_distribution<float> g(0.0f, 1.0f);
        for (int t = 0; t < N; t++) {
            float sc = powf(10.0f, -2.0f + 3.0f * (t % 7) / 6.0f);
            for (int k = 0; k < K; k++) x_in[(size_t)t * K + k] = g(rng) * sc;
        }
        /* rmsnorm weight (K entries) = 1+noise so x*inv*w is non-trivial */
        std::vector<float> rms_w(K);
        for (int k = 0; k < K; k++) rms_w[k] = 1.0f + 0.1f * g(rng);

        /* ---- upload + repack weight to A layout ---- */
        size_t codes_sz = (size_t)(M / 8) * ngrp * 128;
        size_t sc_sz = (size_t)ngrp * M;
        uint8_t* A_codes = (uint8_t*)aligned_alloc(32, codes_sz);
        float* A_sc = (float*)aligned_alloc(32, sc_sz * sizeof(float));
        repack_q2_local(W.data(), M, K, A_codes, A_sc);
        uint8_t* d_A_codes = sycl::malloc_shared<uint8_t>(codes_sz, q);
        float* d_A_sc = sycl::malloc_shared<float>(sc_sz, q);
        q.memcpy(d_A_codes, A_codes, codes_sz).wait();
        q.memcpy(d_A_sc, A_sc, sc_sz * sizeof(float)).wait();
        free(A_codes); free(A_sc);

        /* ---- B operand: produced by gpu_rmsnorm_quant_b (the real producer) ---- */
        float* d_x = sycl::malloc_shared<float>((size_t)N * K, q);
        float* d_rmsw = sycl::malloc_shared<float>(K, q);
        memcpy(d_x, x_in.data(), (size_t)N * K * sizeof(float));
        memcpy(d_rmsw, rms_w.data(), K * sizeof(float));
        /* B_vnni: [K/4][N] int32 = (K/4 * N * 4) bytes.  sum_a: [ngrp][N]. */
        int8_t* d_B = sycl::malloc_shared<int8_t>((size_t)(K / 4) * N * 4, q);
        int32_t* d_sa = sycl::malloc_shared<int32_t>((size_t)ngrp * N, q);
        float* d_asc = sycl::malloc_shared<float>(N, q);
        gpu_rmsnorm_quant(q, d_x, d_rmsw, K, d_B, d_sa, d_asc, N);
        q.wait();

        /* ---- run GEMM ---- */
        float* d_C = sycl::malloc_shared<float>((size_t)M * N, q);
        gemm_gpu_v4(q, d_A_codes, d_A_sc, d_B, d_sa, d_asc, d_C, M, K, N, nullptr, true);

        /* ---- correctness: CPU ref applies rmsnorm (fp32) then W@b per sampled row.
         * GPU path = rmsnorm_quant_b (fp32 rmsnorm + s8 quant) + GEMM.  Expected
         * mean|d| ~ 1e-3 from the s8 quantization noise. ---- */
        int M_SAMPLE = std::min(M, 64);
        std::vector<int> sample_rows(M_SAMPLE);
        for (int i = 0; i < M_SAMPLE; i++) sample_rows[i] = (int)((uint64_t)rng() % M);
        std::vector<float> normed(K);
        double md = 0; float wd = 0; int ncmp = 0;
        for (int n = 0; n < N; n++) {
            rmsnorm_row(x_in.data() + (size_t)n * K, rms_w.data(), normed.data(), K);
            for (int i = 0; i < M_SAMPLE; i++) {
                int m = sample_rows[i];
                const qx_block_q2* row = W.data() + (size_t)m * ngrp;
                float acc = 0;
                for (int gg = 0; gg < ngrp; gg++) {
                    float d = f16_to_f32_local(row[gg].d);
                    float p = 0;
                    for (int k = 0; k < QX_QK2; k++) {
                        int qq = (row[gg].qs[k / 4] >> ((k % 4) * 2)) & 3;
                        p += (float)(qq - 1) * normed[gg * QX_QK2 + k];
                    }
                    acc += d * p;
                }
                double dd = fabs(acc - d_C[(size_t)m * N + n]);
                md += dd; ncmp++;
                if (dd > wd) wd = (float)dd;
            }
        }
        printf("  correctness vs CPU fp32+rmsnorm (sampled %d rows x %d toks): mean|d|=%.3e worst=%.3e\n",
               M_SAMPLE, N, md / ncmp, wd);

        /* ---- perf: sustained TFLOPs over 50 back-to-back calls ---- */
        int REPS = 50;
        q.wait();
        auto t0 = std::chrono::steady_clock::now();
        for (int r = 0; r < REPS; r++)
            gemm_gpu_v4(q, d_A_codes, d_A_sc, d_B, d_sa, d_asc, d_C, M, K, N, nullptr, false);
        q.wait();
        auto t1 = std::chrono::steady_clock::now();
        double ms = std::chrono::duration<double, std::milli>(t1 - t0).count() / REPS;
        double flops = 2.0 * (double)M * K * N;
        double tflops = flops / (ms / 1e3) / 1e12;
        /* tok/s-eq: per-token GEMM cost across all 64 layers is ~25.7B MACs
         * (M1.2 convention). tok/s = TFLOPs * 1e12 / 2 / 25.7e9. */
        double tokseq = tflops * 1e12 / 2.0 / 25.7e9;
        printf("  perf: %.2f ms/call, %.1f TFLOPs (%.0f tok/s-eq)\n", ms, tflops, tokseq);

        sycl::free(d_A_codes, q); sycl::free(d_A_sc, q);
        sycl::free(d_x, q); sycl::free(d_rmsw, q);
        sycl::free(d_B, q); sycl::free(d_sa, q); sycl::free(d_asc, q);
        sycl::free(d_C, q);
    }
    return bad_total ? 1 : 0;
}