/* ffn_batch_test: M1.6 S1 — ffn_forward_b (batched FFN over M tokens via
 * gemm_gpu_v4) vs the decode per-token oracle (ffn_forward called M times).
 *
 * Synthesizes a dev_ffn_blk with random ternary Q2 weights in the exact GGUF
 * layout (same pattern as gemm_engine_test), runs both paths on the same
 * random x[M][D_MODEL], and compares the residual-stream output.
 *
 * Gate (per plan.md S1): mean|d| of FFN output <= ~0.05 vs per-token decode.
 * NOT bit-exact — the batched GEMM reassociates fp32 differently than the
 * per-token gemv reduction (ground rule 2: greedy-token gate, drift OK).
 *
 * Build: meson compile -C build ffn_batch_test && ./build/ffn_batch_test
 */
#include <algorithm>
#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(0xF1FFULL);

/* ---- Q2 weight synthesis + repack (matches gemm_engine_test) ------------ */
static uint16_t f32_to_f16_bits(float f) {
    return sycl::bit_cast<uint16_t>(sycl::half(f));
}
static float f16_to_f32_local(uint16_t bits) {
    return float(sycl::bit_cast<sycl::half>(bits));
}

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 to A_codes [M/8][ngrp][2][64B] + A_sc_t [ngrp][M] (decode upload_q2
 * layout, shared with gemm_gpu_v4 per M1.6 layout discovery). */
static dev_q2 upload_synth(sycl::queue& q, const std::vector<qx_block_q2>& W,
                           int M, int K) {
    int ngrp = K / QX_QK2;
    size_t codes_sz = (size_t)(M / 8) * ngrp * 128;
    size_t sc_sz = (size_t)ngrp * M;
    uint8_t* codes_h = (uint8_t*)aligned_alloc(32, codes_sz);
    float*    sc_h   = (float*)aligned_alloc(32, sc_sz * sizeof(float));
    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_h + (((size_t)(m8 * ngrp + gg) * 2 + half) * 64 + r * 8),
                           W[(m8 * 8 + r) * ngrp + gg].qs + half * 8, 8);
    for (int m = 0; m < M; m++)
        for (int gg = 0; gg < ngrp; gg++)
            sc_h[gg * M + m] = f16_to_f32_local(W[m * ngrp + gg].d);
    dev_q2 d;
    d.codes = sycl::malloc_shared<uint8_t>(codes_sz, q);
    d.scales = sycl::malloc_shared<float>(sc_sz, q);
    q.memcpy(d.codes, codes_h, codes_sz).wait();
    q.memcpy(d.scales, sc_h, sc_sz * sizeof(float)).wait();
    free(codes_h); free(sc_h);
    d.M = M; d.K = K; d.src = nullptr;
    return d;
}

/* Fill an [n_tok][dim] tile with varied row scales (matches the
 * gemm_engine_test distribution: rows pick scales across decades, with a
 * zero-row and an outlier to exercise amax edge cases). */
static void fill_tile(float* p, int n_tok, int dim, int seed_off) {
    std::normal_distribution<float> g(0.0f, 1.0f);
    for (int t = 0; t < n_tok; t++) {
        float sc = (((t + seed_off) % 17) == 3) ? 0.0f
                   : powf(10.0f, -2.0f + 3.0f * (((t + seed_off) % 13) / 12.0f));
        for (int i = 0; i < dim; i++) p[(size_t)t * dim + i] = g(rng) * sc;
        if (((t + seed_off) % 31) == 5) p[(size_t)t * dim + 7] = 123.0f;
    }
}

struct Cfg { int M; const char* note; };

int main() {
    sycl::queue q(sycl::gpu_selector_v, sycl::property::queue::in_order{});
    printf("ffn_batch_test (M1.6 S1): ffn_forward_b vs decode ffn_forward oracle\n");

    /* ---- synthesize FFN block 0 shapes ---- */
    const int Md = QX_D_MODEL;     /* 5120 */
    const int Mf = QX_FF_INTER;    /* 17408 */
    auto W_gate = synth_q2_weight(Mf, Md);
    auto W_up   = synth_q2_weight(Mf, Md);
    auto W_down = synth_q2_weight(Md, Mf);

    dev_ffn_blk f;
    f.gate = upload_synth(q, W_gate, Mf, Md);
    f.up   = upload_synth(q, W_up,   Mf, Md);
    f.down = upload_synth(q, W_down, Md, Mf);
    std::vector<float> post_norm(Md);
    {
        std::normal_distribution<float> g(0.0f, 1.0f);
        for (int i = 0; i < Md; i++) post_norm[i] = 1.0f + 0.1f * g(rng);
    }
    f.post_norm = sycl::malloc_shared<float>(Md, q);
    q.memcpy(f.post_norm, post_norm.data(), Md * sizeof(float)).wait();

    /* ---- batched workspace (sized for the largest M we test) ---- */
    const int Mmax = 256;
    /* bvnni_d layout = [K/4][M] int32 = K*M bytes; max K = FF_INTER = 17408. */
    int8_t* bvn_d = sycl::malloc_shared<int8_t>((size_t)Mf * Mmax, q);
    int32_t* sum_d = sycl::malloc_shared<int32_t>((size_t)(Mf / QX_QK2) * Mmax, q);
    float* asc_d = sycl::malloc_shared<float>(Mmax, q);
    float* g_b = sycl::malloc_shared<float>((size_t)Mmax * Mf, q);
    float* u_b = sycl::malloc_shared<float>((size_t)Mmax * Mf, q);
    float* t_b = sycl::malloc_shared<float>((size_t)Mmax * Mf, q);
    float* x_b = sycl::malloc_shared<float>((size_t)Mmax * Md, q);

    /* ---- decode oracle workspace (single-token) ---- */
    /* The decode path uses globals via engine_set_quant_stage.  Device-sized
     * (malloc_device would suffice; shared is fine and simpler to inspect). */
    int8_t* dec_bvn = sycl::malloc_shared<int8_t>(Mf, q);
    int32_t* dec_sum = sycl::malloc_shared<int32_t>(Mf / QX_QK2, q);
    float*   dec_asc = sycl::malloc_shared<float>(1, q);
    engine_set_quant_stage(dec_bvn, dec_sum, dec_asc);
    float* dec_part = sycl::malloc_shared<float>((size_t)8 * Mf, q);
    float* x_dec = sycl::malloc_shared<float>(Md, q);
    float* dec_g = sycl::malloc_shared<float>(Mf, q);
    float* dec_u = sycl::malloc_shared<float>(Mf, q);
    float* dec_t = sycl::malloc_shared<float>(Mf, q);

    Cfg cfgs[] = {
        {16,  "smallest valid GEMM block"},
        {64,  "mid chunk"},
        {256, "full chunk"},
    };

    int fail = 0;
    for (auto c : cfgs) {
        int M = c.M;
        printf("\n=== M=%d  (%s) ===\n", M, c.note);

        /* seed x identically for both paths. */
        fill_tile(x_b, M, Md, 1);
        std::vector<float> x_save((size_t)M * Md);
        memcpy(x_save.data(), x_b, (size_t)M * Md * sizeof(float));

        /* ---- batched: ffn_forward_b on the M tokens ---- */
        ffn_forward_b(q, &f, x_b, M, bvn_d, sum_d, asc_d, g_b, u_b, t_b);
        q.wait();

        /* ---- oracle: decode ffn_forward M times, row by row ---- */
        std::vector<float> x_oracle((size_t)M * Md);
        for (int tok = 0; tok < M; tok++) {
            memcpy(x_dec, x_save.data() + (size_t)tok * Md, Md * sizeof(float));
            ffn_forward(q, &f, x_dec, dec_part,
                        dec_g, dec_u, dec_t);
            q.wait();
            memcpy(x_oracle.data() + (size_t)tok * Md, x_dec, Md * sizeof(float));
        }

        /* ---- compare ---- */
        double sum_abs = 0; float worst = 0; size_t n = (size_t)M * Md;
        for (size_t i = 0; i < n; i++) {
            double d = fabs(x_b[i] - x_oracle[i]);
            sum_abs += d;
            if (d > worst) worst = (float)d;
        }
        double mean_abs = sum_abs / n;
        /* reference magnitude for context */
        double ref_mag = 0;
        for (size_t i = 0; i < n; i++) ref_mag += fabs(x_oracle[i]);
        ref_mag /= n;
        bool pass = mean_abs <= 0.05;
        printf("  mean|d|=%.3e  worst=%.3e  (mean|x|=%.3e)  %s\n",
               mean_abs, worst, ref_mag, pass ? "PASSED" : "FAILED (gate 0.05)");
        if (!pass) fail++;
    }

    /* ---- perf: M=256 back-to-back to amortize clock ramp ---- */
    {
        int M = 256, REPS = 30;
        q.wait();
        auto t0 = std::chrono::steady_clock::now();
        for (int r = 0; r < REPS; r++)
            ffn_forward_b(q, &f, x_b, M, bvn_d, sum_d, asc_d, g_b, u_b, t_b);
        q.wait();
        auto t1 = std::chrono::steady_clock::now();
        double ms = std::chrono::duration<double, std::milli>(t1 - t0).count() / REPS;
        /* 3 GEMMs: 2x (Mf*Md) + 1x (Md*Mf), per token. */
        double macs = (double)(2 * Mf * Md + Md * Mf);
        double tflops = macs * M / (ms / 1e3) / 1e12;
        double tokseq = M / (ms / 1e3);
        printf("\nperf M=256: %.2f ms/FFN-block, %.1f TFLOPs, %.0f tok/s (FFN only)\n",
               ms, tflops, tokseq);
    }

    /* cleanup */
    auto free_devq2 = [&](dev_q2& d){ sycl::free(d.codes, q); sycl::free(d.scales, q); };
    free_devq2(f.gate); free_devq2(f.up); free_devq2(f.down);
    sycl::free(f.post_norm, q);
    sycl::free(bvn_d, q); sycl::free(sum_d, q); sycl::free(asc_d, q);
    sycl::free(g_b, q); sycl::free(u_b, q); sycl::free(t_b, q); sycl::free(x_b, q);
    sycl::free(dec_bvn, q); sycl::free(dec_sum, q); sycl::free(dec_asc, q);
    sycl::free(dec_part, q); sycl::free(x_dec, q);
    sycl::free(dec_g, q); sycl::free(dec_u, q); sycl::free(dec_t, q);

    printf(fail ? "\nFAILED: %d config(s)\n"
                : "\nALL PASSED — mean|d| within gate vs decode oracle\n", fail);
    return fail ? 1 : 0;
}