// vae.h: AutoencoderOobleck decoder (audio VAE) via ggml
//
// Architecture: conv1(64->2048,k=7) -> 5xblock(snake+convT+3xresunit) -> snake+conv2(128->2,k=7)
// ResUnit(ch, dil): skip=x -> snake->conv(k=7,dil)->snake->conv(k=1)->+skip
// Snake: x + sin^2(e^a * x) * (1/e^b)  [decomposed: mul->sin->sqr->mul->add]
// ConvT: ggml_conv_transpose_1d (standard upstream op)
// Weight norm fused at load: w = g*v/||v||
// Upsample: 10x6x4x4x2 = 1920x

#pragma once
#include "backend.h"
#include "ggml-backend.h"
#include "ggml.h"
#include "gguf-weights.h"
#include "safetensors.h"

#include <cmath>
#include <cstdio>
#include <string>
#include <unordered_map>
#include <vector>

// Cross-platform case-insensitive string compare
#ifndef _MSC_VER
#include <strings.h>
#define _stricmp strcasecmp
#endif

// Structs
struct VAEResUnit {
    struct ggml_tensor *s1a, *s1b;  // snake1 exp(alpha), exp(beta) [1, C]
    struct ggml_tensor *c1w, *c1b;  // conv1 fused [7, C, C], bias [C]
    struct ggml_tensor *s2a, *s2b;  // snake2
    struct ggml_tensor *c2w, *c2b;  // conv2 fused [1, C, C], bias [C]
    int                 dilation;
};

struct VAEBlock {
    struct ggml_tensor *sa, *sb;    // snake exp(a/b) [1, in_ch]
    struct ggml_tensor *ctw, *ctb;  // conv_transpose F16 [IC, K*OC] pre-transposed for GEMM, bias [out_ch]
    int                 in_ch, out_ch, stride, kernel;
    VAEResUnit          ru[3];
};

struct VAEGGML {
    struct ggml_tensor * c1w, *c1b;  // conv1 [7, 64, 2048], bias [2048]
    VAEBlock             blk[5];
    struct ggml_tensor * sa, *sb;    // final snake [1, 128]
    struct ggml_tensor * c2w;        // conv2 [7, 128, 2] (no bias)

    ggml_backend_t        backend;
    ggml_backend_t        cpu_backend;
    ggml_backend_sched_t  sched;
    ggml_backend_buffer_t buf;
    struct ggml_context * weight_ctx;  // holds weight tensor metadata

    // Graph cache for tiled decode (avoids rebuild per tile)
    struct ggml_context * graph_ctx;
    uint8_t *             graph_buf;  // heap-allocated backing for graph_ctx
    struct ggml_cgraph *  graph;
    struct ggml_tensor *  graph_input;
    struct ggml_tensor *  graph_output;
    int                   graph_T;  // cached T_latent (0 = no cache)

    // Scratch buffer (reused across tiles, grown as needed)
    std::vector<float> scratch_in;  // transposed input [64 * T]
};

// Auto-remap PyTorch Sequential layers.N naming to named-module naming.
// Builds a remap table from the safetensors entries on first call.
// Returns the remapped name, or the original name if no remap needed.
struct VaeKeyRemap {
    bool                                              built = false;
    std::unordered_map<std::string, std::string>      map;   // named-module -> layers.N

    // Build the remap table by scanning all safetensors entries.
    // Only builds if sequential naming (layers.N) is detected.
    void build(const STFile & st) {
        if (built) return;
        built = true;

        // Detect sequential naming: look for "decoder.layers." or "encoder.layers."
        bool has_sequential = false;
        for (const auto & e : st.entries) {
            if (e.name.find("decoder.layers.") == 0 || e.name.find("encoder.layers.") == 0) {
                has_sequential = true;
                break;
            }
        }
        if (!has_sequential) return;

        fprintf(stderr, "[VAE] Detected PyTorch Sequential naming, building auto-remap table\n");

        // Build reverse map: for each sequential key, compute the named-module equivalent.
        // The C++ loader will ask for named-module keys; we map them to the sequential keys.
        for (const auto & e : st.entries) {
            std::string k = e.name;
            std::string remapped = remap_sequential_key(k);
            if (remapped != k) {
                map[remapped] = k;  // named-module -> sequential
            }
        }

        fprintf(stderr, "[VAE] Auto-remap: %zu key mappings built\n", map.size());
    }

    // Remap a single PyTorch Sequential key to named-module convention.
    // decoder.layers.0.* -> decoder.conv1.*
    // decoder.layers.1..5.* -> decoder.block.0..4.*
    // decoder.layers.6.* -> decoder.snake1.*
    // decoder.layers.7.* -> decoder.conv2.*
    // block.N.layers.0.* -> block.N.snake1.*
    // block.N.layers.1.* -> block.N.conv_t1.*
    // block.N.layers.2..4.* -> block.N.res_unit1..3.*
    // res_unitR.layers.0..3.* -> res_unitR.snake1/conv1/snake2/conv2.*
    // Same pattern for encoder (with conv1 for downsample instead of conv_t1).
    static std::string remap_sequential_key(const std::string & k) {
        std::string r = k;

        // --- Top-level decoder layers ---
        r = replace_prefix(r, "decoder.layers.0.", "decoder.conv1.");
        r = replace_prefix(r, "decoder.layers.6.", "decoder.snake1.");
        r = replace_prefix(r, "decoder.layers.7.", "decoder.conv2.");
        for (int n = 1; n <= 5; n++) {
            std::string from = "decoder.layers." + std::to_string(n) + ".";
            std::string to   = "decoder.block." + std::to_string(n - 1) + ".";
            r = replace_prefix(r, from, to);
        }

        // --- Top-level encoder layers ---
        r = replace_prefix(r, "encoder.layers.0.", "encoder.conv1.");
        r = replace_prefix(r, "encoder.layers.6.", "encoder.snake1.");
        r = replace_prefix(r, "encoder.layers.7.", "encoder.conv2.");
        for (int n = 1; n <= 5; n++) {
            std::string from = "encoder.layers." + std::to_string(n) + ".";
            std::string to   = "encoder.block." + std::to_string(n - 1) + ".";
            r = replace_prefix(r, from, to);
        }

        // --- Inside blocks: layers.0..4 -> snake1/conv_t1(decoder) or conv1(encoder)/res_unitN ---
        // Decoder blocks: snake1, conv_t1, res_unit1..3
        r = replace_block_layer(r, "decoder.block.", 0, "snake1.");
        r = replace_block_layer(r, "decoder.block.", 1, "conv_t1.");
        for (int ru = 0; ru < 3; ru++) {
            std::string to = "res_unit" + std::to_string(ru + 1) + ".";
            r = replace_block_layer(r, "decoder.block.", ru + 2, to);
        }
        // Encoder blocks: res_unit1..3, snake1, conv1 (downsample)
        for (int ru = 0; ru < 3; ru++) {
            std::string to = "res_unit" + std::to_string(ru + 1) + ".";
            r = replace_block_layer(r, "encoder.block.", ru, to);
        }
        r = replace_block_layer(r, "encoder.block.", 3, "snake1.");
        r = replace_block_layer(r, "encoder.block.", 4, "conv1.");

        // --- Inside res_units: layers.0..3 -> snake1/conv1/snake2/conv2 ---
        static const char * ru_map[] = { "snake1.", "conv1.", "snake2.", "conv2." };
        for (int ru = 1; ru <= 3; ru++) {
            std::string ru_pfx = "res_unit" + std::to_string(ru) + ".";
            for (int l = 0; l < 4; l++) {
                std::string from = ru_pfx + "layers." + std::to_string(l) + ".";
                std::string to   = ru_pfx + ru_map[l];
                r = replace_anywhere(r, from, to);
            }
        }

        return r;
    }

private:
    // Replace prefix at the start of the string
    static std::string replace_prefix(const std::string & s, const std::string & from, const std::string & to) {
        if (s.compare(0, from.size(), from) == 0) {
            return to + s.substr(from.size());
        }
        return s;
    }

    // Replace "<block_prefix>N.layers.<layer_idx>." with "<block_prefix>N.<replacement>"
    // for any block number N (0..9)
    static std::string replace_block_layer(const std::string & s, const std::string & block_prefix,
                                           int layer_idx, const std::string & replacement) {
        std::string layer_str = ".layers." + std::to_string(layer_idx) + ".";
        // Find block prefix
        if (s.compare(0, block_prefix.size(), block_prefix) != 0) return s;
        // Find block number
        size_t num_start = block_prefix.size();
        size_t num_end = s.find('.', num_start);
        if (num_end == std::string::npos) return s;
        std::string block_num = s.substr(num_start, num_end - num_start);
        // Check if followed by .layers.<layer_idx>.
        std::string full_from = block_prefix + block_num + layer_str;
        if (s.compare(0, full_from.size(), full_from) == 0) {
            return block_prefix + block_num + "." + replacement + s.substr(full_from.size());
        }
        return s;
    }

    // Replace first occurrence of 'from' anywhere in the string
    static std::string replace_anywhere(const std::string & s, const std::string & from, const std::string & to) {
        size_t pos = s.find(from);
        if (pos == std::string::npos) return s;
        return s.substr(0, pos) + to + s.substr(pos + from.size());
    }
};

// Abstracts tensor data access for VAE loading from GGUF or safetensors.
struct VaeWeightSource {
    bool            is_st = false;
    GGUFModel *     gf    = nullptr;
    STFile *        st    = nullptr;
    VaeKeyRemap     remap;  // lazy-built auto-remap for sequential naming

    // Resolve a tensor name through the remap table (safetensors only).
    // Returns the original name if no remap is needed.
    const char * resolve(const char * name) {
        if (!is_st) return name;
        remap.build(*st);
        auto it = remap.map.find(name);
        if (it != remap.map.end()) {
            return it->second.c_str();
        }
        return name;
    }

    // Get raw data pointer + type for a named tensor
    const void * data(const char * name, ggml_type & type) {
        if (is_st) {
            const char * resolved = resolve(name);
            const void * p = st_find_data(*st, resolved, &type);
            if (!p) {
                fprintf(stderr, "[VAE] FATAL: tensor '%s' not found in safetensors", name);
                if (strcmp(name, resolved) != 0) {
                    fprintf(stderr, " (tried remap: '%s')", resolved);
                }
                fprintf(stderr, "\n");
            }
            return p;
        } else {
            struct ggml_tensor * mt = ggml_get_tensor(gf->meta, name);
            if (!mt) {
                fprintf(stderr, "[VAE] FATAL: tensor '%s' not found in GGUF\n", name);
                return nullptr;
            }
            type = mt->type;
            return gf_get_data(*gf, name);
        }
    }

    // Get shape info for a named tensor (n_dims and ne[])
    bool shape(const char * name, int & n_dims, int64_t ne[4]) {
        if (is_st) {
            const char * resolved = resolve(name);
            const STEntry * e = st_find(*st, resolved);
            if (!e) return false;
            n_dims = e->n_dims;
            // Reverse shape: safetensors is PyTorch order [dim0..dimN], ggml is [neN..ne0]
            for (int i = 0; i < 4; i++) {
                int src = e->n_dims - 1 - i;
                ne[i] = (src >= 0 && src < e->n_dims) ? e->shape[src] : 1;
            }
            return true;
        } else {
            struct ggml_tensor * mt = ggml_get_tensor(gf->meta, name);
            if (!mt) return false;
            n_dims = ggml_n_dims(mt);
            for (int i = 0; i < 4; i++) ne[i] = mt->ne[i];
            return true;
        }
    }
};

// Load helpers

// Type-aware element reader: extract float from tensor data at index.
// Supports F32, BF16, F16 source tensors (format-agnostic: works for both GGUF and safetensors).
static inline float vae_read_float(const void * data, int idx, ggml_type type) {
    switch (type) {
        case GGML_TYPE_F32:
            return ((const float *) data)[idx];
        case GGML_TYPE_BF16:
            return ggml_bf16_to_fp32(*(const ggml_bf16_t *) &((const uint16_t *) data)[idx]);
        case GGML_TYPE_F16:
            return ggml_fp16_to_fp32(((const ggml_fp16_t *) data)[idx]);
        default:
            fprintf(stderr, "[VAE] FATAL: unsupported tensor type %d\n", (int) type);
            return 0.0f;
    }
}

// Fuse weight_norm: w = g*v/||v||, write f32 into pre-allocated ggml_tensor
// Works for Conv1d [OC,IC,K]: weight_norm normalizes over dim=0 (shape[0]).
// Type-aware: reads F32, BF16, or F16 source tensors.
static void vae_fuse_wn(struct ggml_tensor * dst, VaeWeightSource & ws, const std::string & pfx) {
    ggml_type v_type, g_type;
    const void * v = ws.data((pfx + ".weight_v").c_str(), v_type);
    const void * g = ws.data((pfx + ".weight_g").c_str(), g_type);
    if (!v || !g) { fprintf(stderr, "[VAE] FATAL: missing weight_norm tensors for '%s', aborting load\n", pfx.c_str()); exit(1); }
    int n_dims_v; int64_t ne_v[4];
    ws.shape((pfx + ".weight_v").c_str(), n_dims_v, ne_v);
    int dim0 = (int) ne_v[n_dims_v - 1];
    int64_t total = 1;
    for (int i = 0; i < n_dims_v; i++) total *= ne_v[i];
    int fan = (int) (total / dim0);
    std::vector<float> w(dim0 * fan);
    for (int d = 0; d < dim0; d++) {
        float gv  = vae_read_float(g, d, g_type);
        float nsq = 0;
        for (int i = 0; i < fan; i++) {
            float vv = vae_read_float(v, d * fan + i, v_type);
            nsq += vv * vv;
        }
        float s = gv / (sqrtf(nsq) + 1e-12f);
        for (int i = 0; i < fan; i++) {
            float vv       = vae_read_float(v, d * fan + i, v_type);
            w[d * fan + i] = vv * s;
        }
    }
    if (dst->type == GGML_TYPE_F16) {
        std::vector<ggml_fp16_t> w16(w.size());
        ggml_fp32_to_fp16_row(w.data(), w16.data(), (int) w.size());
        ggml_backend_tensor_set(dst, w16.data(), 0, w16.size() * sizeof(ggml_fp16_t));
    } else {
        ggml_backend_tensor_set(dst, w.data(), 0, w.size() * sizeof(float));
    }
}

// Fuse weight_norm for ConvTranspose1d into [IC, K*OC] layout for GEMM-based transpose conv.
// Source weight_v is [K, OC, IC] (ggml ne[0]=K, ne[1]=OC, ne[2]=IC).
// weight_norm dim0=IC, fan=K*OC.  Data is transposed so mul_mat contracts over IC.
// Type-aware: reads F32, BF16, or F16 source tensors.
static void vae_fuse_wn_ct(struct ggml_tensor * dst, VaeWeightSource & ws, const std::string & pfx) {
    ggml_type v_type, g_type;
    const void * v = ws.data((pfx + ".weight_v").c_str(), v_type);
    const void * g = ws.data((pfx + ".weight_g").c_str(), g_type);
    if (!v || !g) { fprintf(stderr, "[VAE] FATAL: missing weight_norm tensors for '%s', aborting load\n", pfx.c_str()); exit(1); }
    int n_dims_v; int64_t ne_v[4];
    ws.shape((pfx + ".weight_v").c_str(), n_dims_v, ne_v);
    int dim0 = (int) ne_v[n_dims_v - 1];  // IC
    int64_t total = 1;
    for (int i = 0; i < n_dims_v; i++) total *= ne_v[i];
    int fan = (int) (total / dim0);       // K*OC
    std::vector<float> w(dim0 * fan);
    for (int d = 0; d < dim0; d++) {
        float gv  = vae_read_float(g, d, g_type);
        float nsq = 0;
        for (int i = 0; i < fan; i++) {
            float vv = vae_read_float(v, d * fan + i, v_type);
            nsq += vv * vv;
        }
        float s = gv / (sqrtf(nsq) + 1e-12f);
        for (int i = 0; i < fan; i++) {
            float vv        = vae_read_float(v, d * fan + i, v_type);
            w[i * dim0 + d] = vv * s;    // transposed: [k_oc * IC + ic] for GEMM
        }
    }
    if (dst->type == GGML_TYPE_F16) {
        std::vector<ggml_fp16_t> w16(w.size());
        ggml_fp32_to_fp16_row(w.data(), w16.data(), (int) w.size());
        ggml_backend_tensor_set(dst, w16.data(), 0, w16.size() * sizeof(ggml_fp16_t));
    } else {
        ggml_backend_tensor_set(dst, w.data(), 0, w.size() * sizeof(float));
    }
}

// Load snake param [1,C,1] -> exp -> f32 [1, C]
// Type-aware: reads F32, BF16, or F16 source tensors.
static void vae_load_snake(struct ggml_tensor * dst, VaeWeightSource & ws, const std::string & name) {
    int n_dims; int64_t ne[4];
    ws.shape(name.c_str(), n_dims, ne);
    int C = (int) ne[1];  // PyTorch [1,C,1] -> ggml ne=[1,C,1], middle dim
    ggml_type type;
    const void * raw = ws.data(name.c_str(), type);
    if (!raw) { fprintf(stderr, "[VAE] FATAL: missing tensor '%s', aborting load\n", name.c_str()); exit(1); }
    std::vector<float> d(C);
    for (int i = 0; i < C; i++) {
        d[i] = expf(vae_read_float(raw, i, type));
    }
    ggml_backend_tensor_set(dst, d.data(), 0, C * sizeof(float));
}

// Load snake param [1,C,1] -> 1/exp -> f32 [1, C] (reciprocal for mul fusion)
// Type-aware: reads F32, BF16, or F16 source tensors.
static void vae_load_snake_inv(struct ggml_tensor * dst, VaeWeightSource & ws, const std::string & name) {
    int n_dims; int64_t ne[4];
    ws.shape(name.c_str(), n_dims, ne);
    int C = (int) ne[1];
    ggml_type type;
    const void * raw = ws.data(name.c_str(), type);
    if (!raw) { fprintf(stderr, "[VAE] FATAL: missing tensor '%s', aborting load\n", name.c_str()); exit(1); }
    std::vector<float> d(C);
    for (int i = 0; i < C; i++) {
        d[i] = 1.0f / expf(vae_read_float(raw, i, type));
    }
    ggml_backend_tensor_set(dst, d.data(), 0, C * sizeof(float));
}

// Load bias [C] -> f32
// Type-aware: reads F32, BF16, or F16 source tensors.
static void vae_load_bias(struct ggml_tensor * dst, VaeWeightSource & ws, const std::string & name) {
    int n_dims; int64_t ne[4];
    ws.shape(name.c_str(), n_dims, ne);
    int C = (int) ne[0];  // 1D: ne[0] = C
    ggml_type type;
    const void * raw = ws.data(name.c_str(), type);
    if (!raw) { fprintf(stderr, "[VAE] FATAL: missing tensor '%s', aborting load\n", name.c_str()); exit(1); }
    std::vector<float> d(C);
    for (int i = 0; i < C; i++) {
        d[i] = vae_read_float(raw, i, type);
    }
    ggml_backend_tensor_set(dst, d.data(), 0, C * sizeof(float));
}

// Load model
static void vae_ggml_load(VAEGGML * m, const char * path) {
    const char * ext = strrchr(path, '.');
    bool is_st = (ext && _stricmp(ext, ".safetensors") == 0);

    GGUFModel gf = {};
    STFile    st = {};
    if (is_st) {
        if (!st_open(&st, path)) {
            fprintf(stderr, "[VAE] FATAL: cannot load %s\n", path);
            exit(1);
        }
    } else {
        if (!gf_load(&gf, path)) {
            fprintf(stderr, "[VAE] FATAL: cannot load %s\n", path);
            exit(1);
        }
    }
    VaeWeightSource ws;
    ws.is_st = is_st;
    ws.gf    = &gf;
    ws.st    = &st;

    static const int strides[]   = { 10, 6, 4, 4, 2 };
    static const int in_ch[]     = { 2048, 1024, 512, 256, 128 };
    static const int out_ch[]    = { 1024, 512, 256, 128, 128 };
    static const int dilations[] = { 1, 3, 9 };

    // Phase 1: create tensor metadata (no_alloc context)
    size_t                  ctx_size = ggml_tensor_overhead() * 200;
    struct ggml_init_params p        = { ctx_size, NULL, true };
    m->weight_ctx                    = ggml_init(p);
    struct ggml_context * ctx        = m->weight_ctx;

    m->c1w = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, 7, 64, 2048);
    m->c1b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2048);

    for (int i = 0; i < 5; i++) {
        VAEBlock & b = m->blk[i];
        b.in_ch      = in_ch[i];
        b.out_ch     = out_ch[i];
        b.stride     = strides[i];
        b.kernel     = strides[i] * 2;
        int C        = out_ch[i];
        b.sa         = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, in_ch[i]);
        b.sb         = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, in_ch[i]);
        b.ctw        = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, in_ch[i], b.kernel * out_ch[i]);
        b.ctb        = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_ch[i]);
        for (int r = 0; r < 3; r++) {
            VAEResUnit & ru = b.ru[r];
            ru.dilation     = dilations[r];
            ru.s1a          = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, C);
            ru.s1b          = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, C);
            ru.c1w          = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, 7, C, C);
            ru.c1b          = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, C);
            ru.s2a          = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, C);
            ru.s2b          = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, C);
            ru.c2w          = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, 1, C, C);
            ru.c2b          = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, C);
        }
    }
    m->sa  = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, 128);
    m->sb  = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, 128);
    m->c2w = ggml_new_tensor_3d(ctx, GGML_TYPE_F16, 7, 128, 2);

    // Phase 2: allocate backend buffer
    BackendPair bp = backend_init("VAE");
    m->backend     = bp.backend;
    m->cpu_backend = bp.cpu_backend;
    m->sched       = backend_sched_new(bp, 8192);
    m->buf         = ggml_backend_alloc_ctx_tensors(ctx, m->backend);
    if (!m->buf) {
        fprintf(stderr, "[VAE] FATAL: failed to allocate weight buffer\n");
        exit(1);
    }
    fprintf(stderr, "[VAE] Backend: %s, Weight buffer: %.1f MB\n", ggml_backend_name(m->backend),
            (float) ggml_backend_buffer_get_size(m->buf) / (1024 * 1024));

    // Phase 3: load & fuse weights
    vae_fuse_wn(m->c1w, ws, "decoder.conv1");
    vae_load_bias(m->c1b, ws, "decoder.conv1.bias");

    for (int i = 0; i < 5; i++) {
        VAEBlock &  b       = m->blk[i];
        std::string blk_pfx = "decoder.block." + std::to_string(i);
        vae_load_snake(b.sa, ws, blk_pfx + ".snake1.alpha");
        vae_load_snake_inv(b.sb, ws, blk_pfx + ".snake1.beta");
        vae_fuse_wn_ct(b.ctw, ws, blk_pfx + ".conv_t1");
        vae_load_bias(b.ctb, ws, blk_pfx + ".conv_t1.bias");
        for (int r = 0; r < 3; r++) {
            VAEResUnit & ru = b.ru[r];
            std::string  rp = blk_pfx + ".res_unit" + std::to_string(r + 1);
            vae_load_snake(ru.s1a, ws, rp + ".snake1.alpha");
            vae_load_snake_inv(ru.s1b, ws, rp + ".snake1.beta");
            vae_fuse_wn(ru.c1w, ws, rp + ".conv1");
            vae_load_bias(ru.c1b, ws, rp + ".conv1.bias");
            vae_load_snake(ru.s2a, ws, rp + ".snake2.alpha");
            vae_load_snake_inv(ru.s2b, ws, rp + ".snake2.beta");
            vae_fuse_wn(ru.c2w, ws, rp + ".conv2");
            vae_load_bias(ru.c2b, ws, rp + ".conv2.bias");
        }
    }
    vae_load_snake(m->sa, ws, "decoder.snake1.alpha");
    vae_load_snake_inv(m->sb, ws, "decoder.snake1.beta");
    vae_fuse_wn(m->c2w, ws, "decoder.conv2");

    fprintf(stderr, "[VAE] Loaded (%s): 5 blocks, upsample=1920x, F32 activations\n", is_st ? "safetensors" : "GGUF");
    if (is_st) st_close(&st); else gf_close(&gf);
}

// Graph building
// Snake activation (decomposed): y = x + sin^2(a * x) * inv_b
// Uses standard GGML ops (mul->sin->sqr->mul->add) for vanilla ggml-org compatibility.
// x: [T, C], exp_a: [1, C], inv_b: [1, C] (pre-computed at load)
static struct ggml_tensor * vae_snake(struct ggml_context * ctx,
                                      struct ggml_tensor *  x,
                                      struct ggml_tensor *  exp_a,
                                      struct ggml_tensor *  inv_b) {
    struct ggml_tensor * ax  = ggml_mul(ctx, x, exp_a);   // a * x  (broadcast [1,C] over [T,C])
    struct ggml_tensor * s   = ggml_sin(ctx, ax);          // sin(a * x)
    struct ggml_tensor * s2  = ggml_sqr(ctx, s);           // sin^2(a * x)
    struct ggml_tensor * s2b = ggml_mul(ctx, s2, inv_b);   // sin^2(a * x) * inv_b
    return ggml_add(ctx, x, s2b);                          // x + sin^2(a * x) * inv_b
}

// Conv1d + bias: data [T, IC] -> [T_out, OC]
static struct ggml_tensor * vae_conv1d(struct ggml_context * ctx,
                                       struct ggml_tensor *  w,  // [K, IC, OC] (F16, pre-cast at load)
                                       struct ggml_tensor *  b,  // [OC] or NULL
                                       struct ggml_tensor *  x,  // [T, IC]
                                       int                   stride,
                                       int                   padding,
                                       int                   dilation) {
    struct ggml_tensor * y = ggml_conv_1d(ctx, w, x, stride, padding, dilation);
    // ggml_conv_1d returns [OL, OC, N=1], squeeze to 2d
    y                      = ggml_reshape_2d(ctx, y, y->ne[0], y->ne[1]);
    if (b) {
        // bias [OC] -> [1, OC] for broadcast over OL dimension
        struct ggml_tensor * b2d = ggml_reshape_2d(ctx, b, 1, b->ne[0]);
        y                        = ggml_add(ctx, y, b2d);
    }
    return y;
}

// ConvTranspose1d via GEMM + col2im using standard GGML ops only.
// w: [IC, K*OC] F16, pre-transposed at load for mul_mat (contracts over IC).
// x: [T_in, IC]
// Returns: [T_out_cropped, OC]
//
// PERF: ggml_conv_transpose_1d uses a naive O(N*IC*T) CUDA kernel. This GEMM
// approach uses cuBLAS (tensor-core accelerated), then standard reshape/pad/acc
// ops for the col2im scatter-add. ~100x faster for our tensor sizes.
//
// Requires K = 2*stride (true for all VAE blocks: kernel = stride * 2).
static struct ggml_tensor * vae_conv_t1d(struct ggml_context * ctx,
                                         struct ggml_tensor *  w,  // [IC, K*OC] F16 pre-transposed
                                         struct ggml_tensor *  b,  // [OC] or NULL
                                         struct ggml_tensor *  x,  // [T_in, IC]
                                         int                   stride,
                                         int                   padding,
                                         int                   oc) {
    int64_t T_in = x->ne[0];
    int64_t IC   = x->ne[1];
    int     S    = stride;
    int     K    = 2 * S;  // always true for our VAE
    int64_t T_out = (T_in - 1) * S + K;  // full output length
    (void) IC;

    // Step 1: Transpose x from [T_in, IC] to [IC, T_in]
    struct ggml_tensor * xt = ggml_cont(ctx, ggml_transpose(ctx, x));

    // Step 2: GEMM via cuBLAS — contracts over IC (ne[0] of both)
    // w: [IC, K*OC]  xt: [IC, T_in]  ->  col: [K*OC, T_in]
    struct ggml_tensor * col = ggml_mul_mat(ctx, w, xt);
    // col: ne0=K*OC, ne1=T_in

    // Step 3: col2im via reshape + permute + pad + acc
    // k_oc indexing: k_oc = k + oc*K (k is fast dim from weight layout)
    // Split K=2S kernel positions into two halves:
    //   set_A (k=0..S-1): fills output positions i*S+0 .. i*S+(S-1)
    //   set_B (k=S..2S-1): fills output positions i*S+S .. i*S+(2S-1)

    // Reshape col [K*OC, T_in] -> [K, OC, T_in] (k is ne0, fast dim)
    struct ggml_tensor * col3 = ggml_reshape_3d(ctx, col, K, oc, T_in);
    // col3: ne0=K, ne1=OC, ne2=T_in

    // --- Set A: first S kernel positions (k=0..S-1) ---
    // View [S, OC, T_in] starting at k=0
    struct ggml_tensor * setA = ggml_view_3d(ctx, col3,
                                              S, oc, T_in,
                                              col3->nb[1], col3->nb[2],
                                              0);
    // Permute [S, OC, T_in] -> [S, T_in, OC] for correct interleaving
    // permute(0, 2, 1): ne0'=S(ax0), ne1'=T_in(ax2), ne2'=OC(ax1)
    setA = ggml_cont(ctx, ggml_permute(ctx, setA, 0, 2, 1, 3));
    // Reshape to [S*T_in, OC] — contiguous interleaved time samples
    struct ggml_tensor * outA = ggml_reshape_2d(ctx, setA, S * T_in, oc);
    // Pad to full T_out: add S zeros at end -> [S*T_in + S, OC] = [T_out, OC]
    struct ggml_tensor * outA_pad = ggml_pad(ctx, outA, S, 0, 0, 0);

    // --- Set B: last S kernel positions (k=S..2S-1) ---
    // View [S, OC, T_in] starting at k=S (offset in ne0 dim)
    struct ggml_tensor * setB = ggml_view_3d(ctx, col3,
                                              S, oc, T_in,
                                              col3->nb[1], col3->nb[2],
                                              S * col3->nb[0]);  // S elements into ne0
    // Same permute + reshape as set A
    setB = ggml_cont(ctx, ggml_permute(ctx, setB, 0, 2, 1, 3));
    struct ggml_tensor * outB = ggml_reshape_2d(ctx, setB, S * T_in, oc);

    // Accumulate set B into padded set A at offset S (shifted by one stride)
    // ggml_acc: dst = a; view(dst, nb1, nb2, nb3, offset) += b
    struct ggml_tensor * y = ggml_acc(ctx, outA_pad, outB,
                                      outA_pad->nb[1],    // nb1: stride between rows in dst
                                      outA_pad->nb[2],    // nb2
                                      outA_pad->nb[3],    // nb3
                                      S * sizeof(float));  // offset: S samples into dim0

    // Step 4: Crop padding from output
    // ConvTranspose1d "padding" = remove `padding` samples from each end
    if (padding > 0) {
        int64_t T_crop = T_out - 2 * padding;
        y = ggml_view_2d(ctx, y, T_crop, oc,
                         y->nb[1],
                         padding * sizeof(float));
        y = ggml_cont(ctx, y);
    }

    // Step 5: Add bias
    if (b) {
        struct ggml_tensor * b2d = ggml_reshape_2d(ctx, b, 1, b->ne[0]);
        y                        = ggml_add(ctx, y, b2d);
    }
    return y;
}

// ResUnit forward
static struct ggml_tensor * vae_res_unit(struct ggml_context * ctx,
                                         VAEResUnit *          ru,
                                         struct ggml_tensor *  x) {  // [T, C]
    struct ggml_tensor * skip = x;

    // snake1 -> dilated conv(k=7) -> snake2 -> conv(k=1)
    int pad = 3 * ru->dilation;  // (k-1)*dil/2 = 3*dil
    x       = vae_snake(ctx, x, ru->s1a, ru->s1b);
    x       = vae_conv1d(ctx, ru->c1w, ru->c1b, x, 1, pad, ru->dilation);
    x       = vae_snake(ctx, x, ru->s2a, ru->s2b);
    x       = vae_conv1d(ctx, ru->c2w, ru->c2b, x, 1, 0, 1);

    return ggml_add(ctx, skip, x);
}

// Build full VAE decode graph
// latent: [T_latent, 64] -> audio: [T_audio, 2]
static struct ggml_tensor * vae_ggml_build_graph(struct ggml_context * ctx,
                                                 VAEGGML *             m,
                                                 struct ggml_tensor *  latent) {  // [T, 64] input

    // conv1: [T, 64] -> [T, 2048]
    struct ggml_tensor * x = vae_conv1d(ctx, m->c1w, m->c1b, latent, 1, 3, 1);

    // 5 decoder blocks
    for (int i = 0; i < 5; i++) {
        VAEBlock & b = m->blk[i];
        // snake -> conv_transpose (upsample)
        x            = vae_snake(ctx, x, b.sa, b.sb);
        int pad      = (b.kernel - b.stride) / 2;
        x            = vae_conv_t1d(ctx, b.ctw, b.ctb, x, b.stride, pad, b.out_ch);
        // 3 res units
        for (int r = 0; r < 3; r++) {
            x = vae_res_unit(ctx, &b.ru[r], x);
        }
    }

    // Final: snake -> conv2(128->2, k=7, pad=3)
    x = vae_snake(ctx, x, m->sa, m->sb);
    x = vae_conv1d(ctx, m->c2w, NULL, x, 1, 3, 1);

    return x;  // [T_audio, 2]
}

// Core compute: ensure graph cached, set input, run. Returns T_audio or -1.
// Output remains in m->graph_output for caller to read as needed.
static int vae_ggml_compute(VAEGGML *     m,
                            const float * latent,    // [T_full, 64] time-major
                            int           T_latent,  // window length to decode
                            int           win_start = 0) {     // offset into latent

    // Build graph only when T_latent changes (cached for tiled decode reuse)
    if (m->graph_T != T_latent) {
        if (m->graph_ctx) {
            ggml_backend_sched_reset(m->sched);
            ggml_free(m->graph_ctx);
            free(m->graph_buf);
        }

        // Graph context (generous fixed allocation)
        size_t ctx_size = ggml_tensor_overhead() * 1024 + ggml_graph_overhead_custom(8192, false);
        m->graph_buf    = (uint8_t *) malloc(ctx_size);
        if (!m->graph_buf) {
            fprintf(stderr, "[VAE] FATAL: OOM allocating graph context (%zu bytes) for T=%d\n", ctx_size, T_latent);
            m->graph_T = 0;
            return -1;
        }
        struct ggml_init_params p   = { ctx_size, m->graph_buf, true };
        struct ggml_context *   ctx = ggml_init(p);

        m->graph_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, T_latent, 64);
        ggml_set_name(m->graph_input, "vae_input");
        ggml_set_input(m->graph_input);

        m->graph_output = vae_ggml_build_graph(ctx, m, m->graph_input);
        ggml_set_name(m->graph_output, "vae_output");
        ggml_set_output(m->graph_output);

        m->graph = ggml_new_graph_custom(ctx, 8192, false);
        ggml_build_forward_expand(m->graph, m->graph_output);

        if (!ggml_backend_sched_alloc_graph(m->sched, m->graph)) {
            fprintf(stderr, "[VAE] FATAL: graph alloc failed for T=%d\n", T_latent);
            ggml_free(ctx);
            free(m->graph_buf);
            m->graph_ctx = NULL;
            m->graph_buf = NULL;
            m->graph_T   = 0;
            return -1;
        }

        m->graph_ctx = ctx;
        m->graph_T   = T_latent;
        fprintf(stderr, "[VAE] Graph: %d nodes, T_latent=%d\n", ggml_graph_n_nodes(m->graph), T_latent);
    }

    // Extract window + transpose: [T, 64] time-major -> ggml [T, 64] channel-major
    size_t in_size = 64 * T_latent;
    if (m->scratch_in.size() < in_size) {
        m->scratch_in.resize(in_size);
    }
    for (int c = 0; c < 64; c++) {
        for (int t = 0; t < T_latent; t++) {
            m->scratch_in[c * T_latent + t] = latent[(win_start + t) * 64 + c];
        }
    }
    ggml_backend_tensor_set(m->graph_input, m->scratch_in.data(), 0, in_size * sizeof(float));

    ggml_backend_sched_graph_compute(m->sched, m->graph);

    return (int) m->graph_output->ne[0];
}

// Decode API: latent [T_latent, 64] -> audio [2, T_audio] flat.
// Returns T_audio (or -1 on error).
static int vae_ggml_decode(VAEGGML * m, const float * latent, int T_latent, float * audio_out, int max_T_audio) {
    int T_audio = T_latent * 1920;
    if (T_audio > max_T_audio) {
        fprintf(stderr, "[VAE] T_audio %d exceeds max %d\n", T_audio, max_T_audio);
        return -1;
    }

    int T_out = vae_ggml_compute(m, latent, T_latent, 0);
    if (T_out < 0) {
        return -1;
    }

    ggml_backend_tensor_get(m->graph_output, audio_out, 0, T_out * 2 * sizeof(float));

    fprintf(stderr, "[VAE] Decoded: T_latent=%d -> T_audio=%d (%.2fs @ 48kHz)\n", T_latent, T_out,
            (float) T_out / 48000.0f);
    return T_out;
}

// Tiled decode: overlap-discard chunking for bounded VRAM usage.
// stride = chunk_size - 2*overlap
// For each tile: decode latent window with overlap context, trim to core, concatenate.
// Default chunk=256/overlap=64 matches reference code. Larger chunks (e.g. 1024)
// reduce tile count and improve throughput; adjust chunk/overlap to tune.
// Returns T_audio (total samples per channel) or -1 on error.
static int vae_ggml_decode_tiled(VAEGGML *     m,
                                 const float * latent,     // [T_latent, 64] flat time-major (DiT output layout)
                                 int           T_latent,
                                 float *       audio_out,  // [2, T_audio] flat (caller allocs)
                                 int           max_T_audio,
                                 int           chunk_size = 256,
                                 int           overlap    = 64,
                                 bool (*cancel)(void *)   = nullptr,
                                 void * cancel_data       = nullptr) {
    // Ensure positive stride (matches Python effective_overlap reduction)
    while (chunk_size - 2 * overlap <= 0 && overlap > 0) {
        overlap /= 2;
    }

    // Short sequence: decode directly
    if (T_latent <= chunk_size) {
        return vae_ggml_decode(m, latent, T_latent, audio_out, max_T_audio);
    }

    int stride    = chunk_size - 2 * overlap;
    int num_steps = (T_latent + stride - 1) / stride;

    fprintf(stderr, "[VAE] Tiled decode: %d tiles (chunk=%d, overlap=%d, stride=%d)\n", num_steps, chunk_size, overlap,
            stride);

    float upsample_factor = 0.0f;
    int   audio_write_pos = 0;

    for (int i = 0; i < num_steps; i++) {
        if (cancel && cancel(cancel_data)) {
            fprintf(stderr, "[VAE] Cancelled at tile %d/%d\n", i, num_steps);
            return -1;
        }
        // Core range in latent frames (the part we keep)
        int core_start = i * stride;
        int core_end   = core_start + stride;
        if (core_end > T_latent) {
            core_end = T_latent;
        }

        // Window range with overlap context
        int win_start = core_start - overlap;
        if (win_start < 0) {
            win_start = 0;
        }
        int win_end = core_end + overlap;
        if (win_end > T_latent) {
            win_end = T_latent;
        }
        int win_len = win_end - win_start;

        // Compute tile (graph cached, extract+transpose fused)
        int tile_T = vae_ggml_compute(m, latent, win_len, win_start);
        if (tile_T < 0) {
            fprintf(stderr, "[VAE] FATAL: tile %d decode failed\n", i);
            return -1;
        }

        // Determine upsample factor from first tile
        if (i == 0) {
            upsample_factor = (float) tile_T / (float) win_len;
            fprintf(stderr, "[VAE] Upsample factor: %.2f (expected ~1920)\n", upsample_factor);
        }

        // Compute trim in audio samples (matches Python int(round(...)))
        int added_start = core_start - win_start;
        int trim_start  = (int) roundf((float) added_start * upsample_factor);
        int added_end   = win_end - core_end;
        int trim_end    = (int) roundf((float) added_end * upsample_factor);

        int end_idx  = (trim_end > 0) ? (tile_T - trim_end) : tile_T;
        int core_len = end_idx - trim_start;
        if (core_len <= 0) {
            continue;
        }

        // Check output bounds
        if (audio_write_pos + core_len > max_T_audio) {
            fprintf(stderr, "[VAE] FATAL: tiled output exceeds max_T_audio\n");
            return -1;
        }

        // Read trimmed ch0 and ch1 directly from backend tensor into final audio_out
        // Layout: [ch0: tile_T floats, ch1: tile_T floats]
        ggml_backend_tensor_get(m->graph_output, audio_out + audio_write_pos, trim_start * sizeof(float),
                                core_len * sizeof(float));
        ggml_backend_tensor_get(m->graph_output, audio_out + max_T_audio + audio_write_pos,
                                (tile_T + trim_start) * sizeof(float), core_len * sizeof(float));
        audio_write_pos += core_len;
    }

    // Compact ch1 from offset max_T_audio to offset audio_write_pos
    memmove(audio_out + audio_write_pos, audio_out + max_T_audio, audio_write_pos * sizeof(float));

    fprintf(stderr, "[VAE] Tiled decode done: %d tiles -> T_audio=%d (%.2fs @ 48kHz)\n", num_steps, audio_write_pos,
            (float) audio_write_pos / 48000.0f);

    return audio_write_pos;
}

// Free
static void vae_ggml_free(VAEGGML * m) {
    if (m->graph_ctx) {
        ggml_backend_sched_reset(m->sched);
        ggml_free(m->graph_ctx);
        free(m->graph_buf);
    }
    if (m->sched) {
        ggml_backend_sched_free(m->sched);
    }
    if (m->buf) {
        ggml_backend_buffer_free(m->buf);
    }
    if (m->weight_ctx) {
        ggml_free(m->weight_ctx);
    }
    backend_release(m->backend, m->cpu_backend);
    *m = {};
}