// Latent ("soft") thinking for LFM3 on llama.cpp, with continuous batching.
//
// Every input (prompt, thinking, answer) is fed as an embedding vector, so all modes share one code path and the
// discrete mode is an exact control:
//   discrete  next input = E[sampled token]                                                  (control)
//   soft      next input = sum_j p_j E[j] over the top --topk-soft candidates                (Soft Thinking)
//   gumbel    next input = sum_j softmax((log p_j + g_j) / tau) E[j], g ~ Gumbel             (stochastic soft thinking)
//   moi       next input = sum_j w_j E[j], w = (H p + (beta + 1 - H) onehot(sampled)) / (beta + 1),
//             H = entropy / log(n_vocab)                                                     (Mixture of Inputs)
//   mti       discrete sampling, but where the entropy of the next-token distribution exceeds --tau-mti the logits are
//             replaced by classifier-free guidance against a negative branch: the context continued by a short
//             negative cue (--neg, " OUTPUT ERROR") in a scratch copy of the sequence; log p = (1 - w) log p_neg +
//             w log p, w = --omega (Minimal Test-Time Intervention, ACL 2026). Applied while thinking and answering.
// Thinking ends when </think> is the chosen token (sampled for discrete/moi/mti, argmax of the weights for soft/gumbel),
// at --budget (forced), or for soft/gumbel when the entropy stays below --cold-tau for --cold-len steps ("cold stop").
// The answer is then sampled discretely in every mode.
// Sampling chain as in our llama-server runs: top_k 50 -> top_p 0.95 (temperature 1) -> temperature 0.6.
//
// --prompts: one line per prompt:  id <TAB> token ids separated by spaces
// --out:     id <TAB> think_steps <TAB> forced <TAB> cold <TAB> thinking ids <TAB> answer ids   (appended)
// --selftest compares the token path with the embedding path, alone and batched.

#include "ggml-backend.h"
#include "llama.h"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <fstream>
#include <random>
#include <set>
#include <sstream>
#include <string>
#include <vector>

static const llama_token THINK_END = 124902;
static const llama_token EOS_TOK   = 124900;

enum Mode { DISCRETE, SOFT, GUMBEL, MOI, MTI };

struct Args {
    std::string model, embd, prompts, out, mode = "discrete";
    int   budget = 6144, answer = 768, par = 64, top_k = 50, topk_soft = 15, cold_len = 256, ngl = 99, limit = 0;
    float top_p = 0.95f, temp = 0.6f, tau = 0.5f, beta = 1.0f, cold_tau = 0.01f, tau_mti = 1.0f, omega = 1.5f;
    int   kvu = -1;  // unified KV buffer: -1 = only for mti (cheap scratch copies), 0 / 1 = force
    std::vector<llama_token> neg;
    unsigned long long seed = 0;
    bool  selftest = false;
    bool  tokens   = false;  // discrete mode only: feed token ids instead of their embeddings (speed check)
};

struct Prompt {
    std::string              id;
    std::vector<llama_token> toks;
};

struct Seq {
    int  p     = -1;  // prompt index, -1 = idle
    int  pos   = 0;
    int  phase = 0;   // 0 thinking, 1 answer, 2 finished
    bool prefill = false;
    int  think_steps = 0, low_run = 0;
    bool forced = false, cold = false;
    std::vector<llama_token> think, answer;
    std::vector<float>       next;  // next input embedding
    llama_token              next_tok = -1;  // the same input as a token id, when it is one token
    std::mt19937_64          rng;
    int  out = -1;                  // index of its logits in the batch
    // MTI: guidance at this step, index of the negative branch's logits, conditional top-k ids and log-probabilities
    bool cfg = false;
    int  out2 = -1;
    std::vector<int>   cid;
    std::vector<float> clp;
};

static int                n_embd  = 0;
static int                n_vocab = 0;
static std::vector<float> E;  // n_vocab x n_embd

static const float * row_ptr(llama_token t) { return &E[(size_t) t * n_embd]; }

static void set_row(llama_token t, std::vector<float> & out) { out.assign(row_ptr(t), row_ptr(t) + n_embd); }

static void mix(const std::vector<int> & ids, const std::vector<float> & w, size_t k, std::vector<float> & out) {
    out.assign(n_embd, 0.0f);
    for (size_t j = 0; j < k; j++) {
        if (w[j] == 0.0f) continue;
        const float * e = row_ptr(ids[j]);
        for (int d = 0; d < n_embd; d++) out[d] += w[j] * e[d];
    }
}

// top_k -> top_p (temperature 1) -> temperature; ids sorted by probability, probabilities renormalised
static void filtered(const float * logits, const Args & a, std::vector<int> & ids, std::vector<float> & ps) {
    std::vector<std::pair<float, int>> h;
    h.reserve(a.top_k + 1);
    auto gt = [](const std::pair<float, int> & x, const std::pair<float, int> & y) { return x.first > y.first; };
    for (int i = 0; i < n_vocab; i++) {
        const float v = logits[i];
        if ((int) h.size() < a.top_k) {
            h.push_back({ v, i });
            std::push_heap(h.begin(), h.end(), gt);
        } else if (v > h.front().first) {
            std::pop_heap(h.begin(), h.end(), gt);
            h.back() = { v, i };
            std::push_heap(h.begin(), h.end(), gt);
        }
    }
    std::sort(h.begin(), h.end(), gt);
    const double        mx = h[0].first;
    std::vector<double> p1(h.size());
    double              s = 0;
    for (size_t i = 0; i < h.size(); i++) s += (p1[i] = std::exp(h[i].first - mx));
    size_t keep = 0;
    double cum  = 0;
    for (size_t i = 0; i < h.size(); i++) {
        keep = i + 1;
        cum += p1[i] / s;
        if (cum >= a.top_p) break;
    }
    ids.resize(keep);
    ps.resize(keep);
    double z = 0;
    for (size_t i = 0; i < keep; i++) {
        ids[i] = h[i].second;
        z += (p1[i] = std::exp((h[i].first - mx) / a.temp));
    }
    for (size_t i = 0; i < keep; i++) ps[i] = (float) (p1[i] / z);
}

static void step(Seq & s, const float * logits, const Args & a, Mode mode) {
    std::vector<int>   ids;
    std::vector<float> ps;
    filtered(logits, a, ids, ps);
    std::uniform_real_distribution<double> U(0.0, 1.0);
    auto sample = [&]() {
        const double u = U(s.rng);
        double       c = 0;
        for (size_t j = 0; j < ps.size(); j++) {
            c += ps[j];
            if (u < c) return ids[j];
        }
        return ids.back();
    };
    if (s.phase == 1) {  // answer: always discrete
        const llama_token t = sample();
        if (t == EOS_TOK || (int) s.answer.size() >= a.answer) {
            s.phase = 2;
            return;
        }
        s.answer.push_back(t);
        set_row(t, s.next);
        s.next_tok = t;
        return;
    }
    llama_token        chosen;
    std::vector<float> w;
    size_t             k = 0;
    if (mode == DISCRETE || mode == MTI) {
        chosen = sample();
    } else if (mode == MOI) {
        chosen   = sample();
        double H = 0;
        for (float p : ps)
            if (p > 0) H -= p * std::log((double) p);
        H /= std::log((double) n_vocab);
        k = ps.size();
        w.resize(k);
        for (size_t j = 0; j < k; j++)
            w[j] = (float) ((H * ps[j] + (ids[j] == chosen ? a.beta + 1.0 - H : 0.0)) / (a.beta + 1.0));
    } else {
        k = std::min((size_t) a.topk_soft, ps.size());
        w.assign(ps.begin(), ps.begin() + k);
        if (mode == GUMBEL) {
            std::vector<double> z(k);
            double              mz = -1e300;
            for (size_t j = 0; j < k; j++) {
                const double g = -std::log(-std::log(std::max(U(s.rng), 1e-300)));
                z[j]           = (std::log(std::max((double) w[j], 1e-30)) + g) / a.tau;
                mz             = std::max(mz, z[j]);
            }
            double sum = 0;
            for (size_t j = 0; j < k; j++) sum += (z[j] = std::exp(z[j] - mz));
            for (size_t j = 0; j < k; j++) w[j] = (float) (z[j] / sum);
        } else {
            double sum = 0;
            for (size_t j = 0; j < k; j++) sum += w[j];
            for (size_t j = 0; j < k; j++) w[j] = (float) (w[j] / sum);
        }
        chosen = ids[std::max_element(w.begin(), w.end()) - w.begin()];
        double H = 0;
        for (float p : ps)
            if (p > 0) H -= p * std::log((double) p);
        s.low_run = H < a.cold_tau ? s.low_run + 1 : 0;
    }
    bool end = chosen == THINK_END;
    if (!end && s.think_steps >= a.budget) end = s.forced = true;
    if (!end && (mode == SOFT || mode == GUMBEL) && s.low_run >= a.cold_len) end = s.cold = true;
    if (end) {
        s.phase = 1;
        set_row(THINK_END, s.next);
        s.next_tok = THINK_END;
        return;
    }
    s.think_steps++;
    s.think.push_back(chosen);
    if (mode == DISCRETE || mode == MTI) set_row(chosen, s.next), s.next_tok = chosen;
    else mix(ids, w, k, s.next), s.next_tok = -1;
}

// entropy (nats) of softmax(logits) over the whole vocabulary
static double entropy(const float * l) {
    double m = l[0];
    for (int i = 1; i < n_vocab; i++) m = std::max(m, (double) l[i]);
    double z = 0, s = 0;
    for (int i = 0; i < n_vocab; i++) {
        const double e = std::exp(l[i] - m);
        z += e;
        s += e * (l[i] - m);
    }
    return std::log(z) - s / z;
}

static double logsumexp(const float * l) {
    double m = l[0];
    for (int i = 1; i < n_vocab; i++) m = std::max(m, (double) l[i]);
    double z = 0;
    for (int i = 0; i < n_vocab; i++) z += std::exp(l[i] - m);
    return m + std::log(z);
}

// conditional top-k candidates and their log-probabilities (guidance is applied among the plausible candidates only)
static void top_logprobs(const float * l, int k, std::vector<int> & ids, std::vector<float> & lp) {
    std::vector<std::pair<float, int>> h;
    auto gt = [](const std::pair<float, int> & x, const std::pair<float, int> & y) { return x.first > y.first; };
    for (int i = 0; i < n_vocab; i++) {
        if ((int) h.size() < k) {
            h.push_back({ l[i], i });
            std::push_heap(h.begin(), h.end(), gt);
        } else if (l[i] > h.front().first) {
            std::pop_heap(h.begin(), h.end(), gt);
            h.back() = { l[i], i };
            std::push_heap(h.begin(), h.end(), gt);
        }
    }
    const double lse = logsumexp(l);
    ids.resize(h.size());
    lp.resize(h.size());
    for (size_t j = 0; j < h.size(); j++) ids[j] = h[j].second, lp[j] = (float) (h[j].first - lse);
}

static unsigned long long fnv(const std::string & s) {
    unsigned long long h = 1469598103934665603ULL;
    for (unsigned char c : s) h = (h ^ c) * 1099511628211ULL;
    return h;
}

static void add_embd(llama_batch & b, const float * e, llama_pos pos, llama_seq_id sq, bool out) {
    const int i = b.n_tokens++;
    std::memcpy(b.embd + (size_t) i * n_embd, e, sizeof(float) * n_embd);
    b.pos[i]       = pos;
    b.n_seq_id[i]  = 1;
    b.seq_id[i][0] = sq;
    b.logits[i]    = out;
}

static Args parse(int argc, char ** argv) {
    Args a;
    for (int i = 1; i < argc; i++) {
        std::string k = argv[i];
        auto        v = [&]() { return std::string(argv[++i]); };
        if (k == "--model") a.model = v();
        else if (k == "--embd") a.embd = v();
        else if (k == "--prompts") a.prompts = v();
        else if (k == "--out") a.out = v();
        else if (k == "--mode") a.mode = v();
        else if (k == "--budget") a.budget = std::stoi(v());
        else if (k == "--answer") a.answer = std::stoi(v());
        else if (k == "--par") a.par = std::stoi(v());
        else if (k == "--topk-soft") a.topk_soft = std::stoi(v());
        else if (k == "--cold-len") a.cold_len = std::stoi(v());
        else if (k == "--cold-tau") a.cold_tau = std::stof(v());
        else if (k == "--tau") a.tau = std::stof(v());
        else if (k == "--beta") a.beta = std::stof(v());
        else if (k == "--tau-mti") a.tau_mti = std::stof(v());
        else if (k == "--omega") a.omega = std::stof(v());
        else if (k == "--kvu") a.kvu = std::stoi(v());
        else if (k == "--neg") {
            std::istringstream ss(v());
            int                t;
            while (ss >> t) a.neg.push_back(t);
        }
        else if (k == "--seed") a.seed = std::stoull(v());
        else if (k == "--ngl") a.ngl = std::stoi(v());
        else if (k == "--limit") a.limit = std::stoi(v());
        else if (k == "--selftest") a.selftest = true;
        else if (k == "--tokens") a.tokens = true;
        else {
            fprintf(stderr, "unknown argument %s\n", k.c_str());
            exit(1);
        }
    }
    return a;
}

static std::vector<Prompt> read_prompts(const std::string & path, int limit, const std::set<std::string> & skip) {
    std::vector<Prompt> P;
    std::ifstream       f(path);
    std::string         line;
    while (std::getline(f, line)) {
        const size_t tab = line.find('\t');
        if (tab == std::string::npos) continue;
        Prompt p;
        p.id = line.substr(0, tab);
        std::istringstream ss(line.substr(tab + 1));
        int                t;
        while (ss >> t) p.toks.push_back(t);
        if (limit && (int) P.size() + (int) skip.size() >= limit) break;
        if (!skip.count(p.id)) P.push_back(p);
    }
    return P;
}

static std::vector<float> logits_copy(llama_context * ctx, int i) {
    const float * l = llama_get_logits_ith(ctx, i);
    return std::vector<float>(l, l + n_vocab);
}

// KL(p_x || p_y) at temperature 1 (the logit noise that matters for sampling)
static float maxdiff(const std::vector<float> & x, const std::vector<float> & y) {
    const double mx = *std::max_element(x.begin(), x.end()), my = *std::max_element(y.begin(), y.end());
    double       zx = 0, zy = 0;
    for (size_t i = 0; i < x.size(); i++) zx += std::exp(x[i] - mx), zy += std::exp(y[i] - my);
    double kl = 0;
    for (size_t i = 0; i < x.size(); i++) {
        const double lpx = x[i] - mx - std::log(zx), lpy = y[i] - my - std::log(zy);
        kl += std::exp(lpx) * (lpx - lpy);
    }
    return (float) kl;
}

static int selftest(llama_context * ctx, const std::vector<Prompt> & P) {
    llama_memory_t mem = llama_get_memory(ctx);
    const auto &   t0  = P[0].toks;
    const auto &   t1  = P[1].toks;
    // A: token path, sequence 0 alone
    llama_batch bt = llama_batch_init((int) t0.size() + 8, 0, 1);
    bt.n_tokens    = 0;
    for (size_t j = 0; j < t0.size(); j++) {
        const int i = bt.n_tokens++;
        bt.token[i] = t0[j], bt.pos[i] = (llama_pos) j, bt.n_seq_id[i] = 1, bt.seq_id[i][0] = 0;
        bt.logits[i] = j + 1 == t0.size();
    }
    if (llama_decode(ctx, bt)) return fprintf(stderr, "decode A failed\n"), 1;
    auto A = logits_copy(ctx, bt.n_tokens - 1);
    // greedy continuation, token path
    std::vector<std::vector<float>> contA;
    std::vector<llama_token>        cont;
    llama_token                     g = (llama_token) (std::max_element(A.begin(), A.end()) - A.begin());
    for (int s = 0; s < 6; s++) {
        cont.push_back(g);
        bt.n_tokens = 1;
        bt.token[0] = g, bt.pos[0] = (llama_pos) (t0.size() + s), bt.n_seq_id[0] = 1, bt.seq_id[0][0] = 0, bt.logits[0] = 1;
        if (llama_decode(ctx, bt)) return fprintf(stderr, "decode A+ failed\n"), 1;
        contA.push_back(logits_copy(ctx, 0));
        g = (llama_token) (std::max_element(contA.back().begin(), contA.back().end()) - contA.back().begin());
    }
    llama_memory_clear(mem, true);
    // B: embedding path, sequence 0 alone
    llama_batch be = llama_batch_init((int) (t0.size() + t1.size()) + 8, n_embd, 1);
    be.n_tokens    = 0;
    for (size_t j = 0; j < t0.size(); j++) add_embd(be, row_ptr(t0[j]), (llama_pos) j, 0, j + 1 == t0.size());
    if (llama_decode(ctx, be)) return fprintf(stderr, "decode B failed\n"), 1;
    auto B = logits_copy(ctx, be.n_tokens - 1);
    llama_memory_clear(mem, true);
    // C: embedding path, sequences 0 and 1 in one batch, then 6 steps together
    be.n_tokens = 0;
    for (size_t j = 0; j < t1.size(); j++) add_embd(be, row_ptr(t1[j]), (llama_pos) j, 1, j + 1 == t1.size());
    for (size_t j = 0; j < t0.size(); j++) add_embd(be, row_ptr(t0[j]), (llama_pos) j, 0, j + 1 == t0.size());
    if (llama_decode(ctx, be)) return fprintf(stderr, "decode C failed\n"), 1;
    auto  C  = logits_copy(ctx, be.n_tokens - 1);
    float dc = 0;
    for (int s = 0; s < 6; s++) {
        be.n_tokens = 0;
        add_embd(be, row_ptr(cont[s]), (llama_pos) (t1.size() + s), 1, true);
        add_embd(be, row_ptr(cont[s]), (llama_pos) (t0.size() + s), 0, true);
        if (llama_decode(ctx, be)) return fprintf(stderr, "decode C+ failed\n"), 1;
        dc = std::max(dc, maxdiff(contA[s], logits_copy(ctx, 1)));
    }
    float scale = 0;
    for (float x : A) scale = std::max(scale, std::fabs(x));
    printf("KL divergence (nats), logit scale %.2f | prefill: token vs embedding %.4g, alone vs batched %.4g | 6 decode "
           "steps, token alone vs embedding batched (max) %.4g\n",
           scale, maxdiff(A, B), maxdiff(A, C), dc);
    llama_memory_clear(mem, true);
    return 0;
}

int main(int argc, char ** argv) {
    Args a = parse(argc, argv);
    const Mode mode = a.mode == "soft" ? SOFT : a.mode == "gumbel" ? GUMBEL : a.mode == "moi" ? MOI : a.mode == "mti" ? MTI
                                                                                                                  : DISCRETE;
    if (a.mode != "soft" && a.mode != "gumbel" && a.mode != "moi" && a.mode != "mti" && a.mode != "discrete") {
        fprintf(stderr, "unknown mode\n");
        return 1;
    }
    if (mode == MTI && a.neg.empty()) {
        fprintf(stderr, "mti needs --neg token ids\n");
        return 1;
    }
    std::set<std::string> done;
    if (!a.selftest && !a.out.empty()) {
        std::ifstream f(a.out);
        std::string   line;
        while (std::getline(f, line)) done.insert(line.substr(0, line.find('\t')));
    }
    std::vector<Prompt> P = read_prompts(a.prompts, a.limit, done);
    std::sort(P.begin(), P.end(), [](const Prompt & x, const Prompt & y) { return x.toks.size() > y.toks.size(); });
    if (P.empty()) {
        printf("nothing to do\n");
        return 0;
    }
    size_t maxlen = 0;
    for (auto & p : P) maxlen = std::max(maxlen, p.toks.size());

    llama_backend_init();
    ggml_backend_load_all();
    auto mp         = llama_model_default_params();
    mp.n_gpu_layers = a.ngl;
    llama_model * model = llama_model_load_from_file(a.model.c_str(), mp);
    if (!model) return fprintf(stderr, "model load failed\n"), 1;
    n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
    n_embd  = llama_model_n_embd(model);
    {
        std::ifstream f(a.embd, std::ios::binary | std::ios::ate);
        const size_t  bytes = (size_t) f.tellg();
        if (bytes < (size_t) n_vocab * n_embd * sizeof(float))
            return fprintf(stderr, "embedding file too small: %zu bytes for %d x %d\n", bytes, n_vocab, n_embd), 1;
        f.seekg(0);
        E.resize((size_t) n_vocab * n_embd);
        f.read((char *) E.data(), E.size() * sizeof(float));
    }
    const int n_batch  = 2048 + a.par;
    auto      cp       = llama_context_default_params();
    cp.n_seq_max       = mode == MTI ? 2 * a.par : a.par;  // MTI: scratch copy par + i of sequence i
    cp.n_ctx           = (uint32_t) (a.par * (a.budget + a.answer + maxlen + 16 + a.neg.size()));
    cp.n_batch         = std::max(n_batch, (int) maxlen + a.par);
    cp.n_ubatch        = 512;
    cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
    // a unified buffer makes every sequence attend over all cells (masked): slow for many independent sequences
    cp.kv_unified      = a.kvu < 0 ? mode == MTI : a.kvu == 1;
    cp.n_threads       = 8;
    cp.n_threads_batch = 8;
    cp.no_perf         = true;
    llama_context * ctx = llama_init_from_model(model, cp);
    if (!ctx) return fprintf(stderr, "context failed\n"), 1;
    printf("model n_vocab %d n_embd %d | %zu prompts, mode %s, par %d, n_ctx %u\n", n_vocab, n_embd, P.size(),
           a.mode.c_str(), a.par, cp.n_ctx);
    fflush(stdout);
    if (a.selftest) return selftest(ctx, P);

    llama_memory_t   mem = llama_get_memory(ctx);
    llama_batch      b   = llama_batch_init(cp.n_batch, n_embd, 1);
    llama_batch      b2  = llama_batch_init(a.par * (int) a.neg.size() + 1, n_embd, 1);
    llama_batch      bt  = llama_batch_init(cp.n_batch, 0, 1);
    std::vector<Seq> S(a.par);
    FILE *           fo = fopen(a.out.c_str(), "ab");
    size_t           next_p = 0, n_done = 0;
    long long        n_fed = 0, n_steps = 0, n_cfg = 0, n_seqsteps = 0;
    const auto       t0 = std::chrono::steady_clock::now();
    std::vector<const float *> L(a.par);
    double t_dec = 0, t_smp = 0;
    auto   now = [] { return std::chrono::steady_clock::now(); };
    auto   ms  = [](auto x, auto y) { return std::chrono::duration<double, std::milli>(y - x).count(); };
    while (true) {
        int room = 2048;  // prompt tokens admitted per step
        for (int i = 0; i < a.par && next_p < P.size(); i++) {
            if (S[i].p >= 0) continue;
            const int n = (int) P[next_p].toks.size();
            if (n > room && room < 2048) break;
            S[i]         = Seq();
            S[i].p       = (int) next_p;
            S[i].prefill = true;
            S[i].rng.seed(a.seed * 1000003ULL + fnv(P[next_p].id));
            room -= n;
            next_p++;
        }
        const bool use_tok = a.tokens && mode == DISCRETE;
        b.n_tokens = bt.n_tokens = 0;
        auto add_tok = [&](llama_token t, llama_pos pos, llama_seq_id sq, bool out) {
            const int k = bt.n_tokens++;
            bt.token[k] = t, bt.pos[k] = pos, bt.n_seq_id[k] = 1, bt.seq_id[k][0] = sq, bt.logits[k] = out;
        };
        for (int i = 0; i < a.par; i++) {
            Seq & s = S[i];
            s.out   = -1;
            if (s.p < 0) continue;
            if (s.prefill) {
                const auto & t = P[s.p].toks;
                for (size_t j = 0; j < t.size(); j++)
                    if (use_tok) add_tok(t[j], (llama_pos) j, i, j + 1 == t.size());
                    else add_embd(b, row_ptr(t[j]), (llama_pos) j, i, j + 1 == t.size());
                s.pos     = (int) t.size();
                s.prefill = false;
            } else if (use_tok) {
                add_tok(s.next_tok, s.pos++, i, true);
            } else {
                add_embd(b, s.next.data(), s.pos++, i, true);
            }
            s.out = (use_tok ? bt.n_tokens : b.n_tokens) - 1;
        }
        if (use_tok) b.n_tokens = bt.n_tokens;
        if (b.n_tokens == 0) break;
        n_fed += b.n_tokens;
        for (int i = 0; i < a.par; i++) n_seqsteps += S[i].out >= 0;
        auto      tA = now();
        const int rc = llama_decode(ctx, use_tok ? bt : b);
        if (rc != 0) return fprintf(stderr, "llama_decode failed: %d\n", rc), 1;
        for (int i = 0; i < a.par; i++) L[i] = S[i].out >= 0 ? llama_get_logits_ith(ctx, S[i].out) : nullptr;
        auto tB = now();
        t_dec += ms(tA, tB);
        if (mode == MTI) {
#pragma omp parallel for schedule(dynamic)
            for (int i = 0; i < a.par; i++) {
                S[i].cfg = L[i] && entropy(L[i]) > a.tau_mti;
                if (S[i].cfg) top_logprobs(L[i], a.top_k, S[i].cid, S[i].clp);
                else if (L[i]) step(S[i], L[i], a, mode);
            }
            b2.n_tokens = 0;
            for (int i = 0; i < a.par; i++) {
                if (!S[i].cfg) continue;
                llama_memory_seq_cp(mem, i, a.par + i, -1, -1);
                for (size_t j = 0; j < a.neg.size(); j++)
                    add_embd(b2, row_ptr(a.neg[j]), (llama_pos) (S[i].pos + j), a.par + i, j + 1 == a.neg.size());
                S[i].out2 = b2.n_tokens - 1;
                n_cfg++;
            }
            if (b2.n_tokens) {
                const int rc2 = llama_decode(ctx, b2);
                if (rc2 != 0) return fprintf(stderr, "llama_decode (negative branch) failed: %d\n", rc2), 1;
                for (int i = 0; i < a.par; i++) L[i] = S[i].cfg ? llama_get_logits_ith(ctx, S[i].out2) : nullptr;
#pragma omp parallel for schedule(dynamic)
                for (int i = 0; i < a.par; i++) {
                    if (!S[i].cfg) continue;
                    const double       lse = logsumexp(L[i]);
                    std::vector<float> g(n_vocab, -1e30f);
                    for (size_t j = 0; j < S[i].cid.size(); j++) {
                        const int t = S[i].cid[j];
                        g[t] = (float) ((1.0 - a.omega) * (L[i][t] - lse) + a.omega * S[i].clp[j]);
                    }
                    step(S[i], g.data(), a, mode);
                }
                for (int i = 0; i < a.par; i++)
                    if (S[i].cfg) llama_memory_seq_rm(mem, a.par + i, -1, -1);
            }
        } else {
#pragma omp parallel for schedule(dynamic)
            for (int i = 0; i < a.par; i++)
                if (L[i]) step(S[i], L[i], a, mode);
        }
        t_smp += ms(tB, now());
        for (int i = 0; i < a.par; i++) {
            Seq & s = S[i];
            if (s.p < 0 || s.phase != 2) continue;
            fprintf(fo, "%s\t%d\t%d\t%d\t", P[s.p].id.c_str(), s.think_steps, (int) s.forced, (int) s.cold);
            for (size_t j = 0; j < s.think.size(); j++) fprintf(fo, j ? " %d" : "%d", s.think[j]);
            fprintf(fo, "\t");
            for (size_t j = 0; j < s.answer.size(); j++) fprintf(fo, j ? " %d" : "%d", s.answer[j]);
            fprintf(fo, "\n");
            fflush(fo);
            llama_memory_seq_rm(mem, i, -1, -1);
            s.p = -1;
            n_done++;
        }
        if (++n_steps % 200 == 0) {
            const double el = std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count();
            fprintf(stderr, "step %lld: %zu/%zu done  %.0f tok/s  %.1f min  guidance at %.1f%% of steps  (decode %.1f ms, "
                    "sampling %.1f ms per step)\n", n_steps, n_done, P.size(), n_fed / el, el / 60,
                    100.0 * n_cfg / std::max(1LL, n_seqsteps), t_dec / n_steps, t_smp / n_steps);
        }
    }
    const double el = std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count();
    printf("%zu/%zu done  %.0f tok/s  %.1f min  guidance at %.1f%% of steps\n", n_done, P.size(), n_fed / el, el / 60,
           100.0 * n_cfg / std::max(1LL, n_seqsteps));
    fclose(fo);
    llama_batch_free(b);
    llama_batch_free(b2);
    llama_free(ctx);
    llama_model_free(model);
    return 0;
}