#include "cp_inference.h"
#include <ggml.h>
#include <gguf.h>
#include <ggml-cpu.h>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>

// Hyperparams CP (Qwen3-TTS Code Predictor) — fixés par l'architecture du modèle.
static const int N_EMBD   = 1024;
static const int N_LAYER  = CP_N_LAYER;  // 5, déclaré dans le header pour la taille KV cache
static const int N_HEAD   = 16;
static const int N_KV     = 8;
static const int HEAD_DIM = 128;
static const int N_VOCAB  = 2048;
static const int N_CB     = 15;
static const float ROPE_BASE = 1000000.0f;
static const float RMS_EPS   = 1e-6f;
// Positions max dans une frame CP : 2 (prefill: hidden + cb0_emb) + 14 (decode CB1..CB14
// nourris pour prédire CB2..CB15) = 16. Borne dure utilisée pour le KV cache.
static const int T_MAX    = 16;

static std::vector<float> read_floats(const char* path, size_t expect) {
    FILE* f = fopen(path, "rb");
    if (!f) { fprintf(stderr, "cp_load: open fail %s\n", path); return {}; }
    std::vector<float> v(expect);
    if (fread(v.data(), sizeof(float), expect, f) != expect) {
        fprintf(stderr, "cp_load: short read %s\n", path); fclose(f); return {};
    }
    fclose(f); return v;
}

bool cp_load(CPState& s, const char* gguf_path, const char* heads_path, const char* embs_path) {
    ggml_context* meta = nullptr;
    gguf_init_params p; p.no_alloc = true; p.ctx = &meta;
    gguf_context* g = gguf_init_from_file(gguf_path, p);
    if (!g) { fprintf(stderr, "cp_load: gguf open fail %s\n", gguf_path); return false; }
    int64_t n = gguf_get_n_tensors(g);
    size_t bytes = 0;
    for (int64_t i = 0; i < n; i++) {
        ggml_tensor* mt = ggml_get_tensor(meta, gguf_get_tensor_name(g, i));
        bytes += ggml_nelements(mt) * sizeof(float);
    }
    ggml_init_params ip = { bytes + (size_t)n * ggml_tensor_overhead() + (1u << 20), nullptr, false };
    s.weights_ctx = ggml_init(ip);

    FILE* f = fopen(gguf_path, "rb");
    const size_t off = gguf_get_data_offset(g);
    std::vector<uint8_t> tmp;
    for (int64_t i = 0; i < n; i++) {
        const char* name = gguf_get_tensor_name(g, i);
        ggml_tensor* mt = ggml_get_tensor(meta, name);
        ggml_tensor* t32 = ggml_new_tensor(s.weights_ctx, GGML_TYPE_F32, ggml_n_dims(mt), mt->ne);
        size_t nb = ggml_nbytes(mt);
        int64_t ne = ggml_nelements(mt);
        tmp.resize(nb);
        fseek(f, off + gguf_get_tensor_offset(g, i), SEEK_SET);
        if (fread(tmp.data(), 1, nb, f) != nb) { fprintf(stderr, "cp_load read fail %s\n", name); fclose(f); return false; }
        if (mt->type == GGML_TYPE_F32)      memcpy(t32->data, tmp.data(), nb);
        else if (mt->type == GGML_TYPE_F16) ggml_fp16_to_fp32_row((const ggml_fp16_t*)tmp.data(), (float*)t32->data, ne);
        else { fprintf(stderr, "cp_load unsupported type %s for %s\n", ggml_type_name(mt->type), name); fclose(f); return false; }
        s.tensors[name] = t32;
    }
    fclose(f);
    gguf_free(g);

    const size_t TAB = (size_t)N_CB * N_VOCAB * N_EMBD;
    s.heads      = read_floats(heads_path, TAB); if (s.heads.empty()) return false;
    s.codec_embs = read_floats(embs_path,  TAB); if (s.codec_embs.empty()) return false;
    fprintf(stderr, "cp_load: %lld tensors + heads(%.0f MB) + codec_embs(%.0f MB) OK\n",
            (long long)n, TAB * 4 / 1048576.0, TAB * 4 / 1048576.0);
    return true;
}

void cp_free(CPState& s) {
    if (s.weights_ctx) ggml_free(s.weights_ctx);
    s.weights_ctx = nullptr;
    s.tensors.clear();
    s.heads.clear();
    s.codec_embs.clear();
    if (s.cache_ctx) ggml_free(s.cache_ctx);
    s.cache_ctx = nullptr;
    for (int i = 0; i < N_LAYER; i++) { s.K_cache[i] = nullptr; s.V_cache[i] = nullptr; }
}

// Alloue le KV cache (paresseux, au 1er appel cp_predict_cached). Tensors persistants entre
// étapes/frames (le data pointer reste stable), c'est ce qui permet à un graph par-étape de
// lire/écrire dans le même buffer via ggml_view_3d + ggml_cpy.
static bool cp_cache_init(CPState& s) {
    if (s.cache_ctx) return true;
    const size_t per_tensor_bytes = (size_t)HEAD_DIM * N_KV * T_MAX * sizeof(float);
    const size_t total = (size_t)N_LAYER * 2 * (per_tensor_bytes + ggml_tensor_overhead()) + (1u << 16);
    ggml_init_params p = { total, nullptr, false };
    s.cache_ctx = ggml_init(p);
    if (!s.cache_ctx) { fprintf(stderr, "cp_cache_init: ggml_init failed\n"); return false; }
    for (int L_i = 0; L_i < N_LAYER; L_i++) {
        s.K_cache[L_i] = ggml_new_tensor_3d(s.cache_ctx, GGML_TYPE_F32, HEAD_DIM, N_KV, T_MAX);
        s.V_cache[L_i] = ggml_new_tensor_3d(s.cache_ctx, GGML_TYPE_F32, HEAD_DIM, N_KV, T_MAX);
        // Pas besoin de zéro initial : les positions utilisées sont écrites avant lecture
        // (cpy ops topo-ordonnés avant les attention reads, cf cp_forward_cached_step).
    }
    return true;
}

// Forward un transformer 5L sur X[N_EMBD, L] -> out_hn = hidden après output_norm à la position `last`.
// Architecture identique à cp_runner.cpp (la référence bit-exact validée).
static void cp_forward_lastpos(CPState& s, const std::vector<float>& embeds_flat, int L, int last,
                               std::vector<float>& out_hn) {
    size_t mem = 128ULL * 1024 * 1024;
    ggml_init_params p = { mem, nullptr, false };
    ggml_context* ctx = ggml_init(p);

    ggml_tensor* x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, N_EMBD, L);
    memcpy(x->data, embeds_flat.data(), (size_t)N_EMBD * L * sizeof(float));
    ggml_tensor* pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, L);
    for (int i = 0; i < L; i++) ((int32_t*)pos->data)[i] = i;
    ggml_tensor* mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, L, L);
    {
        float* m = (float*)mask->data;
        for (int q = 0; q < L; q++)
            for (int k = 0; k < L; k++)
                m[k + L * q] = (k > q) ? -INFINITY : 0.0f;
    }
    const float scale = 1.0f / sqrtf((float)HEAD_DIM);

    auto W = [&](const std::string& k) -> ggml_tensor* {
        auto it = s.tensors.find(k);
        if (it == s.tensors.end()) { fprintf(stderr, "CP missing tensor %s\n", k.c_str()); abort(); }
        return it->second;
    };

    for (int L_i = 0; L_i < N_LAYER; L_i++) {
        char pre[32]; snprintf(pre, sizeof(pre), "blk.%d", L_i);
        std::string b = pre;
        ggml_tensor* res = x;

        ggml_tensor* xn = ggml_rms_norm(ctx, x, RMS_EPS);
        xn = ggml_mul(ctx, xn, W(b + ".attn_norm.weight"));

        ggml_tensor* Q = ggml_mul_mat(ctx, W(b + ".attn_q.weight"), xn);
        ggml_tensor* K = ggml_mul_mat(ctx, W(b + ".attn_k.weight"), xn);
        ggml_tensor* V = ggml_mul_mat(ctx, W(b + ".attn_v.weight"), xn);

        Q = ggml_reshape_3d(ctx, Q, HEAD_DIM, N_HEAD, L);
        K = ggml_reshape_3d(ctx, K, HEAD_DIM, N_KV,   L);
        V = ggml_reshape_3d(ctx, V, HEAD_DIM, N_KV,   L);

        // q/k-norm AVANT RoPE (gotcha)
        Q = ggml_mul(ctx, ggml_rms_norm(ctx, Q, RMS_EPS), W(b + ".attn_q_norm.weight"));
        K = ggml_mul(ctx, ggml_rms_norm(ctx, K, RMS_EPS), W(b + ".attn_k_norm.weight"));

        Q = ggml_rope_ext(ctx, Q, pos, NULL, HEAD_DIM, GGML_ROPE_TYPE_NEOX, 0,
                          ROPE_BASE, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
        K = ggml_rope_ext(ctx, K, pos, NULL, HEAD_DIM, GGML_ROPE_TYPE_NEOX, 0,
                          ROPE_BASE, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);

        Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3));
        K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3));
        V = ggml_cont(ctx, ggml_permute(ctx, V, 1, 2, 0, 3));

        ggml_tensor* KQ = ggml_mul_mat(ctx, K, Q);
        KQ = ggml_soft_max_ext(ctx, KQ, mask, scale, 0.0f);
        ggml_tensor* KQV = ggml_mul_mat(ctx, V, KQ);
        KQV = ggml_cont(ctx, ggml_permute(ctx, KQV, 0, 2, 1, 3));
        KQV = ggml_reshape_2d(ctx, KQV, HEAD_DIM * N_HEAD, L);

        ggml_tensor* attn = ggml_mul_mat(ctx, W(b + ".attn_output.weight"), KQV);
        x = ggml_add(ctx, res, attn);

        res = x;
        ggml_tensor* xn2 = ggml_rms_norm(ctx, x, RMS_EPS);
        xn2 = ggml_mul(ctx, xn2, W(b + ".ffn_norm.weight"));
        ggml_tensor* g = ggml_silu(ctx, ggml_mul_mat(ctx, W(b + ".ffn_gate.weight"), xn2));
        ggml_tensor* u = ggml_mul_mat(ctx, W(b + ".ffn_up.weight"), xn2);
        ggml_tensor* ff = ggml_mul_mat(ctx, W(b + ".ffn_down.weight"), ggml_mul(ctx, g, u));
        x = ggml_add(ctx, res, ff);
    }
    x = ggml_rms_norm(ctx, x, RMS_EPS);
    x = ggml_mul(ctx, x, W("output_norm.weight"));

    ggml_cgraph* gf = ggml_new_graph_custom(ctx, 8192, false);
    ggml_build_forward_expand(gf, x);
    ggml_graph_compute_with_ctx(ctx, gf, s.n_threads);

    out_hn.resize(N_EMBD);
    memcpy(out_hn.data(), (float*)x->data + (size_t)N_EMBD * last, N_EMBD * sizeof(float));
    ggml_free(ctx);
}

// Forward UNE étape avec KV cache. Input : embeds_new[n_new, N_EMBD] f32, positions
// [pos_off .. pos_off+n_new). Lecture cache [0..pos_off), écriture cache [pos_off..pos_off+n_new).
// Sortie : out_hn = hidden state (après output_norm) à la dernière position (pos_off+n_new-1).
//
// Gotcha topo : K_cpy/V_cpy n'apparaissent PAS dans le DAG de l'output x (les vues
// K_full/V_full ne dépendent que du leaf K_cache/V_cache, pas du cpy). On les pousse en tête
// via build_forward_expand AVANT l'expand(x) pour qu'ils soient exécutés en premier dans
// l'ordre topologique du graph. Sans ça, l'attention lirait des positions non-écrites.
static void cp_forward_cached_step(CPState& s, const float* embeds_new, int n_new,
                                    int pos_off, std::vector<float>& out_hn) {
    const int T_full = pos_off + n_new;
    // Buffer de calcul : on rebuild un graph par étape, mais chacun ne touche que n_new tokens
    // sur 5 couches. ~5-10 MB par graph en pratique ; on laisse de la marge.
    size_t mem = 32ULL * 1024 * 1024;
    ggml_init_params p = { mem, nullptr, false };
    ggml_context* ctx = ggml_init(p);

    ggml_tensor* x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, N_EMBD, n_new);
    memcpy(x->data, embeds_new, (size_t)N_EMBD * n_new * sizeof(float));
    ggml_tensor* pos_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_new);
    for (int i = 0; i < n_new; i++) ((int32_t*)pos_ids->data)[i] = pos_off + i;

    // Masque causal [T_full (n_k), n_new (n_q)] : la requête q (à position pos_off+q) voit
    // les clés 0..pos_off+q. KQ aura shape [T_full, n_new, N_HEAD] -> broadcast sur dim 2.
    ggml_tensor* mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, T_full, n_new);
    {
        float* m = (float*)mask->data;
        for (int q = 0; q < n_new; q++)
            for (int k = 0; k < T_full; k++)
                m[k + T_full * q] = (k > pos_off + q) ? -INFINITY : 0.0f;
    }
    const float scale = 1.0f / sqrtf((float)HEAD_DIM);

    auto W = [&](const std::string& key) -> ggml_tensor* {
        auto it = s.tensors.find(key);
        if (it == s.tensors.end()) { fprintf(stderr, "CP missing tensor %s\n", key.c_str()); abort(); }
        return it->second;
    };

    // On collecte tous les cpy ops dans l'ordre de couche pour les expand en tête du graph.
    std::vector<ggml_tensor*> cpy_ops; cpy_ops.reserve((size_t)N_LAYER * 2);

    for (int L_i = 0; L_i < N_LAYER; L_i++) {
        char pre[32]; snprintf(pre, sizeof(pre), "blk.%d", L_i);
        std::string b = pre;
        ggml_tensor* res = x;

        ggml_tensor* xn = ggml_rms_norm(ctx, x, RMS_EPS);
        xn = ggml_mul(ctx, xn, W(b + ".attn_norm.weight"));

        ggml_tensor* Q = ggml_mul_mat(ctx, W(b + ".attn_q.weight"), xn);
        ggml_tensor* K = ggml_mul_mat(ctx, W(b + ".attn_k.weight"), xn);
        ggml_tensor* V = ggml_mul_mat(ctx, W(b + ".attn_v.weight"), xn);

        Q = ggml_reshape_3d(ctx, Q, HEAD_DIM, N_HEAD, n_new);
        K = ggml_reshape_3d(ctx, K, HEAD_DIM, N_KV,   n_new);
        V = ggml_reshape_3d(ctx, V, HEAD_DIM, N_KV,   n_new);

        // q/k-norm AVANT RoPE (gotcha identique à cp_forward_lastpos).
        Q = ggml_mul(ctx, ggml_rms_norm(ctx, Q, RMS_EPS), W(b + ".attn_q_norm.weight"));
        K = ggml_mul(ctx, ggml_rms_norm(ctx, K, RMS_EPS), W(b + ".attn_k_norm.weight"));

        Q = ggml_rope_ext(ctx, Q, pos_ids, NULL, HEAD_DIM, GGML_ROPE_TYPE_NEOX, 0,
                          ROPE_BASE, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
        K = ggml_rope_ext(ctx, K, pos_ids, NULL, HEAD_DIM, GGML_ROPE_TYPE_NEOX, 0,
                          ROPE_BASE, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);

        // Écriture cache : K, V (post-RoPE, shape [HEAD_DIM, N_KV, n_new]) -> slot
        // [pos_off..pos_off+n_new) du cache (même layout, contigu, juste un offset sur dim 2).
        ggml_tensor* K_cache = s.K_cache[L_i];
        ggml_tensor* V_cache = s.V_cache[L_i];
        ggml_tensor* K_dst = ggml_view_3d(ctx, K_cache, HEAD_DIM, N_KV, n_new,
                                          K_cache->nb[1], K_cache->nb[2],
                                          (size_t)pos_off * K_cache->nb[2]);
        ggml_tensor* V_dst = ggml_view_3d(ctx, V_cache, HEAD_DIM, N_KV, n_new,
                                          V_cache->nb[1], V_cache->nb[2],
                                          (size_t)pos_off * V_cache->nb[2]);
        ggml_tensor* K_cpy = ggml_cpy(ctx, K, K_dst);
        ggml_tensor* V_cpy = ggml_cpy(ctx, V, V_dst);
        cpy_ops.push_back(K_cpy);
        cpy_ops.push_back(V_cpy);

        // Lecture cache : [0..T_full) pour l'attention. Vues à offset 0.
        ggml_tensor* K_full = ggml_view_3d(ctx, K_cache, HEAD_DIM, N_KV, T_full,
                                           K_cache->nb[1], K_cache->nb[2], 0);
        ggml_tensor* V_full = ggml_view_3d(ctx, V_cache, HEAD_DIM, N_KV, T_full,
                                           V_cache->nb[1], V_cache->nb[2], 0);

        Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3));
        ggml_tensor* K_p = ggml_cont(ctx, ggml_permute(ctx, K_full, 0, 2, 1, 3));
        ggml_tensor* V_p = ggml_cont(ctx, ggml_permute(ctx, V_full, 1, 2, 0, 3));

        ggml_tensor* KQ = ggml_mul_mat(ctx, K_p, Q);
        KQ = ggml_soft_max_ext(ctx, KQ, mask, scale, 0.0f);
        ggml_tensor* KQV = ggml_mul_mat(ctx, V_p, KQ);
        KQV = ggml_cont(ctx, ggml_permute(ctx, KQV, 0, 2, 1, 3));
        KQV = ggml_reshape_2d(ctx, KQV, HEAD_DIM * N_HEAD, n_new);

        ggml_tensor* attn = ggml_mul_mat(ctx, W(b + ".attn_output.weight"), KQV);
        x = ggml_add(ctx, res, attn);

        // FFN : pointwise par-token, n_new colonnes.
        res = x;
        ggml_tensor* xn2 = ggml_rms_norm(ctx, x, RMS_EPS);
        xn2 = ggml_mul(ctx, xn2, W(b + ".ffn_norm.weight"));
        ggml_tensor* g = ggml_silu(ctx, ggml_mul_mat(ctx, W(b + ".ffn_gate.weight"), xn2));
        ggml_tensor* u = ggml_mul_mat(ctx, W(b + ".ffn_up.weight"), xn2);
        ggml_tensor* ff = ggml_mul_mat(ctx, W(b + ".ffn_down.weight"), ggml_mul(ctx, g, u));
        x = ggml_add(ctx, res, ff);
    }
    x = ggml_rms_norm(ctx, x, RMS_EPS);
    x = ggml_mul(ctx, x, W("output_norm.weight"));

    ggml_cgraph* gf = ggml_new_graph_custom(ctx, 8192, false);
    // ORDRE CRITIQUE : pousser tous les cpy AVANT x. Comme K_cpy/V_cpy ne sont pas des
    // dépendances DAG de x (les vues K_full ne lisent qu'un leaf), l'expand normal ne les
    // mettrait jamais dans le graph. Et même si on les ajoutait après, leur position topo
    // pourrait être après les reads. Ici l'expand récursif pose les chaînes K_proj/RoPE/cpy
    // en premier, puis l'expand(x) ajoute le reste (Q, attention, FFN, …) après — les reads
    // K_full/V_full apparaissent donc topo-après les writes K_cpy/V_cpy correspondants.
    for (auto* c : cpy_ops) ggml_build_forward_expand(gf, c);
    ggml_build_forward_expand(gf, x);

    ggml_graph_compute_with_ctx(ctx, gf, s.n_threads);

    out_hn.resize(N_EMBD);
    memcpy(out_hn.data(), (float*)x->data + (size_t)N_EMBD * (n_new - 1),
           N_EMBD * sizeof(float));
    ggml_free(ctx);
}

void cp_predict(CPState& s, const float* hidden, const float* cb0_emb, int32_t* out_codes) {
    // embeds = [hidden, cb0_emb, codec_embs[0][cb1], codec_embs[1][cb2], ..., codec_embs[13][cb14]]
    // step s (1..15) : forward sur les (s+1) tokens, sortir hidden à la position `s`, head[s-1] -> sample.
    std::vector<float> embeds; embeds.reserve((size_t)17 * N_EMBD);
    embeds.insert(embeds.end(), hidden,  hidden  + N_EMBD);
    embeds.insert(embeds.end(), cb0_emb, cb0_emb + N_EMBD);

    std::vector<float> hn;
    std::vector<float> scores(N_VOCAB);
    int local_history[N_CB];  // tokens samplés ce step (pour rep_penalty intra-frame)
    int n_local = 0;
    for (int step = 1; step <= N_CB; step++) {
        const int L = step + 1;
        cp_forward_lastpos(s, embeds, L, step, hn);

        // tête[step-1] : scores[j] = sum_k hn[k] * heads[step-1, j, k]
        const float* Wh = &s.heads[(size_t)(step - 1) * N_VOCAB * N_EMBD];
        for (int j = 0; j < N_VOCAB; j++) {
            const float* wj = Wh + (size_t)j * N_EMBD;
            float dot = 0.f;
            for (int k = 0; k < N_EMBD; k++) dot += hn[k] * wj[k];
            scores[j] = dot;
        }
        // Sample via sampler partagé (HF-style). rep_penalty appliquée sur les tokens déjà
        // émis CE step (intra-frame), pas l'historique long.
        int chosen = sampler_sample_local(s.sampler, scores.data(), N_VOCAB,
                                          local_history, n_local);
        out_codes[step - 1] = chosen;
        local_history[n_local++] = chosen;

        if (step < N_CB) {
            const float* e = &s.codec_embs[(size_t)(step - 1) * N_VOCAB * N_EMBD + (size_t)chosen * N_EMBD];
            embeds.insert(embeds.end(), e, e + N_EMBD);
        }
    }
}

void cp_predict_cached(CPState& s, const float* hidden, const float* cb0_emb, int32_t* out_codes) {
    if (!cp_cache_init(s)) { fprintf(stderr, "cp_predict_cached: cache init failed\n"); abort(); }

    std::vector<float> hn;
    std::vector<float> scores(N_VOCAB);
    int local_history[N_CB];  // rep_penalty intra-frame, identique à cp_predict
    int n_local = 0;

    auto sample_head = [&](int head_idx) -> int {
        const float* Wh = &s.heads[(size_t)head_idx * N_VOCAB * N_EMBD];
        for (int j = 0; j < N_VOCAB; j++) {
            const float* wj = Wh + (size_t)j * N_EMBD;
            float dot = 0.f;
            for (int k = 0; k < N_EMBD; k++) dot += hn[k] * wj[k];
            scores[j] = dot;
        }
        return sampler_sample_local(s.sampler, scores.data(), N_VOCAB, local_history, n_local);
    };

    // Prefill : 2 tokens [hidden, cb0_emb] aux positions 0,1. hn = sortie à pos 1 -> head[0] -> CB1.
    std::vector<float> prefill_emb((size_t)2 * N_EMBD);
    memcpy(prefill_emb.data(),           hidden,  N_EMBD * sizeof(float));
    memcpy(prefill_emb.data() + N_EMBD,  cb0_emb, N_EMBD * sizeof(float));
    cp_forward_cached_step(s, prefill_emb.data(), /*n_new=*/2, /*pos_off=*/0, hn);
    int chosen = sample_head(0);
    out_codes[0] = chosen;
    local_history[n_local++] = chosen;

    // Decode : 14 étapes. À l'étape `step` (1..14), input = codec_embs[step-1][out_codes[step-1]]
    // placé à pos = step+1 ; hn à cette pos -> head[step] -> out_codes[step].
    for (int step = 1; step < N_CB; step++) {
        const int prev = out_codes[step - 1];
        const float* e = &s.codec_embs[(size_t)(step - 1) * N_VOCAB * N_EMBD + (size_t)prev * N_EMBD];
        const int pos_off = 1 + step;  // 2..15
        cp_forward_cached_step(s, e, /*n_new=*/1, pos_off, hn);
        chosen = sample_head(step);
        out_codes[step] = chosen;
        local_history[n_local++] = chosen;
    }
}