// Test bit-exact du Talker engine vs Python.
//
// Stratégie = teacher-forcing : on rejoue les MÊMES inputs_embeds (prefill +
// step) que Python a effectivement utilisés, et on compare logits/hidden à
// chaque step. Pas de sampling côté engine — on consomme step_inputs[i] tel
// qu'observé côté Python.
//
// Si match (RMSE/cos serrés sur logits ET hidden) → Patch 2 (M-RoPE IMROPE) +
// Patch 1 (embeds API) sont corrects. Sinon : où ça diverge donne la piste.
//
// Fixtures (dans <dump_dir>/):
//   manifest.txt                         (T_prefill, N_steps, n_embd, vocab)
//   talker_prefill_embeds.bin            float32 [T_prefill * n_embd]
//   talker_step_inputs.bin               float32 [N_steps  * n_embd]
//   talker_step_logits.bin               float32 [N_steps  * vocab]
//   talker_step_hidden.bin               float32 [N_steps  * n_embd]
//   talker_last_hidden_pf.bin            float32 [n_embd]   (past_hidden après prefill)
//
// Usage : test_talker_replay <gguf> <dump_dir> [cpu|htp]
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <cstdint>
#include <vector>
#include <string>
#include <fstream>
#include "llama.h"
#include "ggml-backend.h"

static ggml_backend_dev_t find_htp() {
    for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
        auto d = ggml_backend_dev_get(i);
        if (!strcmp(ggml_backend_dev_name(d), "HTP0")) return d;
    }
    return nullptr;
}

static std::vector<float> read_f32(const std::string& p, size_t n_expected) {
    std::ifstream f(p, std::ios::binary | std::ios::ate);
    if (!f) { fprintf(stderr, "open %s\n", p.c_str()); exit(1); }
    size_t n = (size_t)f.tellg() / sizeof(float);
    if (n != n_expected) { fprintf(stderr, "%s: %zu f32, attendu %zu\n", p.c_str(), n, n_expected); exit(1); }
    f.seekg(0);
    std::vector<float> v(n);
    f.read((char*)v.data(), n * sizeof(float));
    return v;
}

// RMSE + cos + max_abs entre deux vecteurs de même taille.
struct Cmp { float rmse, cos, mae, max_abs_diff; int argmax_a, argmax_b; };
static Cmp compare(const float* a, const float* b, int n) {
    double se = 0, dot = 0, na = 0, nb = 0, mae = 0;
    float mad = 0;
    int am = 0, bm = 0;
    float av = a[0], bv = b[0];
    for (int i = 0; i < n; ++i) {
        double d = a[i] - b[i];
        se += d * d;
        mae += std::abs(d);
        if (std::abs((float)d) > mad) mad = std::abs((float)d);
        dot += (double)a[i] * b[i];
        na  += (double)a[i] * a[i];
        nb  += (double)b[i] * b[i];
        if (a[i] > av) { av = a[i]; am = i; }
        if (b[i] > bv) { bv = b[i]; bm = i; }
    }
    return Cmp{
        (float)std::sqrt(se / n),
        (float)(dot / (std::sqrt(na) * std::sqrt(nb) + 1e-30)),
        (float)(mae / n),
        mad,
        am, bm,
    };
}

static void print_cmp(const char* tag, int step, const Cmp& c, int extra_argmax = -1) {
    printf("  %s[%2d] rmse=%.4e cos=%.6f mae=%.4e max=%.4e argmax(eng,py", tag, step, c.rmse, c.cos, c.mae, c.max_abs_diff);
    if (extra_argmax >= 0) printf(",dump");
    printf(")=(%d,%d", c.argmax_a, c.argmax_b);
    if (extra_argmax >= 0) printf(",%d", extra_argmax);
    printf(")%s\n", c.argmax_a == c.argmax_b ? " ✓" : "");
}

int main(int argc, char** argv) {
    if (argc < 3) { printf("usage: %s <gguf> <dump_dir> [cpu|htp]\n", argv[0]); return 1; }
    const char* gguf = argv[1];
    std::string D = argv[2]; if (D.back() != '/') D += '/';
    bool force_cpu = (argc >= 4 && !strcmp(argv[3], "cpu"));

    // manifest
    int T_prefill = 0, N_steps = 0, n_embd = 1024, n_vocab = 3072;
    {
        std::ifstream f(D + "manifest.txt"); if (!f) { fprintf(stderr, "no manifest\n"); return 1; }
        std::string line;
        while (std::getline(f, line)) {
            if (line.rfind("T_prefill:", 0) == 0) T_prefill = atoi(line.c_str() + 10);
            else if (line.rfind("N_steps:", 0) == 0) N_steps = atoi(line.c_str() + 8);
            else if (line.rfind("n_embd:", 0) == 0)  n_embd  = atoi(line.c_str() + 7);
            else if (line.rfind("vocab:", 0) == 0)   n_vocab = atoi(line.c_str() + 6);
        }
        printf("manifest: T_prefill=%d N_steps=%d n_embd=%d n_vocab=%d\n", T_prefill, N_steps, n_embd, n_vocab);
    }

    // load fixtures
    auto prefill_embeds = read_f32(D + "talker_prefill_embeds.bin", (size_t)T_prefill * n_embd);
    auto step_inputs    = read_f32(D + "talker_step_inputs.bin",    (size_t)N_steps  * n_embd);
    auto step_logits_py = read_f32(D + "talker_step_logits.bin",    (size_t)N_steps  * n_vocab);
    auto step_hidden_py = read_f32(D + "talker_step_hidden.bin",    (size_t)N_steps  * n_embd);
    auto hidden_pf_py   = read_f32(D + "talker_last_hidden_pf.bin", (size_t)n_embd);
    printf("fixtures loaded (Python)\n");

    // engine
    setenv("GGML_HEXAGON_USE_HMX", "0", 1);
    llama_backend_init();
    auto mp = llama_model_default_params();
    ggml_backend_dev_t devs[2] = { force_cpu ? nullptr : find_htp(), nullptr };
    if (devs[0]) { mp.n_gpu_layers = 99; mp.devices = devs; printf("device: HTP0\n"); }
    else         { mp.n_gpu_layers = 0;  printf("device: CPU%s\n", force_cpu ? " (forcé)" : ""); }
    auto m = llama_model_load_from_file(gguf, mp);
    if (!m) { printf("model load FAILED\n"); return 1; }

    auto cp = llama_context_default_params();
    cp.n_ctx = std::max(512, T_prefill + N_steps + 16);
    cp.n_batch = 1024; cp.n_threads = 4;
    cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
    cp.embeddings = true;
    auto ctx = llama_init_from_model(m, cp);
    if (!ctx) { printf("ctx FAILED\n"); return 1; }
    printf("ctx ready (n_ctx=%d, n_embd=%d, n_vocab=%d)\n", cp.n_ctx, llama_model_n_embd(m), llama_vocab_n_tokens(llama_model_get_vocab(m)));

    // M-RoPE (talker) : 4 positions par token (t/h/w/e). Embeds-mode = pas de
    // conversion auto par le framework, on fournit explicitement T*npe entiers.
    auto rt = llama_model_rope_type(m);
    const int npe = (rt == LLAMA_ROPE_TYPE_MROPE || rt == LLAMA_ROPE_TYPE_IMROPE) ? 4 : 1;
    printf("rope_type=%d, n_pos_per_embd=%d\n", (int)rt, npe);

    // === PREFILL ===
    {
        std::vector<llama_pos>     pos(T_prefill * npe, 0);
        std::vector<int32_t>       nsd(T_prefill, 1);
        std::vector<llama_seq_id>  sid0(T_prefill, 0);
        std::vector<llama_seq_id*> sids(T_prefill);
        std::vector<int8_t>        lg(T_prefill, 0);
        for (int i = 0; i < T_prefill; ++i) {
            if (npe == 4) { pos[i] = i; pos[T_prefill + i] = i; pos[2*T_prefill + i] = i; pos[3*T_prefill + i] = 0; }
            else          { pos[i] = i; }
            sids[i] = &sid0[i];
        }
        lg[T_prefill - 1] = 1;
        llama_batch b{};
        b.n_tokens = T_prefill; b.embd = prefill_embeds.data();
        b.pos = pos.data(); b.n_seq_id = nsd.data(); b.seq_id = sids.data(); b.logits = lg.data();
        if (llama_decode(ctx, b) != 0) { printf("prefill FAILED\n"); return 1; }

        const float* h = llama_get_embeddings_ith(ctx, -1);
        Cmp ch = compare(h, hidden_pf_py.data(), n_embd);
        printf("PREFILL OK (T=%d)\n", T_prefill);
        printf("  hidden_pf: rmse=%.4e cos=%.6f mae=%.4e max=%.4e\n", ch.rmse, ch.cos, ch.mae, ch.max_abs_diff);
    }

    // === STEPS (teacher-forced) ===
    int n_ok_argmax = 0;
    double sum_rmse_h = 0, sum_rmse_l = 0, sum_cos_h = 0, sum_cos_l = 0;
    for (int s = 0; s < N_steps; ++s) {
        llama_pos     pos1[4] = {0,0,0,0};
        const llama_pos p = T_prefill + s;
        if (npe == 4) { pos1[0] = p; pos1[1] = p; pos1[2] = p; pos1[3] = 0; }
        else          { pos1[0] = p; }
        int32_t       n = 1;
        llama_seq_id  sd = 0; llama_seq_id* sp = &sd;
        int8_t        l = 1;
        llama_batch b{};
        b.n_tokens = 1; b.embd = step_inputs.data() + (size_t)s * n_embd;
        b.pos = pos1; b.n_seq_id = &n; b.seq_id = &sp; b.logits = &l;
        if (llama_decode(ctx, b) != 0) { printf("step %d FAILED\n", s); return 1; }

        const float* h_eng  = llama_get_embeddings_ith(ctx, -1);
        const float* lg_eng = llama_get_logits_ith(ctx, -1);
        Cmp ch = compare(h_eng,  step_hidden_py.data() + (size_t)s * n_embd,  n_embd);
        Cmp cl = compare(lg_eng, step_logits_py.data() + (size_t)s * n_vocab, n_vocab);
        if (ch.argmax_a == ch.argmax_b) n_ok_argmax++;  // matching argmax over hidden (just a curiosity)
        sum_rmse_h += ch.rmse; sum_cos_h += ch.cos;
        sum_rmse_l += cl.rmse; sum_cos_l += cl.cos;
        // n'imprime que les 3 premiers + les 3 derniers steps (diff de pos / cumul KV)
        if (s < 3 || s >= N_steps - 3) {
            print_cmp("hidden", s, ch);
            print_cmp("logits", s, cl);
        }
        else if (s == 3) printf("  ... (%d steps intermédiaires) ...\n", N_steps - 6);
    }
    printf("STEPS SUMMARY (N=%d):\n", N_steps);
    printf("  hidden: mean_rmse=%.4e mean_cos=%.6f\n", sum_rmse_h/N_steps, sum_cos_h/N_steps);
    printf("  logits: mean_rmse=%.4e mean_cos=%.6f (argmax-on-hidden-match=%d/%d, info only)\n",
           sum_rmse_l/N_steps, sum_cos_l/N_steps, n_ok_argmax, N_steps);

    llama_free(ctx);
    llama_model_free(m);
    return 0;
}