/* qxmx_diff: per-layer diff between the Arc Battlemage GPU engine and the CPU ref oracle.
 *
 * Runs both engines on the same prompt, dumps the residual stream (x[5120])
 * after each layer's attention+FFN, and reports where they diverge.
 *
 *   ./qxmx_diff [-ctk fp16|q8_0|fp8] model.gguf "prompt"
 */
#include "qxmx_gpu.h"
#include "qxmx_model.h"
#include "qxmx_ref.h"
#include "qxmx_tokenizer.h"

#include <cmath>
#include <cstdio>
#include <cstring>
#include <string>
#include <vector>

namespace qx { extern qx_model* g_model; }

int main(int argc, char** argv) {
    qx::engine::KCT k_ctk = qx::engine::KCT::q8_0;
    std::vector<const char*> pos;
    for (int i = 1; i < argc; i++) {
        std::string a = argv[i];
        if (a == "-ctk" && i + 1 < argc) {
            std::string v = argv[++i];
            if      (v == "fp16") k_ctk = qx::engine::KCT::fp16;
            else if (v == "q8_0") k_ctk = qx::engine::KCT::q8_0;
            else if (v == "fp8")  k_ctk = qx::engine::KCT::fp8;
            else { std::fprintf(stderr, "-ctk: unknown type '%s'\n", v.c_str()); return 1; }
        } else {
            pos.push_back(argv[i]);
        }
    }
    if (pos.size() < 2) { std::fprintf(stderr, "usage: %s [-ctk fp16|q8_0|fp8] model.gguf \"prompt\"\n", argv[0]); return 1; }
    qx_model m; char err[256];
    if (qx_model_open(&m, pos[0], err, sizeof err)) { std::fprintf(stderr, "load: %s\n", err); return 1; }
    qx::tokenizer tok;
    if (!tok.load(&m.gguf, err, sizeof err)) { std::fprintf(stderr, "tok: %s\n", err); return 1; }

    std::vector<int32_t> ids;
    tok.encode(pos[1], &ids);
    std::printf("prompt: %zu tokens\n", ids.size());
    if (ids.empty()) { std::fprintf(stderr, "empty prompt\n"); return 1; }

    /* --- CPU ref oracle --- */
    qx::ref_session ref_sess;
    ref_sess.init(&m, std::max((int)ids.size() + 16, 1 << 14));

    /* --- GPU engine --- */
    qx::g_model = &m;
    qx::engine eng;
    eng.k_ctk = k_ctk;
    eng.init(&m, std::max((int)ids.size() + 16, 1 << 14));
    int slot = eng.open_slot();

    /* Feed the first prompt token through both, compare per-layer.
     * The CPU ref's ref_decode runs all 64 layers at once — we need to
     * instrument it to dump per-layer.  Instead, replicate the ref's
     * per-layer loop here using the ref's block-forward functions.
     * But those are static in qxmx_ref.cpp... so we'll compare at the
     * output level: run the first token through both, compare the final
     * logits and the argmax token. */

    /* Run first token through CPU ref.  ref_l2 (declared here, populated in
     * the prefill block when ids.size()>=5) holds the full-prompt ref logits;
     * the final per-logit summary picks the matching one. */
    std::vector<float> ref_x(QX_D_MODEL), ref_logits(QX_N_VOCAB);
    std::vector<float> ref_l2(QX_N_VOCAB);
    qx::ref_dequant_row(&m.token_embd, ids[0], ref_x.data());
    qx::ref_decode(&m, &ref_sess, ref_x.data(), ref_logits.data());

    /* Find CPU ref argmax over the full vocab. */
    int ref_best_full = 0; float ref_bv_full = -1e30f;
    for (int v = 0; v < QX_N_VOCAB; v++) if (ref_logits[v] > ref_bv_full) { ref_bv_full = ref_logits[v]; ref_best_full = v; }
    std::string ref_tok = tok.decode(&ref_best_full, 1, true);
    std::printf("CPU ref: argmax=%d \"%s\" (logit=%.3f)\n", ref_best_full, ref_tok.c_str(), ref_bv_full);

    /* Run first token through B70 engine */
    eng.forward(slot, ids[0]);
    int eng_best = eng.sample(slot, qx::sample_params{});
    std::string eng_tok = tok.decode(&eng_best, 1, true);
    std::printf("GPU eng: argmax=%d \"%s\"\n", eng_best, eng_tok.c_str());

    /* Also run the full 5-token prefill through both and compare the
     * next-token prediction (the one that should produce " Paris"). */
    if (ids.size() >= 5) {
        std::printf("\n=== full %zu-token prefill comparison ===\n", ids.size());
        /* CPU ref: feed all tokens, get logits after last */
        std::vector<float> ref_x2(QX_D_MODEL);
        for (size_t i = 0; i < ids.size(); i++) {
            qx::ref_dequant_row(&m.token_embd, ids[i], ref_x2.data());
            qx::ref_decode(&m, &ref_sess, ref_x2.data(), ref_l2.data());
        }
        int ref_best2 = 0; float ref_bv2 = -1e30f;
        for (int v = 0; v < QX_N_VOCAB; v++) if (ref_l2[v] > ref_bv2) { ref_bv2 = ref_l2[v]; ref_best2 = v; }
        std::string ref_tok2 = tok.decode(&ref_best2, 1, true);
        std::printf("CPU ref after %zu tokens: argmax=%d \"%s\" (logit=%.3f)\n",
                    ids.size(), ref_best2, ref_tok2.c_str(), ref_bv2);

        /* Engine: feed all tokens */
        int eng_next = 0;
        for (size_t i = 0; i < ids.size(); i++) {
            eng.forward(slot, ids[i]);
            eng_next = eng.sample(slot, qx::sample_params{});
        }
        std::string eng_tok2 = tok.decode(&eng_next, 1, true);
        std::printf("GPU eng after %zu tokens: argmax=%d \"%s\" (logit=%.3f)\n",
                    ids.size(), eng_next, eng_tok2.c_str(), eng.logits[eng_next]);

        /* logits diff after full prefill */
        double md2 = 0; float wd2 = 0; int wi2 = 0;
        for (int v = 0; v < QX_N_VOCAB; v++) {
            double d = fabs(ref_l2[v] - eng.logits[v]); md2 += d;
            if (d > wd2) { wd2 = (float)d; wi2 = v; }
        }
        std::printf("logits: mean|diff|=%.3e worst=%.3e @[%d]\n", md2/QX_N_VOCAB, wd2, wi2);
    }

    /* Final per-logit summary.  eng.logits now holds the engine's output after
     * consuming ALL ids.size() tokens (the prefill loop above).  Compare it
     * against the matching ref: ref_l2 (full-prompt ref, when ids.size()>=5)
     * or ref_logits (single-token ref).  Mixing the two — as an earlier
     * version did — compared token-0 ref against token-(N-1) engine logits and
     * produced a bogus ~1.3 mean|diff| that was purely a harness artifact. */
    const float* ref_cmp = (ids.size() >= 5) ? ref_l2.data() : ref_logits.data();
    std::printf("\nlogits comparison (first 20):\n");
    for (int v = 0; v < 20; v++)
        std::printf("  [%5d] ref=%8.3f  eng=%8.3f  diff=%8.3f\n",
                    v, ref_cmp[v], eng.logits[v], ref_cmp[v] - eng.logits[v]);

    /* Compute mean|diff| over all logits */
    double md = 0; float wd = 0; int wi = 0;
    for (int v = 0; v < QX_N_VOCAB; v++) {
        double d = fabs(ref_cmp[v] - eng.logits[v]);
        md += d;
        if (d > wd) { wd = (float)d; wi = v; }
    }
    std::printf("\nlogits: mean|diff|=%.3e  worst=%.3e @[%d]\n", md/QX_N_VOCAB, wd, wi);
    std::printf("  ref[wi]=%.3f  eng[wi]=%.3f\n", ref_cmp[wi], eng.logits[wi]);

    qx_model_close(&m);
    return 0;
}