#include "prediction.hpp"
#include "backend.hpp"
#include "ggml_graph.hpp"
#include "ggml.h"
#include "ggml-backend.h"
#include <cassert>
#include <cstring>
#include <string>
#include <vector>

namespace pk {

PredictionNet::PredictionNet(const ModelLoader& ml) : ml_(ml) {
    H_        = (int)ml.config().pred_hidden;
    vocab_p1_ = (int)ml.config().vocab_size + 1;
    n_layers_ = (int)ml.config().pred_rnn_layers;
    if (n_layers_ <= 0) n_layers_ = 1; // default to a single LSTM layer
    assert(H_ > 0 && "pred_hidden not set");
}

// ---------------------------------------------------------------------------
// Stateful helpers.
// ---------------------------------------------------------------------------
PredState PredictionNet::zero_state() const {
    PredState s;
    s.h.assign((size_t)n_layers_, std::vector<float>((size_t)H_, 0.0f));
    s.c.assign((size_t)n_layers_, std::vector<float>((size_t)H_, 0.0f));
    return s;
}

// ---------------------------------------------------------------------------
// Advance the stacked LSTM by one token, as a single ggml graph that runs on
// whatever backend pk::Backend selected (CPU or device). The embedding table is
// fetched to the host once (ggml_backend_tensor_get, works for device tensors)
// so the SOS/lookup row can seed the layer-0 input; the LSTM weights are
// referenced as zero-copy loader leaves via clone_weight. Each layer's new
// (h', c') is captured for out_state; the top layer's h' is the output g.
//
// LSTM math (must match the prior C++ exactly):
//   z = W_ih·x + b_ih + W_hh·h_in + b_hh                          [4H]
//   i=sigmoid(z[0:H]); f=sigmoid(z[H:2H]); g=tanh(z[2H:3H]); o=sigmoid(z[3H:4H])
//   c' = f*c_in + i*g ;  h' = o*tanh(c')
// ---------------------------------------------------------------------------
void PredictionNet::step(int32_t token_id, bool is_sos,
                          const PredState& in,
                          std::vector<float>& g,
                          PredState& out_state) const {
    const int H = H_;
    const int L = n_layers_;

    // Lazily fetch the embedding table to the host (device-safe). Ensure the
    // loader's weights have a backend buffer first (idempotent) so the tensor
    // is readable via ggml_backend_tensor_get even when step()/forward() is
    // exercised before the encoder graph has realized the weights.
    if (embed_host_.empty()) {
        pk::ensure_weights_realized(ml_);
        ggml_tensor* emb = ml_.tensor("decoder.prediction.embed.weight");
        assert(emb && "missing decoder.prediction.embed.weight");
        embed_host_.resize((size_t)vocab_p1_ * H);
        ggml_backend_tensor_get(emb, embed_host_.data(), 0, ggml_nbytes(emb));
    }

    // Layer-0 input: zeros for SOS, else the embedding row for token_id.
    std::vector<float> x0((size_t)H, 0.0f);
    if (!is_sos) {
        assert(token_id >= 0 && token_id < vocab_p1_ && "embedding id out of range");
        std::memcpy(x0.data(), &embed_host_[(size_t)token_id * H],
                    (size_t)H * sizeof(float));
    }

    out_state.h.assign((size_t)L, std::vector<float>((size_t)H));
    out_state.c.assign((size_t)L, std::vector<float>((size_t)H));

    bool ok = pk::run_graph(0, 0, [&](ggml_context* ctx) -> ggml_tensor* {
        int64_t ne1[1] = { H };
        ggml_tensor* layer_in = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 1, ne1,
                                    x0.data(), (size_t)H * sizeof(float));
        ggml_tensor* top_h = nullptr;
        for (int l = 0; l < L; ++l) {
            const std::string s = "_l" + std::to_string(l);
            ggml_tensor* Wih = pk::clone_weight(ctx, ml_,
                ("decoder.prediction.dec_rnn.lstm.weight_ih" + s).c_str());
            ggml_tensor* Whh = pk::clone_weight(ctx, ml_,
                ("decoder.prediction.dec_rnn.lstm.weight_hh" + s).c_str());
            ggml_tensor* bih = pk::clone_weight(ctx, ml_,
                ("decoder.prediction.dec_rnn.lstm.bias_ih" + s).c_str());
            ggml_tensor* bhh = pk::clone_weight(ctx, ml_,
                ("decoder.prediction.dec_rnn.lstm.bias_hh" + s).c_str());
            ggml_tensor* h_in = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 1, ne1,
                                    in.h[l].data(), (size_t)H * sizeof(float));
            ggml_tensor* c_in = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 1, ne1,
                                    in.c[l].data(), (size_t)H * sizeof(float));
            // z = W_ih·x + b_ih + W_hh·h_in + b_hh                       [4H]
            ggml_tensor* z = ggml_add(ctx,
                ggml_add(ctx, ggml_mul_mat(ctx, Wih, layer_in), bih),
                ggml_add(ctx, ggml_mul_mat(ctx, Whh, h_in),     bhh));
            // Gate slices (i, f, g, o), each [H], contiguous for elementwise ops.
            ggml_tensor* i  = ggml_sigmoid(ctx, ggml_cont(ctx, ggml_view_1d(ctx, z, H, 0)));
            ggml_tensor* f  = ggml_sigmoid(ctx, ggml_cont(ctx, ggml_view_1d(ctx, z, H, (size_t)H * sizeof(float))));
            ggml_tensor* gg = ggml_tanh   (ctx, ggml_cont(ctx, ggml_view_1d(ctx, z, H, (size_t)2 * H * sizeof(float))));
            ggml_tensor* o  = ggml_sigmoid(ctx, ggml_cont(ctx, ggml_view_1d(ctx, z, H, (size_t)3 * H * sizeof(float))));
            // c' = f*c_in + i*g ;  h' = o*tanh(c')
            ggml_tensor* c_out = ggml_add(ctx, ggml_mul(ctx, f, c_in), ggml_mul(ctx, i, gg));
            ggml_tensor* h_out = ggml_mul(ctx, o, ggml_tanh(ctx, c_out));
            pk::capture_graph_output(c_out, &out_state.c[l]);
            pk::capture_graph_output(h_out, &out_state.h[l]);
            layer_in = h_out;
            top_h    = h_out;
        }
        return top_h;
    }, g);
    assert(ok && "pred-net step graph failed");
}

// ---------------------------------------------------------------------------
// Batched single-step advance: the same LSTM math as step(), but with a batch
// axis N. Inputs and state are laid out [H, N] in ggml (item n is column n,
// offset n*H in the flat host buffer). Gate slices become [H, N] views into the
// [4H, N] z, using z->nb[1] as the column stride. N=1 reduces to step().
// ---------------------------------------------------------------------------
void PredictionNet::step_batch(const std::vector<int32_t>& token_ids,
                               const std::vector<uint8_t>& is_sos,
                               const BatchedPredState& in,
                               std::vector<float>& g,
                               BatchedPredState& out_state) const {
    const int H = H_;
    const int L = n_layers_;
    const int N = (int)token_ids.size();
    assert(N > 0 && (int)is_sos.size() == N && "batch size mismatch");

    // Lazily fetch the embedding table to the host (device-safe), exactly as
    // step() does.
    if (embed_host_.empty()) {
        pk::ensure_weights_realized(ml_);
        ggml_tensor* emb = ml_.tensor("decoder.prediction.embed.weight");
        assert(emb && "missing decoder.prediction.embed.weight");
        embed_host_.resize((size_t)vocab_p1_ * H);
        ggml_backend_tensor_get(emb, embed_host_.data(), 0, ggml_nbytes(emb));
    }

    // Layer-0 input [H*N]: zeros for SOS items, else the embedding row.
    std::vector<float> x0((size_t)H * N, 0.0f);
    for (int n = 0; n < N; ++n) {
        if (!is_sos[n]) {
            assert(token_ids[n] >= 0 && token_ids[n] < vocab_p1_ && "embedding id out of range");
            std::memcpy(&x0[(size_t)n * H], &embed_host_[(size_t)token_ids[n] * H],
                        (size_t)H * sizeof(float));
        }
    }

    out_state.h.assign((size_t)L, std::vector<float>((size_t)H * N));
    out_state.c.assign((size_t)L, std::vector<float>((size_t)H * N));

    bool ok = pk::run_graph(0, 0, [&](ggml_context* ctx) -> ggml_tensor* {
        int64_t ne2[2] = { H, N };
        ggml_tensor* layer_in = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, ne2,
                                    x0.data(), (size_t)H * N * sizeof(float));
        ggml_tensor* top_h = nullptr;
        for (int l = 0; l < L; ++l) {
            const std::string s = "_l" + std::to_string(l);
            ggml_tensor* Wih = pk::clone_weight(ctx, ml_,
                ("decoder.prediction.dec_rnn.lstm.weight_ih" + s).c_str());
            ggml_tensor* Whh = pk::clone_weight(ctx, ml_,
                ("decoder.prediction.dec_rnn.lstm.weight_hh" + s).c_str());
            ggml_tensor* bih = pk::clone_weight(ctx, ml_,
                ("decoder.prediction.dec_rnn.lstm.bias_ih" + s).c_str());
            ggml_tensor* bhh = pk::clone_weight(ctx, ml_,
                ("decoder.prediction.dec_rnn.lstm.bias_hh" + s).c_str());
            ggml_tensor* h_in = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, ne2,
                                    in.h[l].data(), (size_t)H * N * sizeof(float));
            ggml_tensor* c_in = pk::graph_input_tensor(ctx, GGML_TYPE_F32, 2, ne2,
                                    in.c[l].data(), (size_t)H * N * sizeof(float));
            // z = W_ih·x + b_ih + W_hh·h_in + b_hh                  [4H, N]
            // (bias [4H] broadcasts over the N columns).
            ggml_tensor* z = ggml_add(ctx,
                ggml_add(ctx, ggml_mul_mat(ctx, Wih, layer_in), bih),
                ggml_add(ctx, ggml_mul_mat(ctx, Whh, h_in),     bhh));
            // Gate slices (i, f, g, o), each [H, N]. The view keeps z's FULL
            // column stride (z->nb[1] = 4H elems), reading only H contiguous
            // elements per column, so consecutive columns skip the other three
            // gate blocks. (Do NOT change the stride to H*sizeof(float).)
            ggml_tensor* i  = ggml_sigmoid(ctx, ggml_cont(ctx, ggml_view_2d(ctx, z, H, N, z->nb[1], 0)));
            ggml_tensor* f  = ggml_sigmoid(ctx, ggml_cont(ctx, ggml_view_2d(ctx, z, H, N, z->nb[1], (size_t)H * sizeof(float))));
            ggml_tensor* gg = ggml_tanh   (ctx, ggml_cont(ctx, ggml_view_2d(ctx, z, H, N, z->nb[1], (size_t)2 * H * sizeof(float))));
            ggml_tensor* o  = ggml_sigmoid(ctx, ggml_cont(ctx, ggml_view_2d(ctx, z, H, N, z->nb[1], (size_t)3 * H * sizeof(float))));
            // c' = f*c_in + i*g ;  h' = o*tanh(c')
            ggml_tensor* c_out = ggml_add(ctx, ggml_mul(ctx, f, c_in), ggml_mul(ctx, i, gg));
            ggml_tensor* h_out = ggml_mul(ctx, o, ggml_tanh(ctx, c_out));
            pk::capture_graph_output(c_out, &out_state.c[l]);
            pk::capture_graph_output(h_out, &out_state.h[l]);
            layer_in = h_out;
            top_h    = h_out;
        }
        return top_h;
    }, g);
    assert(ok && "pred-net step_batch graph failed");
}

// ---------------------------------------------------------------------------
// Full-sequence forward pass (unchanged API; now driven by step() so there is a
// single LSTM implementation). Carries (h, c) state across timesteps; the
// output at step t is the top layer's h'. add_sos prepends a zero SOS step.
// ---------------------------------------------------------------------------
void PredictionNet::forward(const std::vector<int32_t>& ids, bool add_sos,
                            std::vector<float>& out, int& U_out, int& hidden) const {
    const int H   = H_;
    const int seq = add_sos ? (int)ids.size() + 1 : (int)ids.size();
    U_out  = seq;
    hidden = H;
    out.assign((size_t)seq * H, 0.0f);

    PredState st = zero_state();
    std::vector<float> g;
    PredState nxt;
    for (int t = 0; t < seq; ++t) {
        const bool is_sos = add_sos && t == 0;
        const int32_t tok = is_sos ? -1 : ids[add_sos ? t - 1 : t];
        step(tok, is_sos, st, g, nxt);
        std::memcpy(&out[(size_t)t * H], g.data(), (size_t)H * sizeof(float));
        st = nxt;
    }
}

} // namespace pk