#pragma once
#include "model_loader.hpp"
#include "graph_builder.hpp"
#include <vector>

struct ggml_context;
struct ggml_tensor;

namespace pk {

// dw_striding (÷8) convolutional subsampling front end of the FastConformer
// encoder (NeMo ConvSubsampling, dw_striding branch). Builds one ggml graph.
class Subsampling {
public:
    explicit Subsampling(const ModelLoader& ml);

    // GRAPH-BUILDER: append the subsampling ops to a SHARED graph `ctx` and
    // return the output tensor [d_model, T'] (ggml ne0=d_model fastest, i.e.
    // row-major [T', d_model]). `mel` is the host mel [n_mels, T] (feat-major
    // inner = T, mel[m*T+t]); it is transposed host-side into `pool`-owned
    // storage and fed as a graph input, so it must outlive the enclosing
    // Backend::compute. `out_Tp`/`out_valid` receive the subsampled time length
    // and valid (non-pad) frame count. Reused by both the fused encoder and the
    // unit test. Uses pk::clone_weight (zero-copy weights) + pk::graph_input_*
    // for host-built masks (registered into `pool` for lifetime).
    ggml_tensor* build_graph(ggml_context* ctx, const std::vector<float>& mel,
                             int n_mels, int T, GraphInputPool& pool,
                             int& out_Tp, int& out_valid,
                             int in_valid_frames = -1) const;
    // Batched GRAPH-BUILDER. `mel` is contiguous [B][n_mels][T] (mel[(b*n_mels+m)*T+t]),
    // `valid_in` holds per-item valid mel frame counts (size B; element <0 means use
    // the offline T-1 convention for that item). Returns [d_model, T', B]
    // (ne0=d_model, ne1=T', ne2=B). `out_valid` (size B) receives each item's
    // non-pad output-frame count.
    ggml_tensor* build_graph_batched(ggml_context* ctx, const float* mel,
                                     int n_mels, int T, int B, GraphInputPool& pool,
                                     int& out_Tp, std::vector<int>& out_valid,
                                     const std::vector<int>& valid_in) const;
    // mel: row-major [n_mels, T] (feat-major inner = T) — i.e. mel[m*T + t].
    // out: row-major [Tout, d_model] (time-major) matching baseline subsampling_out.
    void forward(const std::vector<float>& mel, int n_mels, int T,
                 std::vector<float>& out, int& Tout, int& d_model) const;

    // Same as forward(), but also returns `valid_len`: the number of non-padding
    // output frames after the conv stride reductions (NeMo's masked conv length).
    // Frames in [valid_len, Tout) are center-pad and must be masked downstream
    // (attention key/query masking, depthwise-conv pad masking). This is the
    // value that should be threaded into ConformerLayer::forward.
    void forward(const std::vector<float>& mel, int n_mels, int T,
                 std::vector<float>& out, int& Tout, int& d_model,
                 int& valid_len) const;

    // Streaming variant: `in_valid_frames` is the number of VALID (real, non-pad)
    // input mel frames for this window. The offline forward assumes the
    // preprocessor's center-pad convention (valid = T-1); for cache-aware
    // streaming the whole chunk window is real audio (valid = T, or the
    // clamped chunk length), so the entry valid length must be supplied
    // explicitly instead of using T-1. When in_valid_frames < 0 the offline
    // (T-1) behaviour is used (identical to the overload above).
    void forward(const std::vector<float>& mel, int n_mels, int T,
                 std::vector<float>& out, int& Tout, int& d_model,
                 int& valid_len, int in_valid_frames) const;

    // Tiled subsampling for long audio: result is identical to forward() within the
    // valid region, but no intermediate tensor exceeds a tile_out_frames-bounded size.
    // tile_out_frames = number of OUTPUT (subsampled) frames per tile. Non-causal only;
    // causal falls back to the single-graph path.
    void forward_tiled(const std::vector<float>& mel, int n_mels, int T,
                       int tile_out_frames, std::vector<float>& out,
                       int& Tout, int& d_model, int& valid_len) const;

    // Number of valid (non-pad) output frames for an input of T mel frames,
    // applying the same per-stage `calc_length` reductions NeMo uses. Pure
    // arithmetic, no graph; exposed so the encoder can derive valid_len.
    // `in_valid_frames` (>=0) overrides the offline T-1 entry valid length
    // with an explicit count (streaming); <0 keeps the offline convention.
    int valid_out_len(int T, int in_valid_frames = -1) const;

    // Number of (subsampled) output spatial frames an input of T mel frames
    // produces, applying ggml conv2d's per-stage OH = floor((in+2p-k)/s)+1 for
    // all three stride-2, k=3 stages. Used for subsampling-tile bookkeeping.
    int subsample_len(int T) const;
private:
    const ModelLoader& ml_;
    int conv_channels_;   // C
    int d_model_;         // out features
    bool causal_ = false; // causal_downsampling: left-heavy conv padding (streaming)
};

} // namespace pk