// Reproduction C++ de la construction Python des talker_input_embeds (Qwen3-TTS,
// mode x_vector_only + non_streaming). Reference : modeling_qwen3_tts.py lignes
// 2124-2233.
//
// Séquence prefill (19 positions pour "Bonjour je m'appelle Kazeia") :
//   pos  0..2 : text_projection( text_embed[input_id[0..3]] )                    role <im_start>assistant\n
//   pos  3..7 : tts_pad_embed + tok_embd[ codec_prefill[0..3] + x_vector_pos ]   5 pos = think/think_bos/lang/think_eos/x_vector
//   pos    8  : tts_bos_embed + tok_embd[ codec_pad_id ]                          1 pos
//   pos  9..N : text_projection( text_embed[input_id[3:-5]] ) + tok_embd[ codec_pad_id ] * N  (N = len text tokens)
//   pos   N+1 : tts_eos_embed + tok_embd[ codec_pad_id ]
//   pos   N+2 : tts_pad_embed + tok_embd[ codec_bos_id ]
//
// tts_*_embed = text_projection( text_embed[ tts_bos/eos/pad_token_id ] ).chunk(3)
//
// Usage : build_prefill <dump_dir> <out.bin>
//   dump_dir contient : text_embed.bin, tp_fc1_w.bin, tp_fc1_b.bin, tp_fc2_w.bin,
//                       tp_fc2_b.bin, talker_tok_embd.bin, damien_xvector.bin,
//                       input_ids_full.bin, manifest_text.txt, talker_prefill_embeds.bin
// Sortie : <out.bin> = prefill_embeds [T, 1024] f32 + validation vs Python.
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <cstdint>
#include <vector>
#include <string>
#include <fstream>

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_expected && 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;
}
static std::vector<int32_t> read_i32(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(int32_t);
    if (n_expected && n != n_expected) { fprintf(stderr, "%s: %zu i32, attendu %zu\n", p.c_str(), n, n_expected); exit(1); }
    f.seekg(0);
    std::vector<int32_t> v(n);
    f.read((char*)v.data(), n * sizeof(int32_t));
    return v;
}

// out[m_out] = bias[m_out] + sum_k W[m_out, k_in] * x[k_in]  (W stocké en [m_out, k_in] row-major,
// = numpy [out_features, in_features])
static void linear(const float* W, const float* b, int out_dim, int in_dim,
                   const float* x, float* out) {
    for (int m = 0; m < out_dim; ++m) {
        float s = b ? b[m] : 0.0f;
        const float* wr = W + (size_t)m * in_dim;
        for (int k = 0; k < in_dim; ++k) s += wr[k] * x[k];
        out[m] = s;
    }
}
static inline float silu(float x) { return x / (1.0f + std::exp(-x)); }

// text_projection : Linear(2048→2048, bias) → SiLU → Linear(2048→1024, bias).
// In: [N, 2048] → Out: [N, 1024]. Buffers réutilisables passés en arg.
static void text_projection(const float* in, int N,
                            const float* fc1_w, const float* fc1_b,
                            const float* fc2_w, const float* fc2_b,
                            float* mid_buf, float* out) {
    for (int n = 0; n < N; ++n) {
        linear(fc1_w, fc1_b, 2048, 2048, in + n * 2048, mid_buf);
        for (int d = 0; d < 2048; ++d) mid_buf[d] = silu(mid_buf[d]);
        linear(fc2_w, fc2_b, 1024, 2048, mid_buf, out + n * 1024);
    }
}

int main(int argc, char** argv) {
    if (argc < 3) { printf("usage: %s <dump_dir> <out.bin>\n", argv[0]); return 1; }
    std::string D = argv[1]; if (D.back() != '/') D += '/';
    const char* OUT = argv[2];

    // manifest_text
    int text_vocab = 151936, text_hidden = 2048, hidden = 1024;
    int tts_bos = 151672, tts_eos = 151673, tts_pad = 151671;
    int codec_bos = 2149, codec_eos = 2150, codec_pad = 2148;
    int codec_think = 2154, codec_nothink = 2155, codec_think_bos = 2156, codec_think_eos = 2157;
    int lang_fr = 2061;
    int input_ids_len = 16;
    {
        std::ifstream f(D + "manifest_text.txt"); if (!f) { fprintf(stderr, "no manifest_text\n"); return 1; }
        std::string line;
        auto eq = [&](const char* k){ return line.rfind(k, 0) == 0; };
        auto val = [&](size_t off){ return atoi(line.c_str() + off); };
        while (std::getline(f, line)) {
            if      (eq("text_vocab_size:"))     text_vocab = val(16);
            else if (eq("text_hidden_size:"))    text_hidden = val(17);
            else if (eq("hidden_size:"))         hidden = val(12);
            else if (eq("tts_bos_token_id:"))    tts_bos = val(17);
            else if (eq("tts_eos_token_id:"))    tts_eos = val(17);
            else if (eq("tts_pad_token_id:"))    tts_pad = val(17);
            else if (eq("codec_bos_id:"))        codec_bos = val(13);
            else if (eq("codec_eos_id:"))        codec_eos = val(13);
            else if (eq("codec_pad_id:"))        codec_pad = val(13);
            else if (eq("codec_think_id:"))      codec_think = val(15);
            else if (eq("codec_nothink_id:"))    codec_nothink = val(17);
            else if (eq("codec_think_bos_id:"))  codec_think_bos = val(19);
            else if (eq("codec_think_eos_id:"))  codec_think_eos = val(19);
            else if (eq("codec_language_french:")) lang_fr = val(22);
            else if (eq("input_ids_len:"))       input_ids_len = val(14);
        }
    }
    printf("manifest: text_vocab=%d text_hid=%d hid=%d input_ids_len=%d lang_fr=%d\n",
           text_vocab, text_hidden, hidden, input_ids_len, lang_fr);

    // fixtures
    printf("loading text_embed (%.1f GB)...\n", (double)text_vocab * text_hidden * 4 / 1e9);
    auto text_embed   = read_f32(D + "text_embed.bin",   (size_t)text_vocab * text_hidden);
    auto tp_fc1_w     = read_f32(D + "tp_fc1_w.bin",     (size_t)text_hidden * text_hidden);
    auto tp_fc1_b     = read_f32(D + "tp_fc1_b.bin",     (size_t)text_hidden);
    auto tp_fc2_w     = read_f32(D + "tp_fc2_w.bin",     (size_t)hidden * text_hidden);
    auto tp_fc2_b     = read_f32(D + "tp_fc2_b.bin",     (size_t)hidden);
    auto tok_embd     = read_f32(D + "talker_tok_embd.bin", (size_t)3072 * hidden);
    auto xvector      = read_f32(D + "damien_xvector.bin",  (size_t)hidden);
    auto input_ids    = read_i32(D + "input_ids_full.bin",  (size_t)input_ids_len);
    auto pf_py        = read_f32(D + "talker_prefill_embeds.bin", 0);  // reference
    const int T_pf_py = pf_py.size() / hidden;
    printf("fixtures loaded. input_ids[16]:");
    for (int i = 0; i < input_ids_len; ++i) printf(" %d", input_ids[i]);
    printf("\nreference prefill_embeds: T=%d\n", T_pf_py);

    // == Step 1 : tts_bos/eos/pad embeds via text_projection(text_embed[tk_ids])
    std::vector<float> tk_special_text(3 * text_hidden);
    int special_ids[3] = { tts_bos, tts_eos, tts_pad };
    for (int i = 0; i < 3; ++i) {
        const float* row = text_embed.data() + (size_t)special_ids[i] * text_hidden;
        std::memcpy(tk_special_text.data() + i * text_hidden, row, text_hidden * sizeof(float));
    }
    std::vector<float> tk_special_proj(3 * hidden);
    std::vector<float> mid(text_hidden);
    text_projection(tk_special_text.data(), 3, tp_fc1_w.data(), tp_fc1_b.data(),
                    tp_fc2_w.data(), tp_fc2_b.data(), mid.data(), tk_special_proj.data());
    const float* tts_bos_emb = tk_special_proj.data() + 0 * hidden;
    const float* tts_eos_emb = tk_special_proj.data() + 1 * hidden;
    const float* tts_pad_emb = tk_special_proj.data() + 2 * hidden;

    // == Step 2 : codec_prefill_list (french) + speaker x_vector
    // codec_prefill = [think_id, think_bos_id, lang_french, think_eos_id]
    int codec_prefill[4] = { codec_think, codec_think_bos, lang_fr, codec_think_eos };
    // codec_input_embedding_0 = tok_embd[codec_prefill]      [4, 1024]
    // codec_input_embedding_1 = tok_embd[[codec_pad, codec_bos]]  [2, 1024]
    // codec_input_embedding   = cat(0, speaker_embed, 1)     [7, 1024]
    // _talker_input_embed     = [pad*5, bos] + codec_input_embedding[:6]  -> [6, 1024]
    // codec_input_embedding[-1] = tok_embd[codec_bos]  (utilisé pour la position "first text token")
    std::vector<float> codec_input_emb(7 * hidden);
    for (int i = 0; i < 4; ++i) {
        std::memcpy(codec_input_emb.data() + i * hidden,
                    tok_embd.data() + (size_t)codec_prefill[i] * hidden,
                    hidden * sizeof(float));
    }
    std::memcpy(codec_input_emb.data() + 4 * hidden, xvector.data(), hidden * sizeof(float));
    std::memcpy(codec_input_emb.data() + 5 * hidden, tok_embd.data() + (size_t)codec_pad * hidden, hidden * sizeof(float));
    std::memcpy(codec_input_emb.data() + 6 * hidden, tok_embd.data() + (size_t)codec_bos * hidden, hidden * sizeof(float));

    // == Step 3 : role = text_projection(text_embed[input_ids[:3]])  shape [3, 1024]
    std::vector<float> role_text(3 * text_hidden);
    for (int i = 0; i < 3; ++i) {
        std::memcpy(role_text.data() + i * text_hidden,
                    text_embed.data() + (size_t)input_ids[i] * text_hidden,
                    text_hidden * sizeof(float));
    }
    std::vector<float> role_proj(3 * hidden);
    text_projection(role_text.data(), 3, tp_fc1_w.data(), tp_fc1_b.data(),
                    tp_fc2_w.data(), tp_fc2_b.data(), mid.data(), role_proj.data());

    // == Step 4 : text body = text_projection(text_embed[input_ids[3 .. input_ids_len - 5]])
    const int text_body_start = 3;
    const int text_body_end   = input_ids_len - 5;
    const int N_text          = text_body_end - text_body_start;
    std::vector<float> body_text(N_text * text_hidden);
    for (int i = 0; i < N_text; ++i) {
        std::memcpy(body_text.data() + i * text_hidden,
                    text_embed.data() + (size_t)input_ids[text_body_start + i] * text_hidden,
                    text_hidden * sizeof(float));
    }
    std::vector<float> body_proj(N_text * hidden);
    text_projection(body_text.data(), N_text, tp_fc1_w.data(), tp_fc1_b.data(),
                    tp_fc2_w.data(), tp_fc2_b.data(), mid.data(), body_proj.data());

    // == Step 5 : assemble prefill_embeds [T, 1024], T = 3 + 6 + (N_text + 1) + 1 = 3 + 6 + N_text + 2 = N_text + 11
    // Sequence:
    //   0..2          : role_proj  (3)
    //   3..7          : tts_pad_emb + codec_input_emb[i]   for i in 0..4  -> [think, think_bos, lang, think_eos, x_vector]
    //   8             : tts_bos_emb + codec_input_emb[5]   = tts_bos + tok_embd[codec_pad]
    //   9..(9+Nt-1)   : body_proj[i] + tok_embd[codec_pad]
    //   9+Nt          : tts_eos_emb + tok_embd[codec_pad]
    //   9+Nt+1        : tts_pad_emb + tok_embd[codec_bos]
    const int T = 3 + 6 + (N_text + 1) + 1;
    std::vector<float> prefill(T * hidden, 0.0f);

    // pos 0..2 : role
    std::memcpy(prefill.data() + 0 * hidden, role_proj.data(), 3 * hidden * sizeof(float));
    // pos 3..7 : tts_pad + codec_input[0..4]
    for (int i = 0; i < 5; ++i) {
        const float* ce = codec_input_emb.data() + i * hidden;
        float* p = prefill.data() + (3 + i) * hidden;
        for (int d = 0; d < hidden; ++d) p[d] = tts_pad_emb[d] + ce[d];
    }
    // pos 8 : tts_bos + codec_input[5]  (codec_input[5] = tok_embd[codec_pad])
    {
        const float* ce = codec_input_emb.data() + 5 * hidden;
        float* p = prefill.data() + 8 * hidden;
        for (int d = 0; d < hidden; ++d) p[d] = tts_bos_emb[d] + ce[d];
    }
    // pos 9..(9+Nt-1) : body_proj + tok_embd[codec_pad]
    const float* codec_pad_emb = tok_embd.data() + (size_t)codec_pad * hidden;
    for (int i = 0; i < N_text; ++i) {
        const float* bp = body_proj.data() + i * hidden;
        float* p = prefill.data() + (9 + i) * hidden;
        for (int d = 0; d < hidden; ++d) p[d] = bp[d] + codec_pad_emb[d];
    }
    // pos 9+Nt : tts_eos + tok_embd[codec_pad]
    {
        float* p = prefill.data() + (9 + N_text) * hidden;
        for (int d = 0; d < hidden; ++d) p[d] = tts_eos_emb[d] + codec_pad_emb[d];
    }
    // pos 9+Nt+1 : tts_pad + tok_embd[codec_bos]
    {
        const float* cbe = tok_embd.data() + (size_t)codec_bos * hidden;
        float* p = prefill.data() + (10 + N_text) * hidden;
        for (int d = 0; d < hidden; ++d) p[d] = tts_pad_emb[d] + cbe[d];
    }

    // == Step 6 : compare to Python reference
    printf("prefill construit : T=%d (Python: %d)\n", T, T_pf_py);
    if (T != T_pf_py) { fprintf(stderr, "MISMATCH T\n"); return 2; }
    double tot_rmse = 0;
    float max_abs = 0;
    for (int i = 0; i < T; ++i) {
        double rmse = 0;
        float local_max = 0;
        for (int d = 0; d < hidden; ++d) {
            float diff = prefill[i * hidden + d] - pf_py[i * hidden + d];
            rmse += diff * diff;
            if (std::abs(diff) > local_max) local_max = std::abs(diff);
        }
        rmse = std::sqrt(rmse / hidden);
        if (local_max > max_abs) max_abs = local_max;
        tot_rmse += rmse;
        printf("  pos %2d: rmse=%.4e max_abs=%.4e %s\n", i, rmse, local_max, local_max < 1e-3 ? "✓" : "");
    }
    printf("SUMMARY: mean_rmse=%.4e max_abs_overall=%.4e\n", tot_rmse / T, max_abs);

    {
        std::ofstream f(OUT, std::ios::binary);
        f.write((const char*)prefill.data(), prefill.size() * sizeof(float));
    }
    printf("prefill -> %s (%d * %d f32)\n", OUT, T, hidden);
    return 0;
}