#include "llama.h"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <limits>
#include <string>
#include <vector>

namespace {

struct Score {
    int index;
    int input_index;
    int target;
    int top1;
    double logprob;
};

bool read_tokens(const char * path, std::vector<llama_token> & tokens) {
    std::ifstream in(path);
    if (!in) {
        std::cerr << "cannot open token file: " << path << "\n";
        return false;
    }
    long long value = 0;
    while (in >> value) {
        if (value < std::numeric_limits<int32_t>::min() || value > std::numeric_limits<int32_t>::max()) {
            std::cerr << "token out of int32 range\n";
            return false;
        }
        tokens.push_back(static_cast<llama_token>(value));
    }
    if (!in.eof()) {
        std::cerr << "malformed token file\n";
        return false;
    }
    return true;
}

bool score_chunk(llama_context * ctx, const std::vector<llama_token> & tokens, int start,
                 int n_vocab, std::vector<Score> & scores, std::string & error) {
    constexpr int chunk = 512;
    llama_batch batch = llama_batch_init(chunk, 0, 1);
    if (!batch.token || !batch.pos || !batch.n_seq_id || !batch.seq_id || !batch.logits) {
        llama_batch_free(batch);
        error = "llama_batch_init returned incomplete batch";
        return false;
    }
    for (int i = 0; i < chunk; ++i) {
        batch.token[i] = tokens[start + i];
        batch.pos[i] = i;
        batch.n_seq_id[i] = 1;
        batch.seq_id[i][0] = 0;
        batch.logits[i] = 1;
    }
    batch.n_tokens = chunk;

    const int rc = llama_decode(ctx, batch);
    if (rc != 0) {
        error = "llama_decode returned " + std::to_string(rc);
        llama_batch_free(batch);
        return false;
    }

    for (int i = 0; i < chunk; ++i) {
        const float * logits = llama_get_logits_ith(ctx, i);
        if (!logits) {
            error = "llama_get_logits_ith returned null at row " + std::to_string(i);
            llama_batch_free(batch);
            return false;
        }
        float max_logit = -std::numeric_limits<float>::infinity();
        int top1 = -1;
        for (int v = 0; v < n_vocab; ++v) {
            if (logits[v] > max_logit) {
                max_logit = logits[v];
                top1 = v;
            }
        }
        double sum_exp = 0.0;
        for (int v = 0; v < n_vocab; ++v) {
            sum_exp += std::exp(static_cast<double>(logits[v]) - static_cast<double>(max_logit));
        }
        const int target = tokens[start + i + 1];
        if (target < 0 || target >= n_vocab || !(sum_exp > 0.0)) {
            error = "invalid target or logit normalization at row " + std::to_string(i);
            llama_batch_free(batch);
            return false;
        }
        const double logprob = static_cast<double>(logits[target]) - static_cast<double>(max_logit) - std::log(sum_exp);
        scores.push_back({start + i + 1, start + i, target, top1, logprob});
    }
    llama_batch_free(batch);
    return true;
}

} // namespace

int main(int argc, char ** argv) {
    if (argc != 4) {
        std::cerr << "usage: llama-teacher-score MODEL TOKEN_IDS OUTPUT_JSONL\n";
        return 2;
    }

    std::vector<llama_token> tokens;
    if (!read_tokens(argv[2], tokens) || tokens.size() < 1025) {
        std::cerr << "need at least 1025 token ids, got " << tokens.size() << "\n";
        return 2;
    }

    llama_backend_init();
    llama_model_params model_params = llama_model_default_params();
    model_params.n_gpu_layers = -1;
    llama_model * model = llama_model_load_from_file(argv[1], model_params);
    if (!model) {
        std::cerr << "llama_model_load_from_file failed\n";
        llama_backend_free();
        return 3;
    }

    llama_context_params context_params = llama_context_default_params();
    context_params.n_ctx = 512;
    context_params.n_batch = 512;
    context_params.n_ubatch = 512;
    context_params.n_seq_max = 1;
    context_params.n_outputs_max = 512;
    context_params.n_outputs_max_per_seq = 512;
    context_params.n_threads = 16;
    context_params.n_threads_batch = 48;
    context_params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO;
    context_params.op_offload = true;
    context_params.kv_unified = false;
    llama_context * ctx = llama_init_from_model(model, context_params);
    if (!ctx) {
        std::cerr << "llama_init_from_model failed\n";
        llama_model_free(model);
        llama_backend_free();
        return 4;
    }

    const llama_vocab * vocab = llama_model_get_vocab(model);
    const int n_vocab = llama_vocab_n_tokens(vocab);
    std::vector<Score> scores;
    scores.reserve(1024);
    std::string error;
    const auto started = std::chrono::steady_clock::now();
    for (int start : {0, 512}) {
        if (start != 0) {
            llama_memory_clear(llama_get_memory(ctx), true);
        }
        if (!score_chunk(ctx, tokens, start, n_vocab, scores, error)) {
            std::cerr << error << "\n";
            llama_free(ctx);
            llama_model_free(model);
            llama_backend_free();
            return 5;
        }
    }
    llama_synchronize(ctx);
    const auto elapsed = std::chrono::duration<double>(std::chrono::steady_clock::now() - started).count();

    double nll = 0.0;
    int top1_correct = 0;
    for (const Score & s : scores) {
        nll -= s.logprob;
        top1_correct += s.top1 == s.target;
    }
    const double mean_nll = nll / static_cast<double>(scores.size());
    const double ppl = std::exp(mean_nll);

    std::ofstream out(argv[3]);
    if (!out) {
        std::cerr << "cannot open output: " << argv[3] << "\n";
        llama_free(ctx);
        llama_model_free(model);
        llama_backend_free();
        return 6;
    }
    out << std::setprecision(12);
    out << "{\"type\":\"meta\",\"commit\":\"9a9394a895b96003ca842a6041cb28ac49a108f7\","
           "\"vocab\":" << n_vocab << ",\"token_count\":" << tokens.size()
        << ",\"chunk_size\":512,\"chunk_starts\":[0,512],\"bos_injected\":false,"
           "\"score_count\":" << scores.size() << ",\"elapsed_seconds\":" << elapsed << "}\n";
    for (const Score & s : scores) {
        out << "{\"type\":\"score\",\"target_index\":" << s.index
            << ",\"input_index\":" << s.input_index
            << ",\"target\":" << s.target
            << ",\"correct_logprob\":" << s.logprob
            << ",\"top1\":" << s.top1 << "}\n";
    }
    out << "{\"type\":\"summary\",\"score_count\":" << scores.size()
        << ",\"mean_nll\":" << mean_nll
        << ",\"ppl\":" << ppl
        << ",\"top1_correct\":" << top1_correct
        << ",\"top1_accuracy\":" << (static_cast<double>(top1_correct) / scores.size())
        << ",\"elapsed_seconds\":" << elapsed << "}\n";
    out.close();

    llama_free(ctx);
    llama_model_free(model);
    llama_backend_free();
    return 0;
}