#pragma once
// audio-tower.h: windowed forward of the Qwen3-ASR audio tower. Splits the mel
// into chunk_mel = n_window * 2 frame chunks, runs one batched conv stem pass
// over the full chunks (plus the shorter tail alone) with the positional table
// reset per chunk, concatenates the after-cnn frames and runs the encoder with
// a block-diagonal attention mask over windows of window_aftercnn =
// chunk_aftercnn * window_chunks frames (100 and 104 on the public
// checkpoints). Ties conv-stem.h and audio-enc.h into one graph.

#include "audio-enc.h"
#include "conv-stem.h"
#include "ggml.h"
#include "gguf-weights.h"

#include <cmath>
#include <cstddef>
#include <limits>
#include <vector>

struct AudioTowerConfig {
    int chunk_mel;      // n_window * 2
    int window_chunks;  // n_window_infer / chunk_mel
};

// Tower windowing read from the GGUF metadata. chunk_mel = n_window * 2,
// window_chunks = n_window_infer / chunk_mel.
static AudioTowerConfig audio_tower_config_load(const GGUFModel & gf) {
    AudioTowerConfig cfg;
    const int        n_window       = (int) gf_get_u32(gf, "qwenasr.audio.n_window");
    const int        n_window_infer = (int) gf_get_u32(gf, "qwenasr.audio.n_window_infer");
    cfg.chunk_mel                   = n_window * 2;
    cfg.window_chunks               = n_window_infer / cfg.chunk_mel;
    return cfg;
}

// after-cnn length of a conv chunk of mel_len frames, three conv2d stride 2
// kernel 3 pad 1. Matches the reference _get_feat_extract_output_lengths for
// any chunk length >= 1.
static int audio_tower_conv_out(int mel_len) {
    int o = mel_len;
    for (int i = 0; i < 3; i++) {
        o = (o - 1) / 2 + 1;
    }
    return o;
}

// Mel chunk plan: ceil(n_frames / chunk_mel) chunks, all chunk_mel long, the
// tail chunk carries the remainder (chunk_mel when it divides evenly).
static std::vector<int> audio_tower_chunk_lengths(int n_frames, const AudioTowerConfig & cfg) {
    std::vector<int> lengths;
    int              off = 0;
    while (off < n_frames) {
        const int len = (n_frames - off) < cfg.chunk_mel ? (n_frames - off) : cfg.chunk_mel;
        lengths.push_back(len);
        off += len;
    }
    return lengths;
}

// Total after-cnn length over all chunks.
static int audio_tower_seq_len(const std::vector<int> & chunk_lengths) {
    int s = 0;
    for (int len : chunk_lengths) {
        s += audio_tower_conv_out(len);
    }
    return s;
}

// Additive block-diagonal mask [S, S], memory key inner. 0 inside a window of
// window_aftercnn frames, -inf across windows.
static std::vector<float> audio_tower_build_mask(int seq_len, int window_aftercnn) {
    const float        neg = -std::numeric_limits<float>::infinity();
    std::vector<float> mask((size_t) seq_len * (size_t) seq_len, neg);
    for (int b0 = 0; b0 < seq_len; b0 += window_aftercnn) {
        const int b1 = (b0 + window_aftercnn) < seq_len ? (b0 + window_aftercnn) : seq_len;
        for (int q = b0; q < b1; q++) {
            for (int k = b0; k < b1; k++) {
                mask[(size_t) q * (size_t) seq_len + (size_t) k] = 0.0f;
            }
        }
    }
    return mask;
}

// Build the windowed tower graph. mel_in [n_frames, n_mels, 1, 1], pe_chunk_in
// [d_model, chunk_aftercnn] is the positional table shared by every chunk,
// mask_in [S, S] is the block-diagonal mask. chunk_lengths is the host chunk
// plan: every chunk is chunk_mel frames except a possible shorter tail. The
// full chunks run as one batched stem pass [chunk_mel, n_mels, 1, n_full] and
// the tail runs alone behind it, so the graph node count stays constant in the
// audio duration. Returns the projected states [output_dim, S]. stem_out, when
// not null, captures the pre encoder stem sequence [d_model, S_stem] for a
// debug dump.
static struct ggml_tensor * audio_tower_build(struct ggml_context *    ctx,
                                              const ConvStem &         stem,
                                              const AudioEnc &         enc,
                                              struct ggml_tensor *     mel_in,
                                              struct ggml_tensor *     pe_chunk_in,
                                              struct ggml_tensor *     mask_in,
                                              const std::vector<int> & chunk_lengths,
                                              struct ggml_tensor **    stem_out = nullptr) {
    const int64_t n_mels  = mel_in->ne[1];
    const int     d_model = enc.cfg.d_model;

    const int n_chunks  = (int) chunk_lengths.size();
    const int chunk_mel = chunk_lengths[0];
    const int tail_len  = chunk_lengths[(size_t) n_chunks - 1];
    const int n_full    = tail_len == chunk_mel ? n_chunks : n_chunks - 1;

    struct ggml_tensor * seq = nullptr;
    if (n_full > 0) {
        struct ggml_tensor * body =
            ggml_cont(ctx, ggml_view_4d(ctx, mel_in, chunk_mel, n_mels, 1, n_full, mel_in->nb[1], mel_in->nb[2],
                                        (size_t) chunk_mel * mel_in->nb[0], 0));
        struct ggml_tensor * s   = conv_stem_build(ctx, stem, body);  // [d_model, t_c, n_full]
        const int64_t        t_c = s->ne[1];
        struct ggml_tensor * pe  = ggml_view_2d(ctx, pe_chunk_in, d_model, t_c, pe_chunk_in->nb[1], 0);
        s                        = ggml_add(ctx, s, pe);  // pe broadcast over the batch dim
        seq                      = ggml_reshape_2d(ctx, s, d_model, t_c * (int64_t) n_full);
    }
    if (tail_len != chunk_mel) {
        const int64_t        off = (int64_t) n_full * chunk_mel;
        struct ggml_tensor * tail =
            ggml_cont(ctx, ggml_view_4d(ctx, mel_in, tail_len, n_mels, 1, 1, mel_in->nb[1], mel_in->nb[2],
                                        mel_in->nb[3], (size_t) off * mel_in->nb[0]));
        struct ggml_tensor * s   = conv_stem_build(ctx, stem, tail);  // [d_model, t_c, 1]
        const int64_t        t_c = s->ne[1];
        struct ggml_tensor * pe  = ggml_view_2d(ctx, pe_chunk_in, d_model, t_c, pe_chunk_in->nb[1], 0);
        s                        = ggml_add(ctx, s, pe);
        s                        = ggml_reshape_2d(ctx, s, d_model, t_c);
        seq                      = seq ? ggml_concat(ctx, seq, s, 1) : s;
    }

    if (stem_out) {
        *stem_out = seq;
    }
    return audio_enc_build(ctx, enc, seq, mask_in);
}