/* qxmx_engine.cpp: base-class sampling + RNG plumbing for engine_i.
 *
 * M2.2: pure-host sampler over `logits` (already synced by forward/prefill).
 * Greedy when temperature==0; otherwise temperature scaling + top-k + top-p +
 * seeded std::mt19937_64.  Implemented once here so both backends share it.
 *
 * Algorithm:
 *   - greedy: single argmax pass (O(VOCAB), ~0.5 ms on 248K).
 *   - temp>0: scale logits by 1/temperature, then top-k (partial sort down to
 *     k candidates), then top-p (cumulative-sum filter on softmax of the k),
 *     then multinomial draw.  When both filters are disabled this degenerates
 *     to a full-vocab softmax + multinomial — the ~1-2 ms the ROADMAP budgets.
 *
 * Scratch (idx_buf, prob_buf) is hoisted on engine_i; the only per-sample
 * allocation is the partial_sort lambda's swaps, all in-place.
 */
#include "qxmx_engine.h"

#include <algorithm>
#include <cmath>
#include <cstdint>

namespace qx {

void engine_i::seed(int slot, uint64_t s) {
    if (slot < 0 || slot >= (int)rngs.size()) return;
    rngs[slot].seed(s);
}

int engine_i::open_slot() {
    int slot = (int)rngs.size();
    rngs.emplace_back();          /* default-constructed mt19937_64 (seed 0) */
    on_open_slot(slot);
    return slot;
}

void engine_i::close_slot(int slot) {
    if (slot < 0 || slot >= (int)rngs.size()) return;
    on_close_slot(slot);
    if (slot == (int)rngs.size() - 1) rngs.pop_back();
    else rngs.erase(rngs.begin() + slot);
}

int engine_i::sample(int slot, const sample_params& p) {
    /* Greedy: single argmax pass.  No scratch, no RNG. */
    if (p.temperature == 0.0f) {
        int best = 0; float bv = -1e30f;
        for (int v = 0; v < QX_N_VOCAB; v++) if (logits[v] > bv) { bv = logits[v]; best = v; }
        return best;
    }

    /* Non-greedy: temperature scale + top-k + top-p + multinomial. */
    std::mt19937_64& rng = rngs.at(slot);
    const int V = QX_N_VOCAB;
    idx_buf.assign(V, 0);
    for (int i = 0; i < V; i++) idx_buf[i] = i;
    prob_buf.assign(V, 0.0f);

    /* Scale logits by 1/temperature into prob_buf. */
    const float inv_t = 1.0f / p.temperature;
    for (int i = 0; i < V; i++) prob_buf[i] = logits[i] * inv_t;

    /* Top-k: keep the k largest logits.  partial_sort is O(V log k) — far
     * cheaper than a full sort, and the ROADMAP's "~1-2 ms" budget holds. */
    int k = V;
    if (p.top_k > 0 && p.top_k < V) k = p.top_k;
    if (k < V) {
        std::partial_sort(idx_buf.begin(), idx_buf.begin() + k, idx_buf.end(),
                         [&](int a, int b){ return prob_buf[a] > prob_buf[b]; });
    } else {
        std::sort(idx_buf.begin(), idx_buf.end(),
                  [&](int a, int b){ return prob_buf[a] > prob_buf[b]; });
    }

    /* Softmax over the k surviving candidates (numerically stable: subtract
     * the max, which is idx_buf[0] after the sort above). */
    const float max_l = prob_buf[idx_buf[0]];
    double z = 0.0;
    for (int i = 0; i < k; i++) {
        float e = std::exp(prob_buf[idx_buf[i]] - max_l);
        prob_buf[idx_buf[i]] = e;
        z += e;
    }
    if (z <= 0.0) return idx_buf[0];   /* degenerate — return the argmax */

    /* Top-p: accumulate over the (already prob-sorted) candidates until the
     * cumulative probability reaches top_p; cut everything after. */
    int keep = k;
    if (p.top_p < 1.0f) {
        double c = 0.0;
        for (int i = 0; i < k; i++) {
            c += prob_buf[idx_buf[i]] / z;
            if (c >= p.top_p) { keep = i + 1; break; }
        }
    }

    /* Multinomial draw over the kept candidates. */
    double r = std::uniform_real_distribution<double>(0.0, 1.0)(rng);
    double c = 0.0;
    for (int i = 0; i < keep; i++) {
        c += prob_buf[idx_buf[i]] / z;
        if (r < c) return idx_buf[i];
    }
    return idx_buf[keep - 1];   /* FP roundoff fallback */
}

}  // namespace qx