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

namespace pk {

// RNN-Transducer joint network — NeMo RNNTJoint.
//
// Architecture (NeMo RNNTJoint.joint):
//   enc_proj[t]  = joint.enc.weight  · enc[t]  + joint.enc.bias   (enc_hidden→joint_hidden)
//   pred_proj[u] = joint.pred.weight · pred[u] + joint.pred.bias  (pred_hidden→joint_hidden)
//   f[t,u]       = ReLU(enc_proj[t] + pred_proj[u])               (broadcast sum, ReLU on sum)
//   logits[t,u]  = joint_net.2.weight · f[t,u] + joint_net.2.bias (joint_hidden→V_plus)
//
// Output logits are RAW (no log_softmax), shape [T, U, V_plus] row-major.
// V_plus = vocab + 1 + num_durations  (full TDT output vector).
//
// enc input convention: row-major [T, enc_hidden], i.e. enc[t*enc_hidden + c].
//   (The Encoder outputs channels-first [enc_hidden, T]; the caller must
//    transpose before calling this function.)
//
// Tensors (verbatim NeMo names, ggml ne = reverse of torch):
//   joint.enc.weight          ggml ne=[enc_hidden=512, joint_hidden=640]
//   joint.enc.bias            ggml ne=[640]
//   joint.pred.weight         ggml ne=[pred_hidden=640, joint_hidden=640]
//   joint.pred.bias           ggml ne=[640]
//   joint.joint_net.2.weight  ggml ne=[joint_hidden=640, V_plus=1030]
//   joint.joint_net.2.bias    ggml ne=[V_plus=1030]
class Joint {
public:
    explicit Joint(const ModelLoader& ml);

    // enc:  row-major [T, enc_hidden],  enc[t*enc_hidden + c]
    // pred: row-major [U, pred_hidden], pred[u*pred_hidden + h]
    // logits out: row-major [T, U, V_plus], logits[t*U*V_plus + u*V_plus + v]
    // V_plus = vocab + 1 + num_durations (written to the out param)
    void forward(const std::vector<float>& enc,  int T, int enc_hidden,
                 const std::vector<float>& pred, int U, int pred_hidden,
                 std::vector<float>& logits, int& V_plus) const;

    // --- Churn-free greedy-decode fast path (used by rnnt/tdt greedy loops) ---
    //
    // The greedy decoders call the joint once per emitted (t,u) step. The old
    // `forward(T=1,U=1,...)` rebuilt the FULL encoder projection per step (two
    // graphs/step, ~187 dispatches/utterance). Instead, precompute the encoder
    // projection over ALL T frames ONCE (one matmul, reused for every u), then
    // run a tight per-step joint (pred projection + ReLU + output projection) on
    // the PERSISTENT backend. Backend::compute reuses the backend + gallocr, so
    // the per-step graphs are churn-free; ggml's multithreaded matmul wins here
    // because the per-step joint is memory-bandwidth bound on the ~2.6 MB output
    // weight (measured ~15x faster/step than a single-threaded C++ matvec).
    // Quantized enc/pred weights are dequantized on the fly by ggml_mul_mat.

    // enc_proj[t][h] = enc.weight·enc[t] + enc.bias, over all T frames.
    // enc: row-major [T, enc_hidden]; out: row-major [T, joint_hidden].
    void precompute_enc_proj(const std::vector<float>& enc, int T, int enc_hidden,
                             std::vector<float>& enc_proj) const;

    // Per-step joint logits for a single (t,u):
    //   pred_proj[h] = pred.weight·g + pred.bias
    //   f[h]         = ReLU(enc_proj_t[h] + pred_proj[h])
    //   logits[v]    = joint_net.2.weight·f + joint_net.2.bias    → [V_plus]
    // enc_proj_t: pointer to enc_proj row for frame t ([joint_hidden]).
    // g:          prediction-net output for step u ([pred_hidden]).
    void step_logits(const float* enc_proj_t,
                     const float* g, int pred_hidden,
                     std::vector<float>& logits) const;

    // Batched per-step joint for N items.
    // enc_proj_gathered: [joint_hidden * N], each item's enc_proj row for its
    //   current frame (item n at offset n*joint_hidden).
    // g:                 [pred_hidden * N], batched pred output (item n at n*pred_hidden).
    // logits out:        [V_plus * N], item n at offset n*V_plus.
    void step_logits_batch(const float* enc_proj_gathered,
                           const float* g, int pred_hidden, int n,
                           std::vector<float>& logits) const;

    int joint_hidden() const { return joint_hidden_; }

    // V_plus = vocab_size + 1 + num_durations
    int V_plus()       const { return V_plus_; }
    int vocab_size()   const { return vocab_size_; }
    int num_durations() const { return num_durations_; }

private:
    const ModelLoader& ml_;
    int joint_hidden_;
    int vocab_size_;
    int num_durations_;
    int V_plus_;
    // Cached for shape asserts on the projection inputs.
    int enc_hidden_  = 0;   // ggml ne[0] of joint.enc.weight
    int pred_hidden_ = 0;   // ggml ne[0] of joint.pred.weight
};

} // namespace pk