#include "rnnt.hpp"
#include "decode_common.hpp"
#include <cassert>
#include <cmath>

namespace pk {

RnntDecodeState rnnt_decode_init(const PredictionNet& pred) {
    RnntDecodeState st;
    st.state      = pred.zero_state();
    st.last_token = -1;     // -1 sentinel: nothing emitted yet -> SOS.
    st.have_token = false;
    st.hyp.clear();
    return st;
}

std::vector<int32_t> rnnt_decode_frames(const PredictionNet& pred, const Joint& joint,
                                        const std::vector<float>& enc_frames,
                                        int Tnew, int enc_hidden,
                                        RnntDecodeState& st,
                                        int blank_id, int max_symbols,
                                        std::vector<int32_t>* emit_frames,
                                        std::vector<TokenInfo>* tokens) {
    assert((int)enc_frames.size() == (size_t)Tnew * enc_hidden);
    assert(joint.num_durations() == 0);

    const int V_plus      = joint.V_plus();   // vocab + 1 (incl. blank), no durations
    const int token_count = V_plus;           // argmax over the full output vector
    assert(token_count == joint.vocab_size() + 1);

    // Tokens emitted in THIS call only (st.hyp accumulates across all calls).
    std::vector<int32_t> emitted_this_call;

    // Precompute the encoder projection over ALL frames ONCE (one matmul on the
    // persistent backend), reused for every step. The per-step joint below is a
    // tight churn-free graph on the same backend. The old code rebuilt the full
    // joint per (t,u) (a fresh graph per step, the bulk of the per-utterance
    // graph dispatches).
    std::vector<float> enc_proj;   // row-major [Tnew, joint_hidden]
    joint.precompute_enc_proj(enc_frames, Tnew, enc_hidden, enc_proj);
    const int H = joint.joint_hidden();

    // Scratch reused across inner steps.
    std::vector<float> g;
    PredState out_state;
    std::vector<float> logits;

    // Prediction-net output cache. `g` (and the resulting `out_state`) depend
    // ONLY on the committed (last_token, lstm_state), which change exclusively on
    // an emit. The RNN-T loop visits every encoder frame and emits blank on most
    // of them (advancing time without changing the hypothesis), so recomputing
    // the LSTM each frame repeats an identical forward pass. Cache it and only
    // recompute after an emit — this is what NeMo's GreedyRNNTInfer does, and it
    // collapses the LSTM from ~T+U forwards to ~U (the dominant decode cost).
    bool g_valid = false;

    int t = 0;
    while (t < Tnew) {
        int emitted = 0;
        while (emitted < max_symbols) {
            // Prediction net step from the committed state — only when the cache
            // is stale (first step, or the previous iteration emitted). On a
            // blank frame the committed state is unchanged, so `g`/`out_state`
            // from the prior step are reused verbatim.
            if (!g_valid) {
                // First step (no token committed yet, EVER across the whole
                // stream) uses SOS; otherwise feed the last EMITTED token. The
                // committed state / last_token / have_token are carried in `st`.
                const bool is_sos = !st.have_token;
                const int32_t last_label = st.have_token ? st.last_token : blank_id;
                pred.step(last_label, is_sos, st.state, g, out_state);
                g_valid = true;
            }

            // Joint for (t,u): precomputed enc_proj[t] x g -> raw logits [V_plus].
            joint.step_logits(enc_proj.data() + (size_t)t * H,
                              g.data(), (int)g.size(), logits);

            const int k = decode_argmax(logits.data(), token_count);

            // Blank -> stop emitting at this frame and advance time.
            if (k == blank_id) break;

            // Non-blank -> emit, commit state + last token, STAY at this frame.
            st.hyp.push_back((int32_t)k);
            emitted_this_call.push_back((int32_t)k);
            if (emit_frames) emit_frames->push_back((int32_t)t);
            if (tokens) {
                // NeMo per-token metadata (GreedyRNNTInfer._greedy_decode +
                // max_prob confidence): frame = the (local) encoder frame t at
                // emission, conf = max_prob over the full joint output vector
                // (N = V_plus = vocab+1), span = 1 (RNN-T advances one frame).
                const float conf = decode_max_prob_conf(logits.data(), token_count, k);
                tokens->push_back(TokenInfo{ (int32_t)k, (int32_t)t, conf, 1 });
            }
            st.last_token = (int32_t)k;
            st.state = out_state;
            st.have_token = true;
            g_valid = false;   // committed state advanced -> recompute g next step
            emitted += 1;
        }

        // Advance exactly one frame (blank, or max_symbols exhausted).
        t += 1;
    }

    return emitted_this_call;
}

std::vector<int32_t> rnnt_greedy(const PredictionNet& pred, const Joint& joint,
                                 const std::vector<float>& enc, int T, int enc_hidden,
                                 int blank_id, int max_symbols,
                                 std::vector<TokenInfo>* tokens) {
    // The whole-encoder greedy decode is exactly the stateful stepper driven
    // once over all T frames from a fresh state (the loop carries nothing but
    // RnntDecodeState across frames, so chunking is irrelevant to the result).
    RnntDecodeState st = rnnt_decode_init(pred);
    if (tokens) tokens->clear();
    rnnt_decode_frames(pred, joint, enc, T, enc_hidden, st, blank_id, max_symbols,
                       /*emit_frames=*/nullptr, tokens);
    return st.hyp;
}

} // namespace pk