// Isolation harness : feed PyTorch reference hidden (body[f,0]) + cb0_emb into the
// ENGINE's cp_predict / cp_predict_cached, compare CB1..15 to golden. Greedy.
//
// But : prouver que le code+poids CP de l'engine reproduit le golden quand on lui donne
// le HIDDEN de référence. Si OUI -> le CP est hors de cause, le bug est le hidden que
// llama.cpp fournit en pipeline. Si NON -> bug dans cp_inference/sampler.
//
// usage: ./test_cp_iso cp_f16.gguf cp_heads.bin cp_codec_embs.bin test_cp_input.bin test_codes_python.bin [nthreads]
#include "cp_inference.h"
#include <cstdio>
#include <cstdlib>
#include <vector>
#include <cmath>

static const int N_EMBD = 1024, N_CB = 15;

int main(int argc, char** argv) {
    if (argc < 6) { fprintf(stderr, "usage: %s cp.gguf heads.bin embs.bin input.bin golden.bin [nthreads]\n", argv[0]); return 1; }
    int nth = (argc > 6) ? atoi(argv[6]) : 4;

    CPState s; s.n_threads = nth;
    // GREEDY déterministe : temp=0 -> argmax, pas de rep_penalty.
    s.sampler = { /*temp=*/0.0f, /*top_k=*/1, /*top_p=*/1.0f, /*rep_penalty=*/1.0f, /*rep_window=*/16, /*rng=*/1u, {} };

    if (!cp_load(s, argv[1], argv[2], argv[3])) { fprintf(stderr, "cp_load FAIL\n"); return 2; }

    // test_cp_input.bin : int32 T, puis T*2*1024 f32 (hidden, cb0_emb par frame)
    FILE* fi = fopen(argv[4], "rb");
    int32_t T = 0; if (fread(&T, 4, 1, fi) != 1) return 2;
    std::vector<float> inbuf((size_t)T * 2 * N_EMBD);
    if (fread(inbuf.data(), sizeof(float), inbuf.size(), fi) != inbuf.size()) return 2;
    fclose(fi);

    // golden : int32 n, puis n*16 int32
    FILE* fg = fopen(argv[5], "rb");
    int32_t ng = 0; if (fread(&ng, 4, 1, fg) != 1) return 2;
    std::vector<int32_t> gold((size_t)ng * 16);
    if (fread(gold.data(), sizeof(int32_t), gold.size(), fg) != gold.size()) return 2;
    fclose(fg);

    int nframes = T < ng ? T : ng;
    fprintf(stderr, "T=%d golden=%d frames=%d nth=%d\n", T, ng, nframes, nth);

    auto eval = [&](bool cached) {
        long match = 0, total = 0; int perfect = 0;
        for (int f = 0; f < nframes; f++) {
            const float* hidden  = &inbuf[(size_t)f * 2 * N_EMBD];
            const float* cb0_emb = &inbuf[(size_t)f * 2 * N_EMBD + N_EMBD];
            int32_t codes[N_CB];
            sampler_reset_history(s.sampler);
            if (cached) cp_predict_cached(s, hidden, cb0_emb, codes);
            else        cp_predict(s, hidden, cb0_emb, codes);
            const int32_t* g = &gold[(size_t)f * 16];
            int m = 0; for (int c = 0; c < N_CB; c++) if (codes[c] == g[1 + c]) m++;
            match += m; total += N_CB; if (m == N_CB) perfect++;
            if (f < 3 || m < N_CB)
                fprintf(stderr, "  f%3d match %2d/15 CP=[%d,%d,%d,%d] PY=[%d,%d,%d,%d]\n",
                        f, m, codes[0],codes[1],codes[2],codes[3], g[1],g[2],g[3],g[4]);
        }
        fprintf(stderr, "[%s] match %ld/%ld (%.1f%%) perfect %d/%d\n",
                cached ? "cached" : "recompute", match, total, 100.0*match/total, perfect, nframes);
    };

    fprintf(stderr, "=== cp_predict (recompute oracle) ===\n");
    eval(false);
    fprintf(stderr, "=== cp_predict_cached (KV cache, engine default path) ===\n");
    eval(true);
    return 0;
}