// Compare repeated draft/accept rounds against a supplied reference JSON.
// Usage: test-dspark-loop <tiny.gguf> <ref.json>

#include "../src/llama-ext.h"
#include "common.h"
#include "llama.h"
#include "speculative.h"
#ifdef LLAMA_DSPARK_MARKOV_CUDA
#    include "dspark-markov.h"

#    include <algorithm>
#    include <cmath>
#    include <limits>
#endif
#ifdef LLAMA_DSPARK_MARKOV_METAL
#    include "dspark-markov-metal.h"

#    include <cmath>
#    include <limits>
#    include <memory>
#endif

#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <fstream>
#include <nlohmann/json.hpp>
#include <string>
#include <vector>

using json = nlohmann::json;

static void fail(const std::string & msg) {
    fprintf(stderr, "FAIL: %s\n", msg.c_str());
    exit(1);
}

// --- synthetic "target tap feature" stand-in ---------------------------
// Matches the reference fixture's
// hash_u32/synth_feat/synth_bonus_token (same constants, same integer ops --
// see that file's header comment for why this is safe to duplicate rather
// than share: it's a closed-form pure function of small integers, not a
// stateful RNG stream, so bit-parity across languages just falls out of
// using the same uint32 wraparound arithmetic).
static uint32_t hash_u32(uint32_t x) {
    x ^= x >> 16;
    x *= 0x7feb352du;
    x ^= x >> 15;
    x *= 0x846ca68bu;
    x ^= x >> 16;
    return x;
}

static float synth_feat(int64_t pos, int64_t d) {
    const uint32_t h = hash_u32((uint32_t) (pos * 131071 + d * 97 + 12345));
    const int32_t  m = (int32_t) (h % 2000u) - 1000;  // [-1000, 999]
    return (float) m / 500.0f;                        // [-2.0, 1.998]
}

static int32_t synth_bonus_token(int32_t round_idx, int32_t vocab_size, int32_t mask_token_id) {
    const uint32_t h = hash_u32((uint32_t) round_idx * 2654435761u + 999983u);
    int32_t        v = (int32_t) (h % (uint32_t) (vocab_size - 1));
    if (v == mask_token_id) {
        v = (v + 1) % vocab_size;
    }
    return v;
}

// mirrored verbatim from dspark_phase2_py_ref.py
static const std::vector<int32_t> ACCEPT_SCHEDULE = { 7, 3, 0, 7, 5, 1, 4 };
static const std::vector<int32_t> PROMPT          = { 1, 2, 3, 4, 5 };

static std::vector<float> synth_feat_rows(int64_t pos_beg, int64_t n_rows, int64_t n_embd_cap) {
    std::vector<float> feat((size_t) n_rows * n_embd_cap);
    for (int64_t i = 0; i < n_rows; ++i) {
        for (int64_t d = 0; d < n_embd_cap; ++d) {
            feat[(size_t) i * n_embd_cap + d] = synth_feat(pos_beg + i, d);
        }
    }
    return feat;
}

#ifdef LLAMA_DSPARK_MARKOV_METAL
static int test_metal_markov() {
    int  checks = 0;
    auto check  = [&](bool good, const char * label) {
        if (!good) {
            fail(label);
        }
        ++checks;
    };
    for (auto shape : {
             std::pair<int, int>{ 17,   1   },
              { 129,  31  },
              { 257,  256 },
              { 1025, 33  },
              { 8193, 256 }
    }) {
        const int          v = shape.first, rank = shape.second, mask = v - 1;
        std::vector<float> a(size_t(v) * rank), b(a.size()), base(size_t(v) * 4);
        for (size_t i = 0; i < a.size(); ++i) {
            a[i] = float(int(hash_u32(uint32_t(i)) % 31) - 15) / 16;
            b[i] = float(int(hash_u32(uint32_t(i) + 987) % 31) - 15) / 16;
        }
        for (size_t i = 0; i < base.size(); ++i) {
            base[i] = float(int(hash_u32(uint32_t(i) + 345) % 31) - 15) / 16;
        }
        std::unique_ptr<dspark_markov_metal, decltype(&dspark_markov_metal_free)> ctx(
            dspark_markov_metal_init(a.data(), b.data(), v, rank, mask), dspark_markov_metal_free);
        check(bool(ctx), "Metal context creation");
        for (int slots : { 1, 4, 2 }) {
            std::vector<int32_t> output(slots);
            check(dspark_markov_metal_resample(ctx.get(), base.data(), 1, slots, output.data()), "Metal resample");
            int previous = 1;
            for (int slot = 0; slot < slots; ++slot) {
                double best   = -std::numeric_limits<double>::infinity();
                int    winner = -1;
                for (int token = 0; token < v; ++token) {
                    if (token == mask) {
                        continue;
                    }
                    double score = base[size_t(slot) * v + token];
                    for (int r = 0; r < rank; ++r) {
                        score += double(a[size_t(previous) * rank + r]) * b[size_t(token) * rank + r];
                    }
                    if (score > best) {
                        best   = score;
                        winner = token;
                    }
                }
                check(output[slot] == winner, "Metal vs independent F64 chain oracle");
                previous = winner;
            }
        }
        int32_t out[4];
        check(!dspark_markov_metal_resample(ctx.get(), base.data(), mask, 1, out), "mask anchor rejected");
        check(!dspark_markov_metal_resample(ctx.get(), base.data(), -1, 1, out), "negative anchor rejected");
        check(!dspark_markov_metal_resample(ctx.get(), base.data(), v, 1, out), "out-of-range anchor rejected");
        check(!dspark_markov_metal_resample(ctx.get(), base.data(), 1, 0, out), "zero slots rejected");
        check(!dspark_markov_metal_resample(ctx.get(), base.data(), 1, 1, nullptr), "null output rejected");
        const float saved = base[0];
        for (float bad : { std::numeric_limits<float>::quiet_NaN(), std::numeric_limits<float>::infinity(),
                           -std::numeric_limits<float>::infinity() }) {
            base[0] = bad;
            check(!dspark_markov_metal_resample(ctx.get(), base.data(), 1, 4, out), "nonfinite score rejected");
        }
        base[0] = saved;
        check(dspark_markov_metal_resample(ctx.get(), base.data(), 1, 4, out), "recovery after rejected round");
    }
    std::vector<float>                                                        zeros(129 * 31), base(129 * 4);
    std::unique_ptr<dspark_markov_metal, decltype(&dspark_markov_metal_free)> ties(
        dspark_markov_metal_init(zeros.data(), zeros.data(), 129, 31, 0), dspark_markov_metal_free);
    check(bool(ties), "tie context");
    for (int slot = 0; slot < 4; ++slot) {
        base[slot * 129] = 1e30f;
    }
    int32_t out[4];
    check(dspark_markov_metal_resample(ties.get(), base.data(), 1, 4, out), "tie and dominant mask round");
    for (int token : out) {
        check(token == 1, "lowest unmasked ID wins ties");
    }
    for (int slot = 0; slot < 4; ++slot) {
        base[slot * 129 + 128] = 1;
    }
    check(dspark_markov_metal_resample(ties.get(), base.data(), 1, 4, out), "ragged final row round");
    for (int token : out) {
        check(token == 128, "ragged final row can win");
    }
    for (float excluded : { -std::numeric_limits<float>::infinity(), std::numeric_limits<float>::quiet_NaN() }) {
        base[0] = excluded;
        check(dspark_markov_metal_resample(ties.get(), base.data(), 1, 4, out), "excluded mask logit ignored");
    }
    check(!dspark_markov_metal_init(zeros.data(), zeros.data(), 129, 31, 129), "invalid mask metadata");
    zeros[0] = std::numeric_limits<float>::quiet_NaN();
    check(!dspark_markov_metal_init(zeros.data(), zeros.data(), 129, 31, 0), "nonfinite weights rejected");
    printf("Metal Markov: %d checks passed\n", checks);
    return 0;
}
#endif

#ifdef LLAMA_DSPARK_MARKOV_CUDA
static int test_cuda_dcut() {
    int checks = 0;
    for (int rank : { 1, 31, 33, 256 }) {
        const int          vocab = 67;
        std::vector<float> a(vocab * rank), b(vocab * rank), base(4 * vocab);
        for (size_t i = 0; i < a.size(); ++i) {
            a[i] = ((int) (i % 11) - 5) * 0.03f;
            b[i] = ((int) (i % 17) - 8) * 0.02f;
        }
        for (size_t i = 0; i < base.size(); ++i) {
            base[i] = std::sin((float) i) * 3.0f;
        }
        for (int k = 0; k < 4; ++k) {
            base[k * vocab + vocab - 1] = 1000.0f;
        }
        auto * ctx = dspark_markov_cuda_init(a.data(), b.data(), vocab, rank, vocab - 1);
        if (!ctx) {
            fail("CUDA D-cut init");
        }
        for (int n : { 4, 1, 3, 2, 4 }) {
            for (int mode = 0; mode < 3; ++mode) {
                float costs[4] = { 1, 1, 1, 1 };
                if (mode == 1) {
                    costs[1] = 20;
                    costs[2] = 30;
                    costs[3] = 40;
                }
                if (mode == 2) {
                    costs[1] = 1.01f;
                    costs[2] = 1.1f;
                    costs[3] = 1.2f;
                }
                int32_t got[4], original[4], kept;
                float   probs[4];
                if (!dspark_markov_cuda_resample(ctx, base.data(), 5, n, original) ||
                    !dspark_markov_cuda_dcut(ctx, base.data(), 5, n, got, costs, &kept, probs)) {
                    fail("CUDA D-cut run");
                }
                int    prev = 5, expected_keep = 1;
                double survival = 1, expected = 1, best_rate = -1;
                for (int k = 0; k < n; ++k) {
                    std::vector<double> logits(vocab);
                    for (int v = 0; v < vocab; ++v) {
                        logits[v] = base[k * vocab + v];
                        for (int r = 0; r < rank; ++r) {
                            logits[v] += (double) a[prev * rank + r] * b[v * rank + r];
                        }
                    }
                    logits[vocab - 1] = -std::numeric_limits<double>::infinity();
                    const int winner = std::max_element(logits.begin(), logits.end()) - logits.begin();
                    double    sum    = 0;
                    for (double logit : logits) {
                        sum += std::exp(logit - logits[winner]);
                    }
                    const double p = 1 / sum;
                    if (got[k] != winner || got[k] != original[k] || std::abs(probs[k] - p) > 1e-5) {
                        fail("CUDA D-cut F64 oracle");
                    }
                    survival *= p;
                    expected += survival;
                    const double rate = expected / costs[k];
                    if (rate > best_rate) {
                        best_rate     = rate;
                        expected_keep = k + 1;
                    }
                    prev = winner;
                    ++checks;
                }
                if (kept != expected_keep) {
                    fail("CUDA D-cut selector oracle");
                }
                ++checks;
            }
        }
        float   invalid[4] = { 1, 2, 3, 0 };
        int32_t got[4], kept;
        if (dspark_markov_cuda_dcut(ctx, base.data(), 0, 4, got, invalid, &kept)) {
            fail("D-cut accepts zero cost");
        }
        invalid[3] = std::numeric_limits<float>::quiet_NaN();
        if (dspark_markov_cuda_dcut(ctx, base.data(), 0, 4, got, invalid, &kept)) {
            fail("D-cut accepts NaN cost");
        }
        checks += 2;
        dspark_markov_cuda_free(ctx);
    }
    printf("CUDA_DCUT_ORACLE PASS checks=%d\n", checks);
    return 0;
}
#endif

int main(int argc, char ** argv) {
#ifdef LLAMA_DSPARK_MARKOV_CUDA
    if (argc == 2 && std::string(argv[1]) == "--cuda-dcut-self-test") {
        return test_cuda_dcut();
    }
#endif
#ifdef LLAMA_DSPARK_MARKOV_METAL
    if (argc == 2 && std::string(argv[1]) == "--metal-markov-self-test") {
        return test_metal_markov();
    }
#endif
    if (argc < 3) {
        fprintf(stderr, "usage: %s <tiny-dspark.gguf> <ref.json>\n", argv[0]);
        return 1;
    }
    const std::string model_path = argv[1];
    const std::string ref_path   = argv[2];

    std::ifstream f(ref_path);
    if (!f) {
        fail("could not open ref file: " + ref_path);
    }
    json ref;
    f >> ref;

    const int64_t n_embd_cap_ref    = ref.at("n_embd_cap").get<int64_t>();
    const int32_t block_size_ref    = ref.at("block_size").get<int32_t>();
    const int32_t vocab_size_ref    = ref.at("vocab_size").get<int32_t>();
    const int32_t mask_token_id_ref = ref.at("mask_token_id").get<int32_t>();
    const int32_t prefill_bonus_ref = ref.at("prefill_bonus").get<int32_t>();
    const auto &  rounds_ref        = ref.at("rounds");

    llama_backend_init();

    llama_model_params mparams = llama_model_default_params();
    mparams.n_gpu_layers       = 0;  // CPU: deterministic, no Metal precision surprises

    llama_model * model = llama_model_load_from_file(model_path.c_str(), mparams);
    if (!model) {
        fail("failed to load model: " + model_path);
    }

    llama_dspark_meta meta;
    if (!llama_model_dspark_get_meta(model, &meta)) {
        fail("llama_model_dspark_get_meta failed -- not a dspark model?");
    }

    printf(
        "meta: n_embd=%lld n_vocab=%lld n_capture=%lld n_embd_cap=%lld block_size=%d mask_token_id=%d "
        "markov_rank=%lld\n",
        (long long) meta.n_embd, (long long) meta.n_vocab, (long long) meta.n_capture, (long long) meta.n_embd_cap,
        meta.block_size, meta.mask_token_id, (long long) meta.markov_rank);

    if (meta.n_embd_cap != n_embd_cap_ref) {
        fail("n_embd_cap mismatch vs ref.json");
    }
    if (meta.block_size != block_size_ref) {
        fail("block_size mismatch vs ref.json");
    }
    if (meta.n_vocab != vocab_size_ref) {
        fail("vocab_size mismatch vs ref.json");
    }
    if (meta.mask_token_id != mask_token_id_ref) {
        fail("mask_token_id mismatch vs ref.json");
    }

    const int64_t n_embd_cap    = meta.n_embd_cap;
    const int32_t block_size    = meta.block_size;
    const int32_t vocab_size    = (int32_t) meta.n_vocab;
    const int32_t mask_token_id = meta.mask_token_id;

    // size the context generously: worst-case round needs ctx_len(<= previous
    // block_size+1) + block_size tokens in one llama_decode call.
    const uint32_t n_ctx_max = (uint32_t) (block_size + 1 + block_size) + 8;

    llama_context_params cparams = llama_context_default_params();
    cparams.n_ctx                = n_ctx_max;
    cparams.n_batch              = n_ctx_max;
    cparams.n_ubatch             = n_ctx_max;
    cparams.n_seq_max            = 1;
    cparams.no_perf              = true;

    llama_context * ctx = llama_init_from_model(model, cparams);
    if (!ctx) {
        fail("llama_init_from_model failed");
    }

    common_params_speculative sparams;
    sparams.types         = { COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK };
    sparams.draft.ctx_dft = ctx;
    // No real target model exists in this synthetic test (see file header
    // comment); dspark's process() path (which reads from ctx_tgt) is never
    // exercised here -- context rows are injected directly via
    // common_speculative_dspark_stage_ctx_test(). ctx_tgt only needs to be a
    // valid, non-null context to satisfy the impl's construction-time assert.
    sparams.draft.ctx_tgt = ctx;
    sparams.draft.n_max   = block_size;
    sparams.draft.n_min   = 0;

    common_speculative * spec = common_speculative_init(sparams, /* n_seq = */ 1);
    if (!spec) {
        fail("common_speculative_init returned null");
    }

    common_speculative_begin(spec, /* seq_id = */ 0, PROMPT);

    // one-time prefill seeding: the whole prompt's (synthetic) tap features,
    // positions [0, N).
    {
        const int64_t        N    = (int64_t) PROMPT.size();
        std::vector<float>   feat = synth_feat_rows(0, N, n_embd_cap);
        std::vector<int32_t> pos(N);
        for (int64_t i = 0; i < N; ++i) {
            pos[i] = (int32_t) i;
        }

        if (!common_speculative_dspark_stage_ctx_test(spec, 0, feat.data(), N, n_embd_cap, pos.data())) {
            fail("common_speculative_dspark_stage_ctx_test (prefill) failed");
        }
    }

    llama_pos   start   = (llama_pos) PROMPT.size();
    llama_token id_last = (llama_token) synth_bonus_token(-1, vocab_size, mask_token_id);
    if (id_last != prefill_bonus_ref) {
        fail("C++/python prefill bonus token disagree -- synth_bonus_token drifted");
    }

    int32_t n_mismatch_rounds = 0;

    for (size_t r = 0; r < rounds_ref.size(); ++r) {
        const auto &               rr          = rounds_ref[r];
        const int32_t              n_accepted  = rr.at("n_accepted").get<int32_t>();
        const int64_t              ctx_len_ref = rr.at("ctx_len").get<int64_t>();
        const int64_t              start_ref   = rr.at("start").get<int64_t>();
        const std::vector<int32_t> sampled_ref = rr.at("sampled").get<std::vector<int32_t>>();

        if ((int32_t) ACCEPT_SCHEDULE[r % ACCEPT_SCHEDULE.size()] != n_accepted) {
            fail("ACCEPT_SCHEDULE drifted out of sync with ref.json at round " + std::to_string(r));
        }
        if ((int64_t) start != start_ref) {
            fail("start bookkeeping disagrees with ref.json at round " + std::to_string(r) +
                 " (cpp=" + std::to_string(start) + " py=" + std::to_string(start_ref) + ")");
        }

        common_speculative_draft_params & dp = common_speculative_get_draft_params(spec, 0);
        dp.drafting                          = true;
        dp.n_max                             = -1;
        dp.n_past                            = start;
        dp.id_last                           = id_last;
        dp.prompt                            = nullptr;  // unused by dspark
        llama_tokens result;
        dp.result = &result;

        common_speculative_draft(spec);

        printf("round %zu: n_past=%d id_last=%d -> result=[", r, start, id_last);
        for (auto t : result) {
            printf("%d ", t);
        }
        printf("] expected=[");
        for (auto t : sampled_ref) {
            printf("%d ", t);
        }
        printf("]\n");

        if (result.size() != (size_t) block_size) {
            fail("round " + std::to_string(r) + ": expected block_size=" + std::to_string(block_size) +
                 " drafted tokens, got " + std::to_string(result.size()));
        }

        bool round_ok = true;
        for (int32_t k = 0; k < block_size; ++k) {
            if (result[k] != sampled_ref[k]) {
                round_ok = false;
            }
        }
        if (!round_ok) {
            n_mismatch_rounds++;
            fprintf(stderr, "  MISMATCH at round %zu\n", r);
        }

        // stage this round's verify-capture: the target would verify the
        // anchor + all block_size drafted tokens in one batch (verify_length
        // = block_size + 1), regardless of how many end up accepted -- accept()
        // below trims this down to the actually-committed prefix.
        {
            std::vector<float>   feat = synth_feat_rows(start, block_size + 1, n_embd_cap);
            std::vector<int32_t> pos(block_size + 1);
            for (int32_t i = 0; i < block_size + 1; ++i) {
                pos[i] = start + i;
            }

            if (!common_speculative_dspark_stage_ctx_test(spec, 0, feat.data(), block_size + 1, n_embd_cap,
                                                          pos.data())) {
                fail("common_speculative_dspark_stage_ctx_test (verify) failed at round " + std::to_string(r));
            }
        }

        common_speculative_accept(spec, 0, (uint16_t) n_accepted);

        const int32_t bonus     = synth_bonus_token((int32_t) r, vocab_size, mask_token_id);
        const int32_t bonus_ref = rr.at("bonus").get<int32_t>();
        if (bonus != bonus_ref) {
            fail("C++/python bonus token disagree at round " + std::to_string(r));
        }

        id_last = (llama_token) bonus;
        start   = start + n_accepted + 1;

        GGML_UNUSED(ctx_len_ref);
    }

    common_speculative_print_stats(spec);
    common_speculative_free(spec);
    llama_free(ctx);
    llama_model_free(model);
    llama_backend_free();

    if (n_mismatch_rounds > 0) {
        fail(std::to_string(n_mismatch_rounds) + "/" + std::to_string(rounds_ref.size()) +
             " rounds mismatched the Python reference");
    }

    printf(
        "\nPhase 2 gate PASSED: %zu/%zu rounds token-for-token identical to the DeepSpec reference "
        "(cache growth/crop, block seeding, RoPE positions, sequential markov resample).\n",
        rounds_ref.size(), rounds_ref.size());
    return 0;
}