// Forward complet du décodeur Qwen3-TTS en pur ggml (CPU baseline).
// Approche staging strict : chaque sous-module = un graphe ggml indépendant
// avec son propre ggml_init, son ggml_graph_compute_with_ctx, et le résultat
// recopié en host array entre étages. C'est verbeux mais reprend exactement
// le code des tests par stage qui sont validés numériquement individuellement.

#include "decoder.h"
#include <ggml-backend.h>
#include <ggml-cpu.h>
#include <cassert>
#include <chrono>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cmath>
#include <vector>

namespace kazeia::tts {

// ----- Stage 1 : Quantizer (codes → hidden [512, T]) -----
// codes_flat layout : [n_q=16][T] flat = codes_flat[layer*T + t]
// hidden output layout (host) : [c][t] flat with stride compatible test_quantizer's ggml output :
//   ggml ne=[512, T] → memory data[c + 512*t] (col-major).
//   On copie memory tel quel dans le host array (480KB pour T=59).
static std::vector<float> stage_quantizer(Decoder & d, const std::vector<int32_t> & codes_flat, int T) {
    const int n_q = d.cfg.num_quantizers;  // 16
    const int dim = d.cfg.codebook_dim / 2; // 256 (rvq dimension)
    const int hidden_dim = d.cfg.codebook_dim;  // 512

    size_t mem_size = 64 * 1024 * 1024;
    struct ggml_init_params p = { mem_size, nullptr, false };
    struct ggml_context * ctx = ggml_init(p);

    auto build_codes = [&](int layer_idx) {
        ggml_tensor * t = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, T);
        memcpy(t->data, &codes_flat[layer_idx * T], T * sizeof(int32_t));
        return t;
    };

    // Sum_first : 1 codebook (rvq_first.cb.0)
    ggml_tensor * emb_first = d.tensors["quantizer.rvq_first.cb.0.embedding"];
    ggml_tensor * sum_first = ggml_get_rows(ctx, emb_first, build_codes(0));  // [256, T]

    // Sum_rest : 15 codebooks (rvq_rest.cb.0..14)
    ggml_tensor * sum_rest = nullptr;
    for (int i = 0; i < 15; i++) {
        char key[64]; snprintf(key, sizeof(key), "quantizer.rvq_rest.cb.%d.embedding", i);
        ggml_tensor * emb = d.tensors[key];
        ggml_tensor * x = ggml_get_rows(ctx, emb, build_codes(1 + i));
        sum_rest = sum_rest ? ggml_add(ctx, sum_rest, x) : x;
    }

    auto apply_proj = [&](const std::string & key, ggml_tensor * x_KT) -> ggml_tensor * {
        ggml_tensor * w = d.tensors[key];
        ggml_tensor * w2 = ggml_reshape_2d(ctx, w, w->ne[1], w->ne[2]);
        return ggml_mul_mat(ctx, w2, x_KT);
    };

    ggml_tensor * out_first = apply_proj("quantizer.rvq_first.output_proj.weight", sum_first);
    ggml_tensor * out_rest  = apply_proj("quantizer.rvq_rest.output_proj.weight",  sum_rest);
    ggml_tensor * out = ggml_add(ctx, out_first, out_rest);  // [512, T]

    struct ggml_cgraph * gf = ggml_new_graph(ctx);
    ggml_build_forward_expand(gf, out);
    ggml_build_forward_expand(gf, out_first);
    ggml_build_forward_expand(gf, out_rest);
    ggml_build_forward_expand(gf, sum_first);
    ggml_build_forward_expand(gf, sum_rest);
    int n_threads = 8;
    if (const char * t_env = std::getenv("GGML_NUM_THREADS")) {
        int n = std::atoi(t_env); if (n >= 1) n_threads = n;
    }
    ggml_graph_compute_with_ctx(ctx, gf, n_threads);

    // Debug dumps
    if (const char * dir = std::getenv("DECODER_DEBUG_DUMP_DIR")) {
        auto dump_tensor = [&](const std::string & name, ggml_tensor * t) {
            std::string path = std::string(dir) + "/" + name + ".bin";
            FILE * f = fopen(path.c_str(), "wb");
            if (f) {
                size_t n_elem = ggml_nelements(t);
                fwrite(t->data, sizeof(float), n_elem, f);
                fclose(f);
                fprintf(stderr, "  dump: %s ne=[%lld,%lld] (%zu floats)\n",
                        path.c_str(), (long long)t->ne[0], (long long)t->ne[1], n_elem);
            }
        };
        dump_tensor("q_sum_first", sum_first);
        dump_tensor("q_sum_rest",  sum_rest);
        dump_tensor("q_out_first", out_first);
        dump_tensor("q_out_rest",  out_rest);
    }

    std::vector<float> result(hidden_dim * T);
    memcpy(result.data(), out->data, hidden_dim * T * sizeof(float));

    ggml_free(ctx);
    (void)dim; (void)n_q;
    return result;
}

// ----- Stage 2 : pre_conv (Conv1d k=3 causal pad 2-left + 1-right) -----
// Reproduit exactement test_pre_conv. Input layout : ggml [Tp=T+3, in=512]
// avec dst[(t+2) + Tp*c] = src_input[c + 512*t]. Output : [out=1024, T] dans
// le sens ggml-natif après slice T_out=T+1 → premier T par channel.
static std::vector<float> stage_pre_conv(Decoder & d, const std::vector<float> & hidden, int T) {
    const int Cin = d.cfg.codebook_dim;     // 512
    const int Cout = d.cfg.latent_dim;      // 1024
    const int Tp = T + 3;

    size_t mem_size = 64 * 1024 * 1024;
    struct ggml_init_params p = { mem_size, nullptr, false };
    struct ggml_context * ctx = ggml_init(p);

    ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, Tp, Cin);
    {
        float * dst = (float*) x->data;
        for (int i = 0; i < Tp * Cin; i++) dst[i] = 0.0f;
        // hidden in [c + 512*t] layout (from quantizer output memcpy).
        // Map to ggml[t, c] = pytorch[c, t-2] → dst[(t+2) + Tp*c] = hidden[c + 512*t]
        for (int c = 0; c < Cin; c++) {
            for (int t = 0; t < T; t++) {
                dst[(t + 2) + Tp * c] = hidden[c + Cin * t];
            }
        }
    }

    ggml_tensor * w = d.tensors["pre_conv.weight"];   // F16
    ggml_tensor * b = d.tensors["pre_conv.bias"];
    ggml_tensor * out = ggml_conv_1d(ctx, w, x, 1, 0, 1);
    // out ne=[T+1, Cout]
    ggml_tensor * b2 = ggml_reshape_2d(ctx, b, 1, Cout);
    out = ggml_add(ctx, out, b2);

    struct ggml_cgraph * gf = ggml_new_graph(ctx);
    ggml_build_forward_expand(gf, out);
    int n_threads = 8;
    if (const char * t_env = std::getenv("GGML_NUM_THREADS")) {
        int n = std::atoi(t_env); if (n >= 1) n_threads = n;
    }
    ggml_graph_compute_with_ctx(ctx, gf, n_threads);

    // out ne=[T+1, Cout] → take first T per channel.
    // Output layout (matches test_pre_conv) : memory [t + T_out*c] but slice to T.
    // Pour la stage suivante (pre_transformer), on veut input dans le format
    // que test_pre_transformer attend : ggml ne=[latent=Cout, T] avec
    //   dst[c + latent*t] = src[t + T*c] (= pytorch [B, latent, T]).
    // Or notre out actuel a memory[t + (T+1)*c] = pytorch_pre_conv[c, t].
    // On extrait de manière à produire host[c + Cout*t] = pytorch[c, t].
    const int T_out_y = (int)out->ne[0];
    const float * data = (const float*) out->data;
    std::vector<float> result(Cout * T);
    for (int c = 0; c < Cout; c++) {
        for (int t = 0; t < T; t++) {
            result[c + Cout * t] = data[t + T_out_y * c];
        }
    }

    ggml_free(ctx);
    return result;
}

// ----- Stage 3 : pre_transformer (input_proj + 8 DiTLayers + norm + output_proj) -----
// Input layout (host) : pc_out[c + Cout*t] = pytorch_pre_conv[c, t]
// = ggml ne=[latent, T] data[c + latent*t]
// Output layout (host) : pt_out[c + latent*t] = pytorch_pre_transformer_out[t, c]
//   (PyTorch last_hidden_state is [B, T, latent], les valeurs sont les mêmes
//   que [B, latent, T] indexées par (c, t) au lieu de (t, c) → mêmes valeurs scalaires).
static std::vector<float> stage_pre_transformer(Decoder & d, const std::vector<float> & pc_out, int T) {
    const int latent      = d.cfg.latent_dim;            // 1024
    const int hidden      = d.cfg.hidden_size;           // 512
    const int head_dim    = d.cfg.head_dim;              // 64
    const int n_heads     = d.cfg.num_attention_heads;   // 16
    const int qkv_dim     = n_heads * head_dim;
    const int n_layers    = d.cfg.num_hidden_layers;     // 8
    const float rope_base = d.cfg.rope_theta;
    const float rms_eps   = d.cfg.rms_norm_eps;

    size_t mem_size = 512ULL * 1024 * 1024;
    struct ggml_init_params p = { mem_size, nullptr, false };
    struct ggml_context * ctx = ggml_init(p);

    // Input ggml ne=[latent, T] avec dst[c + latent*t] = pc_out[c + latent*t]
    ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, latent, T);
    memcpy(x->data, pc_out.data(), latent * T * sizeof(float));

    auto get = [&](const std::string & key) -> ggml_tensor * { return d.tensors[key]; };

    auto linear = [&](ggml_tensor * x, const std::string & w_key, const std::string & b_key = "") -> ggml_tensor * {
        ggml_tensor * y = ggml_mul_mat(ctx, get(w_key), x);
        if (!b_key.empty()) {
            ggml_tensor * b = get(b_key);
            ggml_tensor * b2 = ggml_reshape_2d(ctx, b, b->ne[0], 1);
            y = ggml_add(ctx, y, b2);
        }
        return y;
    };

    x = linear(x, "pre_transformer.input_proj.weight", "pre_transformer.input_proj.bias");

    ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, T);
    {
        int32_t * pp = (int32_t*) pos->data;
        for (int t = 0; t < T; t++) pp[t] = t;
    }
    ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, T, T);
    {
        float * mm = (float*) mask->data;
        for (int q = 0; q < T; q++)
            for (int k = 0; k < T; k++)
                mm[k + T*q] = (k > q) ? -INFINITY : 0.0f;
    }
    const float scale = 1.0f / std::sqrt((float)head_dim);

    for (int L = 0; L < n_layers; L++) {
        char prefix[64]; snprintf(prefix, sizeof(prefix), "pre_transformer.%d", L);
        ggml_tensor * residual = x;

        ggml_tensor * x_n = ggml_rms_norm(ctx, x, rms_eps);
        x_n = ggml_mul(ctx, x_n, get(std::string(prefix) + ".input_norm.weight"));

        ggml_tensor * Q = linear(x_n, std::string(prefix) + ".attn.q.weight");
        ggml_tensor * K = linear(x_n, std::string(prefix) + ".attn.k.weight");
        ggml_tensor * V = linear(x_n, std::string(prefix) + ".attn.v.weight");

        Q = ggml_reshape_3d(ctx, Q, head_dim, n_heads, T);
        K = ggml_reshape_3d(ctx, K, head_dim, n_heads, T);
        V = ggml_reshape_3d(ctx, V, head_dim, n_heads, T);
        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, qkv_dim, T);

        ggml_tensor * attn_out = linear(KQV, std::string(prefix) + ".attn.o.weight");
        attn_out = ggml_mul(ctx, attn_out, get(std::string(prefix) + ".attn_layerscale"));
        x = ggml_add(ctx, residual, attn_out);

        residual = x;
        ggml_tensor * x_n2 = ggml_rms_norm(ctx, x, rms_eps);
        x_n2 = ggml_mul(ctx, x_n2, get(std::string(prefix) + ".post_attn_norm.weight"));
        ggml_tensor * gate = linear(x_n2, std::string(prefix) + ".mlp.gate.weight");
        ggml_tensor * up   = linear(x_n2, std::string(prefix) + ".mlp.up.weight");
        gate = ggml_silu(ctx, gate);
        ggml_tensor * mid = ggml_mul(ctx, gate, up);
        ggml_tensor * mlp_out = linear(mid, std::string(prefix) + ".mlp.down.weight");
        mlp_out = ggml_mul(ctx, mlp_out, get(std::string(prefix) + ".mlp_layerscale"));
        x = ggml_add(ctx, residual, mlp_out);
    }
    x = ggml_rms_norm(ctx, x, rms_eps);
    x = ggml_mul(ctx, x, get("pre_transformer.norm.weight"));
    x = linear(x, "pre_transformer.output_proj.weight", "pre_transformer.output_proj.bias");
    // x ne=[latent, T]
    (void)hidden;

    struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 8192, false);
    ggml_build_forward_expand(gf, x);
    int n_threads = 8;
    if (const char * t_env = std::getenv("GGML_NUM_THREADS")) {
        int n = std::atoi(t_env); if (n >= 1) n_threads = n;
    }
    ggml_graph_compute_with_ctx(ctx, gf, n_threads);

    std::vector<float> result(latent * T);
    memcpy(result.data(), x->data, latent * T * sizeof(float));
    // result[c + latent*t] = pytorch_pre_transformer_out at (logical c, logical t)
    // = pytorch [B, T, latent] at (t, c) = pytorch [B, latent, T] at (c, t) — valeur scalaire.

    ggml_free(ctx);
    return result;
}

// ----- Stage 4 : upsample (2 stages : ConvTranspose1d + ConvNeXtBlock) -----
// Input host : pt_out[c + latent*t] (sortie de pre_transformer en CT layout ggml)
// Output host : us_out[c + latent*t_out] avec t_out = T*4 (après 2× upsample par 2)
//
// Reproduit exactement test_upsample (deux stages cumulés) sauf que l'input
// vient d'un host array au lieu d'un .npy.
static std::vector<float> stage_upsample(Decoder & d, const std::vector<float> & pt_out, int T) {
    const int C = d.cfg.latent_dim;  // 1024
    const int T_final = T * 4;       // après 2 upsampling ratio 2

    size_t mem_size = 1024ULL * 1024 * 1024;
    struct ggml_init_params p = { mem_size, nullptr, false };
    struct ggml_context * ctx = ggml_init(p);

    auto get = [&](const std::string & key) { return d.tensors[key]; };

    // Input ggml ne=[C, T_in] : memcpy direct depuis pt_out (même layout)
    ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, C, T);
    memcpy(x->data, pt_out.data(), C * T * sizeof(float));

    int T_cur = T;
    for (int s = 0; s < 2; s++) {
        char prefix[64]; snprintf(prefix, sizeof(prefix), "upsample.%d", s);

        // Stage A : ConvTranspose1d k=2 s=2
        ggml_tensor * x_tch = ggml_cont(ctx, ggml_transpose(ctx, x));  // [T, C]
        ggml_tensor * y = ggml_conv_transpose_1d(ctx, get(std::string(prefix) + ".convtr.weight"),
                                                  x_tch, 2, 0, 1);
        ggml_tensor * b_ct = get(std::string(prefix) + ".convtr.bias");
        y = ggml_add(ctx, y, ggml_reshape_2d(ctx, b_ct, 1, C));
        y = ggml_cont(ctx, ggml_transpose(ctx, y));  // [C, T*2]
        T_cur *= 2;

        // Stage B : ConvNeXtBlock
        ggml_tensor * residual = y;

        // dwconv k=7 causal pad gauche=6
        ggml_tensor * y_tch = ggml_cont(ctx, ggml_transpose(ctx, y));  // [T*2, C]
        ggml_tensor * zeros = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 6, C);
        memset(zeros->data, 0, 6 * C * sizeof(float));
        ggml_tensor * y_pad = ggml_concat(ctx, zeros, y_tch, 0);  // [pad+T, C]
        ggml_tensor * dw = ggml_conv_1d_dw(ctx, get(std::string(prefix) + ".cnx.dwconv.weight"),
                                            y_pad, 1, 0, 1);
        dw = ggml_add(ctx, dw, ggml_reshape_2d(ctx, get(std::string(prefix) + ".cnx.dwconv.bias"), 1, C));
        dw = ggml_cont(ctx, ggml_transpose(ctx, dw));  // [C, T]

        ggml_tensor * h = ggml_norm(ctx, dw, 1e-6f);
        h = ggml_mul(ctx, h, get(std::string(prefix) + ".cnx.norm.weight"));
        h = ggml_add(ctx, h, get(std::string(prefix) + ".cnx.norm.bias"));
        h = ggml_mul_mat(ctx, get(std::string(prefix) + ".cnx.pwconv1.weight"), h);
        ggml_tensor * pw1_b = get(std::string(prefix) + ".cnx.pwconv1.bias");
        h = ggml_add(ctx, h, ggml_reshape_2d(ctx, pw1_b, pw1_b->ne[0], 1));
        h = ggml_gelu_erf(ctx, h);
        h = ggml_mul_mat(ctx, get(std::string(prefix) + ".cnx.pwconv2.weight"), h);
        ggml_tensor * pw2_b = get(std::string(prefix) + ".cnx.pwconv2.bias");
        h = ggml_add(ctx, h, ggml_reshape_2d(ctx, pw2_b, pw2_b->ne[0], 1));
        h = ggml_mul(ctx, h, get(std::string(prefix) + ".cnx.gamma"));
        x = ggml_add(ctx, residual, h);  // [C, T*2]
    }

    struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 4096, false);
    ggml_build_forward_expand(gf, x);
    int n_threads = 8;
    if (const char * t_env = std::getenv("GGML_NUM_THREADS")) {
        int n = std::atoi(t_env); if (n >= 1) n_threads = n;
    }
    ggml_graph_compute_with_ctx(ctx, gf, n_threads);

    std::vector<float> result(C * T_final);
    memcpy(result.data(), x->data, C * T_final * sizeof(float));

    ggml_free(ctx);
    return result;
}

// ----- Stage 5 : BigVGAN (initial Conv k=7 + 4 blocks + final Snake + final Conv k=7) -----
// Reproduit test_decoder. Input : us_out[c + C*t] (CT layout, T = T_in * 4).
// Output : wav[s] for s=0..T_in*1920-1.
static std::vector<float> stage_bigvgan(Decoder & d, const std::vector<float> & us_out, int T_in_us) {
    const int latent = d.cfg.latent_dim;       // 1024 (= C of upsample output)
    const int decoder_dim = d.cfg.decoder_dim; // 1536
    const std::vector<int> & up_rates = d.cfg.upsample_rates;

    size_t mem_size = 12ULL * 1024 * 1024 * 1024;
    struct ggml_init_params p = { mem_size, nullptr, false };
    struct ggml_context * ctx = ggml_init(p);

    auto get = [&](const std::string & key) { return d.tensors[key]; };

    const float SNAKE_EPS = 1e-9f;
    auto precompute_snake = [&](const std::string & ka, const std::string & kb,
                                 ggml_tensor ** oa, ggml_tensor ** ob) -> bool {
        ggml_tensor * a = d.tensors[ka];
        ggml_tensor * b = d.tensors[kb];
        if (!a || !b) return false;
        const int C = (int)a->ne[0];
        ggml_tensor * a_eff = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, C);
        ggml_tensor * inv_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, C);
        const float * pa = (const float*) a->data;
        const float * pb = (const float*) b->data;
        float * da = (float*) a_eff->data;
        float * db = (float*) inv_b->data;
        for (int c = 0; c < C; c++) {
            da[c] = std::exp(pa[c]);
            db[c] = 1.0f / (std::exp(pb[c]) + SNAKE_EPS);
        }
        *oa = a_eff; *ob = inv_b;
        return true;
    };

    auto apply_snake_TC = [&](ggml_tensor * x_TC, ggml_tensor * a_eff, ggml_tensor * inv_b) {
        const int C = (int)a_eff->ne[0];
        ggml_tensor * a2 = ggml_reshape_2d(ctx, a_eff, 1, C);
        ggml_tensor * b2 = ggml_reshape_2d(ctx, inv_b, 1, C);
        ggml_tensor * xa = ggml_mul(ctx, x_TC, a2);
        ggml_tensor * s = ggml_sin(ctx, xa);
        ggml_tensor * s2 = ggml_mul(ctx, s, s);
        ggml_tensor * mod = ggml_mul(ctx, s2, b2);
        return ggml_add(ctx, x_TC, mod);
    };

    auto causal_conv1d_TC = [&](ggml_tensor * x_TC, const std::string & w_key, const std::string & b_key,
                                int K, int dilation) -> ggml_tensor * {
        ggml_tensor * w = d.tensors[w_key];
        ggml_tensor * b = d.tensors[b_key];
        const int C_out = (int)w->ne[2];
        const int T_cur = (int)x_TC->ne[0];
        const int C_in  = (int)x_TC->ne[1];
        ggml_tensor * x_pad;
        if (K == 1) {
            x_pad = x_TC;
        } else {
            const int pad_left = (K - 1) * dilation;
            ggml_tensor * zeros = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, pad_left, C_in);
            memset(zeros->data, 0, pad_left * C_in * sizeof(float));
            x_pad = ggml_concat(ctx, zeros, x_TC, 0);
        }
        ggml_tensor * y = ggml_conv_1d(ctx, w, x_pad, 1, 0, dilation);
        const int T_out_y = (int)y->ne[0];
        if (T_out_y > T_cur) {
            y = ggml_view_2d(ctx, y, T_cur, C_out, y->nb[1], 0);
        }
        ggml_tensor * b2 = ggml_reshape_2d(ctx, b, 1, C_out);
        return ggml_add(ctx, y, b2);
    };

    auto causal_convtr1d_TC = [&](ggml_tensor * x_TC, const std::string & w_key, const std::string & b_key,
                                   int s) -> ggml_tensor * {
        ggml_tensor * w = d.tensors[w_key];
        ggml_tensor * b = d.tensors[b_key];
        const int C_out = (int)w->ne[1];
        const int T_cur = (int)x_TC->ne[0];
        ggml_tensor * y = ggml_conv_transpose_1d(ctx, w, x_TC, s, 0, 1);
        const int T_full = (int)y->ne[0];
        const int T_target = T_cur * s;
        if (T_full > T_target) {
            y = ggml_view_2d(ctx, y, T_target, C_out, y->nb[1], 0);
        }
        ggml_tensor * b2 = ggml_reshape_2d(ctx, b, 1, C_out);
        return ggml_add(ctx, y, b2);
    };

    // Input : us_out [C=latent, T] (CT layout host) → ggml ne=[C, T] memcpy.
    ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, latent, T_in_us);
    memcpy(x->data, us_out.data(), latent * T_in_us * sizeof(float));

    // Switch to TC layout for whole BigVGAN (Optim #3)
    ggml_tensor * x_TC = ggml_cont(ctx, ggml_transpose(ctx, x));  // [T, latent]

    // initial Conv k=7
    x_TC = causal_conv1d_TC(x_TC, "decoder.initial.weight", "decoder.initial.bias", 7, 1);

    // 4 DecoderBlocks
    for (int B = 0; B < 4; B++) {
        char prefix[64]; snprintf(prefix, sizeof(prefix), "decoder.block.%d", B);
        const int up = up_rates[B];

        ggml_tensor *a0 = nullptr, *b0 = nullptr;
        precompute_snake(std::string(prefix) + ".block.0.alpha",
                         std::string(prefix) + ".block.0.beta", &a0, &b0);
        x_TC = apply_snake_TC(x_TC, a0, b0);
        x_TC = causal_convtr1d_TC(x_TC,
            std::string(prefix) + ".block.1.conv.weight",
            std::string(prefix) + ".block.1.conv.bias", up);

        const int dilations[3] = {1, 3, 9};
        for (int R = 0; R < 3; R++) {
            char rprefix[80]; snprintf(rprefix, sizeof(rprefix), "%s.block.%d", prefix, 2 + R);
            ggml_tensor * res = x_TC;
            ggml_tensor *a1 = nullptr, *b1 = nullptr;
            precompute_snake(std::string(rprefix) + ".act1.alpha",
                             std::string(rprefix) + ".act1.beta", &a1, &b1);
            x_TC = apply_snake_TC(x_TC, a1, b1);
            x_TC = causal_conv1d_TC(x_TC, std::string(rprefix) + ".conv1.conv.weight",
                                            std::string(rprefix) + ".conv1.conv.bias", 7, dilations[R]);
            ggml_tensor *a2 = nullptr, *b2 = nullptr;
            precompute_snake(std::string(rprefix) + ".act2.alpha",
                             std::string(rprefix) + ".act2.beta", &a2, &b2);
            x_TC = apply_snake_TC(x_TC, a2, b2);
            x_TC = causal_conv1d_TC(x_TC, std::string(rprefix) + ".conv2.conv.weight",
                                            std::string(rprefix) + ".conv2.conv.bias", 1, 1);
            x_TC = ggml_add(ctx, res, x_TC);
        }
    }

    // final SnakeBeta + Conv k=7
    ggml_tensor *af = nullptr, *bf = nullptr;
    precompute_snake("decoder.final_act.alpha", "decoder.final_act.beta", &af, &bf);
    x_TC = apply_snake_TC(x_TC, af, bf);
    x_TC = causal_conv1d_TC(x_TC, "decoder.final_conv.weight", "decoder.final_conv.bias", 7, 1);

    struct ggml_cgraph * gf = ggml_new_graph_custom(ctx, 16384, false);
    ggml_build_forward_expand(gf, x_TC);
    int n_threads = 8;
    if (const char * t_env = std::getenv("GGML_NUM_THREADS")) {
        int n = std::atoi(t_env); if (n >= 1) n_threads = n;
    }
    ggml_graph_compute_with_ctx(ctx, gf, n_threads);

    // x_TC ne=[T_final, 1] — single channel audio
    const int T_final = (int)x_TC->ne[0];
    std::vector<float> wav(T_final);
    const float * data = (const float*) x_TC->data;
    for (int t = 0; t < T_final; t++) {
        float v = data[t];
        if (v > 1.0f) v = 1.0f;
        if (v < -1.0f) v = -1.0f;
        wav[t] = v;
    }

    ggml_free(ctx);
    (void)decoder_dim;
    return wav;
}

// ============== Forward complet : chaîne staging strict ==============

// Dump an array of floats to disk if env var DECODER_DEBUG_DUMP_DIR is set.
static void debug_dump(const std::string & name, const std::vector<float> & data) {
    const char * dir = std::getenv("DECODER_DEBUG_DUMP_DIR");
    if (!dir) return;
    std::string path = std::string(dir) + "/" + name + ".bin";
    FILE * f = fopen(path.c_str(), "wb");
    if (!f) {
        fprintf(stderr, "debug_dump: cannot open %s\n", path.c_str());
        return;
    }
    fwrite(data.data(), sizeof(float), data.size(), f);
    fclose(f);
    fprintf(stderr, "debug dump: %s (%zu floats)\n", path.c_str(), data.size());
}

std::vector<float> Decoder::forward(const std::vector<int32_t>& codes_flat, int T) {
    assert((int)codes_flat.size() == cfg.num_quantizers * T);
    // KZTTS_DECODER_PROFILE=1 -> imprime le temps pris par chaque stage.
    const bool prof = (std::getenv("KZTTS_DECODER_PROFILE") &&
                       std::atoi(std::getenv("KZTTS_DECODER_PROFILE")) != 0);
    auto nows = []() {
        return std::chrono::duration<double>(
            std::chrono::steady_clock::now().time_since_epoch()).count();
    };
    double t0 = nows(), t1, t2, t3, t4, t5;

    auto hidden = stage_quantizer(*this, codes_flat, T);                  // [512, T]
    t1 = nows();
    debug_dump("stage1_quantizer", hidden);
    auto pc_out = stage_pre_conv(*this, hidden, T);                       // [1024, T]
    t2 = nows();
    debug_dump("stage2_pre_conv", pc_out);
    auto pt_out = stage_pre_transformer(*this, pc_out, T);                // [1024, T]
    t3 = nows();
    debug_dump("stage3_pre_transformer", pt_out);
    auto us_out = stage_upsample(*this, pt_out, T);                       // [1024, T*4]
    t4 = nows();
    debug_dump("stage4_upsample", us_out);
    auto wav    = stage_bigvgan(*this, us_out, T * 4);                    // T*1920 samples
    t5 = nows();
    debug_dump("stage5_wav", wav);

    if (prof) {
        fprintf(stderr,
                "decoder: T=%d  total=%.3fs  quant=%.3f preconv=%.3f pretr=%.3f"
                " upsample=%.3f bigvgan=%.3f\n",
                T, t5 - t0, t1 - t0, t2 - t1, t3 - t2, t4 - t3, t5 - t4);
    }
    return wav;
}

// Stub legacy
std::vector<float> Decoder::test_quantizer(const std::vector<int32_t>& codes_flat, int T) {
    (void)codes_flat; (void)T;
    return {};
}

}  // namespace kazeia::tts