#include "llama-kv-cache-msa.h"

#include "llama-impl.h"
#include "llama-batch.h"
#include "llama-model.h"

#include <algorithm>
#include <cassert>
#include <cmath>

// llama_kv_cache_msa

llama_kv_cache_msa::llama_kv_cache_msa(
        const llama_model & model,
                ggml_type   type_k,
                ggml_type   type_v,
                     bool   v_trans,
                     bool   offload,
                     bool   unified,
                 uint32_t   kv_size,
                 uint32_t   n_seq_max,
                 uint32_t   n_pad,
                 uint32_t   n_swa,
           llama_swa_type   swa_type,
    const layer_filter_cb & filter,
    const layer_filter_cb & filter_idx,
    const  layer_reuse_cb & reuse) :
    hparams_idx(model.hparams),
    n_stream(unified ? 1 : n_seq_max), n_seq_max(n_seq_max), n_pad(n_pad),
    n_swa(n_swa), swa_type(swa_type) {

    LLAMA_LOG_INFO("%s: creating main KV cache, size = %u cells\n", __func__, kv_size);

    kv_base = std::make_unique<llama_kv_cache>(
            model, model.hparams, type_k, type_v,
            v_trans, offload, unified, kv_size, n_seq_max, n_pad,
            n_swa, swa_type, nullptr, filter, reuse, nullptr);

    // the MSA indexer uses a single key head per layer
    std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1);
    hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size;
    // the rope parameters are kept identical to the main cache

    LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size);

    kv_idx = std::make_unique<llama_kv_cache>(
            model, hparams_idx, type_k, type_v,
            v_trans, offload, unified, kv_size, n_seq_max, n_pad,
            n_swa, swa_type, nullptr, filter_idx, reuse, nullptr);
}

void llama_kv_cache_msa::clear(bool data) {
    kv_base->clear(data);
    kv_idx ->clear(data);
}

bool llama_kv_cache_msa::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
    bool res = true;

    res = res & kv_base->seq_rm(seq_id, p0, p1);
    res = res & kv_idx ->seq_rm(seq_id, p0, p1);

    return res;
}

void llama_kv_cache_msa::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
    kv_base->seq_cp(seq_id_src, seq_id_dst, p0, p1);
    kv_idx ->seq_cp(seq_id_src, seq_id_dst, p0, p1);
}

void llama_kv_cache_msa::seq_keep(llama_seq_id seq_id) {
    kv_base->seq_keep(seq_id);
    kv_idx ->seq_keep(seq_id);
}

void llama_kv_cache_msa::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
    kv_base->seq_add(seq_id, p0, p1, shift);
    kv_idx ->seq_add(seq_id, p0, p1, shift);
}

void llama_kv_cache_msa::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
    kv_base->seq_div(seq_id, p0, p1, d);
    kv_idx ->seq_div(seq_id, p0, p1, d);
}

llama_pos llama_kv_cache_msa::seq_pos_min(llama_seq_id seq_id) const {
    return kv_base->seq_pos_min(seq_id);
}

llama_pos llama_kv_cache_msa::seq_pos_max(llama_seq_id seq_id) const {
    return kv_base->seq_pos_max(seq_id);
}

std::map<ggml_backend_buffer_type_t, size_t> llama_kv_cache_msa::memory_breakdown() const {
    std::map<ggml_backend_buffer_type_t, size_t> mb = kv_base->memory_breakdown();
    for (const auto & buft_size : kv_idx->memory_breakdown()) {
        mb[buft_size.first] += buft_size.second;
    }
    return mb;
}

llama_memory_context_ptr llama_kv_cache_msa::init_batch(
            llama_batch_allocr & balloc,
            uint32_t n_ubatch,
            bool embd_all) {
    GGML_UNUSED(embd_all);

    do {
        balloc.split_reset();

        std::vector<llama_ubatch> ubatches;
        while (true) {
            auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0);

            if (ubatch.n_tokens == 0) {
                break;
            }

            ubatches.push_back(std::move(ubatch));
        }

        if (balloc.get_n_used() < balloc.get_n_tokens()) {
            // failed to find a suitable split
            break;
        }

        auto sinfos_base = kv_base->prepare(ubatches);
        if (sinfos_base.empty()) {
            break;
        }

        auto sinfos_idx = kv_idx->prepare(ubatches);
        if (sinfos_idx.empty()) {
            break;
        }

        assert(sinfos_base.size() == sinfos_idx.size());

        return std::make_unique<llama_kv_cache_msa_context>(
                this, std::move(sinfos_base), std::move(sinfos_idx), std::move(ubatches));
    } while (false);

    return std::make_unique<llama_kv_cache_msa_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
}

llama_memory_context_ptr llama_kv_cache_msa::init_full() {
    return std::make_unique<llama_kv_cache_msa_context>(this);
}

llama_memory_context_ptr llama_kv_cache_msa::init_update(llama_context * lctx, bool optimize) {
    return std::make_unique<llama_kv_cache_msa_context>(this, lctx, optimize);
}

bool llama_kv_cache_msa::get_can_shift() const {
    return kv_base->get_can_shift() &&
           kv_idx ->get_can_shift() &&
           kv_base->get_size() == kv_idx->get_size();
}

void llama_kv_cache_msa::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {
    kv_base->state_write(io, seq_id, flags);
    kv_idx ->state_write(io, seq_id, flags);
}

void llama_kv_cache_msa::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
    kv_base->state_read(io, seq_id, flags);
    kv_idx ->state_read(io, seq_id, flags);
}

llama_kv_cache * llama_kv_cache_msa::get_base() const {
    return kv_base.get();
}

llama_kv_cache * llama_kv_cache_msa::get_idx() const {
    return kv_idx.get();
}

// llama_kv_cache_msa_context

llama_kv_cache_msa_context::llama_kv_cache_msa_context(llama_memory_status status) :
    kv(nullptr), status(status) {}

llama_kv_cache_msa_context::llama_kv_cache_msa_context(
        llama_kv_cache_msa * kv) :
    kv(kv),
    ctx_base(kv->get_base()->init_full()),
    ctx_idx (kv->get_idx ()->init_full()),
    status(llama_memory_status_combine(ctx_base->get_status(), ctx_idx->get_status())) {
}

llama_kv_cache_msa_context::llama_kv_cache_msa_context(
        llama_kv_cache_msa * kv,
        llama_context * lctx,
        bool optimize) :
    kv(kv),
    ctx_base(kv->get_base()->init_update(lctx, optimize)),
    ctx_idx (kv->get_idx ()->init_update(lctx, optimize)),
    status(llama_memory_status_combine(ctx_base->get_status(), ctx_idx->get_status())) {
}

llama_kv_cache_msa_context::llama_kv_cache_msa_context(
        llama_kv_cache_msa * kv,
        slot_info_vec_t sinfos_base,
        slot_info_vec_t sinfos_idx,
        std::vector<llama_ubatch> ubatches) :
    kv(kv),
    ubatches(std::move(ubatches)),
    // here we copy the ubatches. not sure if this is ideal
    ctx_base(new llama_kv_cache_context(kv->get_base(), std::move(sinfos_base), this->ubatches)),
    ctx_idx (new llama_kv_cache_context(kv->get_idx (), std::move(sinfos_idx),  this->ubatches)),
    status(llama_memory_status_combine(ctx_base->get_status(), ctx_idx->get_status())) {
}

llama_kv_cache_msa_context::~llama_kv_cache_msa_context() = default;

bool llama_kv_cache_msa_context::next() {
    assert(status == LLAMA_MEMORY_STATUS_SUCCESS);

    ctx_base->next();
    ctx_idx ->next();

    if (++i_next >= ubatches.size()) {
        return false;
    }

    return true;
}

bool llama_kv_cache_msa_context::apply() {
    assert(!llama_memory_status_is_fail(status));

    bool res = true;

    res = res & ctx_base->apply();
    res = res & ctx_idx ->apply();

    return res;
}

llama_memory_status llama_kv_cache_msa_context::get_status() const {
    return status;
}

const llama_ubatch & llama_kv_cache_msa_context::get_ubatch() const {
    assert(status == LLAMA_MEMORY_STATUS_SUCCESS);

    return ubatches[i_next];
}

const llama_kv_cache_context * llama_kv_cache_msa_context::get_base() const {
    assert(status == LLAMA_MEMORY_STATUS_SUCCESS);

    return static_cast<const llama_kv_cache_context *>(ctx_base.get());
}

const llama_kv_cache_context * llama_kv_cache_msa_context::get_idx() const {
    assert(status == LLAMA_MEMORY_STATUS_SUCCESS);

    return static_cast<const llama_kv_cache_context *>(ctx_idx.get());
}

uint32_t llama_kv_cache_msa_context::get_n_pos() const {
    // pad the value so that the graph remains constant across batches and can be reused
    const uint32_t n_pad_cur = std::max(kv->get_n_pad(), 256u);

    llama_pos pos_max = -1;

    for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) kv->get_n_seq_max(); ++seq_id) {
        pos_max = std::max(pos_max, kv->seq_pos_max(seq_id));
    }

    return std::max(n_pad_cur, GGML_PAD((uint32_t) (pos_max + 1), n_pad_cur));
}

void llama_kv_cache_msa_context::set_input_cell_pos(ggml_tensor * dst, const llama_ubatch * ubatch, int32_t div) const {
    GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
    GGML_ASSERT(dst->type == GGML_TYPE_I32);
    GGML_ASSERT(div > 0);

    const int64_t n_tokens    = ubatch->n_tokens;
    const int64_t n_kv        = dst->ne[0];
    const int64_t n_stream_ub = dst->ne[1];

    GGML_ASSERT(n_tokens % n_stream_ub == 0);
    const int64_t n_tps = n_tokens/n_stream_ub;

    int32_t * data = (int32_t *) dst->data;

    for (int64_t s = 0; s < n_stream_ub; ++s) {
        const llama_seq_id seq_id = ubatch->seq_id[s*n_tps][0];

        const auto & cells = kv->get_base()->get_cells(seq_id);

        for (int64_t j = 0; j < n_kv; ++j) {
            // the value for empty or other-sequence cells is irrelevant as consumers mask them
            data[s*n_kv + j] =
                cells.is_empty(j) || !cells.seq_has(j, seq_id)
                    ? 0
                    : (int32_t) (cells.pos_get(j)/div);
        }
    }
}

void llama_kv_cache_msa_context::set_input_pos_slot(ggml_tensor * dst, const llama_ubatch * ubatch) const {
    GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
    GGML_ASSERT(dst->type == GGML_TYPE_I32 || dst->type == GGML_TYPE_F32);

    const int64_t n_tokens    = ubatch->n_tokens;
    const int64_t n_pos       = dst->ne[0];
    const int64_t n_stream_ub = dst->ne[1];

    GGML_ASSERT(n_tokens % n_stream_ub == 0);
    const int64_t n_tps = n_tokens/n_stream_ub;

    for (int64_t s = 0; s < n_stream_ub; ++s) {
        const llama_seq_id seq_id = ubatch->seq_id[s*n_tps][0];

        const auto & cells = kv->get_base()->get_cells(seq_id);

        std::vector<int32_t> map(n_pos, 0);

        for (uint32_t j = 0; j < cells.size(); ++j) {
            if (cells.is_empty(j) || !cells.seq_has(j, seq_id)) {
                continue;
            }

            const llama_pos p0 = cells.pos_get(j);

            if (p0 < 0 || p0 >= n_pos) {
                continue;
            }

            map[p0] = (int32_t) j;
        }

        if (dst->type == GGML_TYPE_I32) {
            int32_t * data = (int32_t *) dst->data + s*n_pos;
            std::copy(map.begin(), map.end(), data);
        } else {
            float * data = (float *) dst->data + s*n_pos;
            for (int64_t p = 0; p < n_pos; ++p) {
                data[p] = (float) map[p];
            }
        }
    }
}

void llama_kv_cache_msa_context::set_input_pos_mask(ggml_tensor * dst, const llama_ubatch * ubatch) const {
    GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer));
    GGML_ASSERT(dst->type == GGML_TYPE_F32);

    const int64_t n_tokens = ubatch->n_tokens;
    const int64_t n_pos    = dst->ne[0];

    GGML_ASSERT(dst->ne[1] == n_tokens);

    const uint32_t       n_swa    = kv->get_n_swa();
    const llama_swa_type swa_type = kv->get_swa_type();

    float * data = (float *) dst->data;

    std::fill(data, data + n_pos*n_tokens, -INFINITY);

    for (int64_t i = 0; i < n_tokens; ++i) {
        const llama_seq_id seq_id = ubatch->seq_id[i][0];

        const auto & cells = kv->get_base()->get_cells(seq_id);

        const llama_pos p1 = ubatch->pos[i];

        for (uint32_t j = 0; j < cells.size(); ++j) {
            if (cells.is_empty(j) || !cells.seq_has(j, seq_id)) {
                continue;
            }

            const llama_pos p0 = cells.pos_get(j);

            if (p0 < 0 || p0 >= n_pos) {
                continue;
            }

            // causal mask
            if (p0 > p1) {
                continue;
            }

            // apply SWA if any
            if (llama_hparams::is_masked_swa(n_swa, swa_type, p0, p1)) {
                continue;
            }

            data[i*n_pos + p0] = 0.0f;
        }
    }
}