#include "models.h"

#include "llama-kv-cache.h"

// Cosmos: 54D mixture-of-states Hebbian attention.
//
//   x54      = W54 . h
//   H_ij     = exp(-||x54_i - x54_j||^2 / 2*sigma^2), causally masked and row-normalised
//   A_final  = (1-g)*A_std + g*H
//
// Two identities keep this inside the existing attention infrastructure:
//
// 1. dropping the sigma tensor. Expanding the square and normalising the row cancels the
//    exp(-||x54_i||^2/2s^2) factor, which is constant along j, so H is a plain masked
//    softmax over (x54_i . x54_j - ||x54_j||^2/2) once sigma is folded into W54 at
//    conversion time. No exp/clamp/divide in the graph, and H reuses ggml_soft_max_ext.
//
// 2. caching x54. H needs the 54D state of every past token, which the unified KV cache has
//    no slot for. Recovering it from the cached K as W54.Wk^-1.K is exact in real arithmetic
//    but cond(Wk) reaches 6.9e6 in the reference weights, so it does not survive an F16
//    cache. Instead n_embd_head_k is widened by d54 and x54 rides along in the key rows;
//    K is then sliced back apart on read. The cost is d54 floats per head per token.

void llama_model_cosmos::load_arch_hparams(llama_model_loader & ml) {
    ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);

    if (hparams.n_embd_head_k_full <= hparams.n_embd_head_v_full) {
        throw std::runtime_error("cosmos: attention.key_length must exceed value_length by d54");
    }

    type = LLM_TYPE_UNKNOWN;
}

void llama_model_cosmos::load_arch_tensors(llama_model_loader &) {
    LLAMA_LOAD_LOCALS;

    const int64_t d54 = hparams.n_embd_head_k() - hparams.n_embd_head_v();

    tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
    pos_embd = create_tensor(tn(LLM_TENSOR_POS_EMBD,   "weight"), {n_embd, n_ctx_train}, 0);

    output_norm   = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
    output_norm_b = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "bias"),   {n_embd}, 0);
    output        = create_tensor(tn(LLM_TENSOR_OUTPUT,      "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED);

    if (output == NULL) {
        output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
    }

    for (int i = 0; i < n_layer; ++i) {
        auto & layer = layers[i];

        layer.attn_norm   = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
        layer.attn_norm_b = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "bias",   i), {n_embd}, 0);

        layer.wqkv   = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", i), {n_embd, 3*n_embd}, 0);
        layer.wqkv_b = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "bias",   i), {3*n_embd}, 0);

        layer.wo   = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd, n_embd}, 0);
        layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias",   i), {n_embd}, 0);

        layer.attn_54   = create_tensor(tn(LLM_TENSOR_ATTN_54,   "weight", i), {n_embd, d54}, 0);
        layer.attn_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", i), {1}, 0);

        layer.ffn_norm   = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
        layer.ffn_norm_b = create_tensor(tn(LLM_TENSOR_FFN_NORM, "bias",   i), {n_embd}, 0);

        layer.ffn_up     = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), {n_embd, n_ff}, 0);
        layer.ffn_up_b   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "bias",   i), {n_ff}, 0);
        layer.ffn_down   = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
        layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias",   i), {n_embd}, 0);
    }
}

std::unique_ptr<llm_graph_context> llama_model_cosmos::build_arch_graph(const llm_graph_params & params) const {
    return std::make_unique<graph>(*this, params);
}

llama_model_cosmos::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
    const int64_t n_embd_head = hparams.n_embd_head_v();
    const int64_t d54         = hparams.n_embd_head_k() - n_embd_head;
    const float   kq_scale    = 1.0f/sqrtf(float(n_embd_head));

    ggml_tensor * cur;
    ggml_tensor * inpL;

    inpL = build_inp_embd(model.tok_embd);

    ggml_tensor * inp_pos = build_inp_pos();

    auto * inp_attn = build_attn_inp_kv();

    ggml_tensor * pos = ggml_get_rows(ctx0, model.pos_embd, inp_pos);
    cb(pos, "pos_embd", -1);

    inpL = ggml_add(ctx0, inpL, pos);
    cb(inpL, "inpL", -1);

    ggml_tensor * inp_out_ids = build_inp_out_ids();

    for (int il = 0; il < n_layer; ++il) {
        cur = build_norm(inpL,
                model.layers[il].attn_norm,
                model.layers[il].attn_norm_b,
                LLM_NORM, il);
        cb(cur, "attn_norm", il);

        {
            ggml_tensor * qkv = build_lora_mm(model.layers[il].wqkv, cur);
            qkv = ggml_add(ctx0, qkv, model.layers[il].wqkv_b);
            cb(qkv, "wqkv", il);

            ggml_tensor * Qcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd, n_tokens, qkv->nb[1], 0*sizeof(float)*n_embd));
            ggml_tensor * Kcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd, n_tokens, qkv->nb[1], 1*sizeof(float)*n_embd));
            ggml_tensor * Vcur = ggml_cont(ctx0, ggml_view_2d(ctx0, qkv, n_embd, n_tokens, qkv->nb[1], 2*sizeof(float)*n_embd));

            ggml_tensor * x54 = build_lora_mm(model.layers[il].attn_54, cur);
            cb(x54, "x54", il);

            Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head, n_head,    n_tokens);
            Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens);
            Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens);

            // every key head carries a copy of x54 so that slicing it back out on read does
            // not depend on which head the cache hands us first
            ggml_tensor * x54r = ggml_repeat_4d(ctx0, ggml_reshape_3d(ctx0, x54, d54, 1, n_tokens),
                    d54, n_head_kv, n_tokens, 1);
            ggml_tensor * Kext = ggml_concat(ctx0, Kcur, x54r, 0);
            cb(Kext, "Kext", il);

            ggml_build_forward_expand(gf, Qcur);
            ggml_build_forward_expand(gf, Vcur);
            ggml_build_forward_expand(gf, Kext);

            const auto * mctx_cur = inp_attn->mctx;

            ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, Kext, inp_attn->get_k_idxs(), il));
            ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, Vcur, inp_attn->get_v_idxs(), il));

            ggml_tensor * kq_mask = inp_attn->get_kq_mask();

            ggml_tensor * kall = mctx_cur->get_k(ctx0, il);
            ggml_tensor * v    = mctx_cur->get_v(ctx0, il);

            const auto n_stream = kall->ne[3];
            const bool v_trans  = v->nb[1] > v->nb[2];

            // slice the widened key rows back into the real key and the cached 54D state
            ggml_tensor * k = ggml_view_4d(ctx0, kall,
                    n_embd_head, kall->ne[1], kall->ne[2], kall->ne[3],
                    kall->nb[1], kall->nb[2], kall->nb[3], 0);

            ggml_tensor * s54 = ggml_view_4d(ctx0, kall,
                    d54, 1, kall->ne[2], kall->ne[3],
                    kall->nb[1], kall->nb[2], kall->nb[3],
                    n_embd_head*ggml_element_size(kall));

            s54 = ggml_cont_3d(ctx0, s54, d54, kall->ne[2], kall->ne[3]);
            cb(s54, "s54", il);

            ggml_tensor * q = ggml_view_4d(ctx0, Qcur,
                    Qcur->ne[0], Qcur->ne[1], Qcur->ne[2]/n_stream, n_stream,
                    Qcur->nb[1], Qcur->nb[2], Qcur->nb[3]/n_stream, 0);

            q = ggml_permute(ctx0, q, 0, 2, 1, 3);
            k = ggml_permute(ctx0, k, 0, 2, 1, 3);

            ggml_tensor * kq = ggml_mul_mat(ctx0, k, q);
            ggml_mul_mat_set_prec(kq, GGML_PREC_F32);
            cb(kq, "kq", il);

            ggml_tensor * a_std = ggml_soft_max_ext(ctx0, kq, kq_mask, kq_scale, 0.0f);
            cb(a_std, "a_std", il);

            // H as a masked softmax over (x54_i . x54_j - ||x54_j||^2/2)
            ggml_tensor * q54 = ggml_reshape_3d(ctx0, x54, d54, n_tokens/n_stream, n_stream);

            ggml_tensor * hdot = ggml_mul_mat(ctx0, s54, q54);
            ggml_mul_mat_set_prec(hdot, GGML_PREC_F32);

            ggml_tensor * nrm = ggml_sum_rows(ctx0, ggml_sqr(ctx0, s54));
            nrm = ggml_cont(ctx0, ggml_transpose(ctx0, nrm));

            ggml_tensor * hscore = ggml_add(ctx0, hdot, ggml_scale(ctx0, nrm, -0.5f));
            hscore = ggml_reshape_4d(ctx0, hscore, hscore->ne[0], hscore->ne[1], 1, hscore->ne[2]);

            ggml_tensor * H = ggml_soft_max_ext(ctx0, hscore, kq_mask, 1.0f, 0.0f);
            cb(H, "hebbian", il);

            // A_final = A_std + g*(H - A_std), so g == 0 is bit-for-bit standard attention
            H = ggml_repeat_4d(ctx0, H, a_std->ne[0], a_std->ne[1], a_std->ne[2], a_std->ne[3]);

            ggml_tensor * a = ggml_add(ctx0, a_std,
                    ggml_mul(ctx0, ggml_sub(ctx0, H, a_std), model.layers[il].attn_gate));
            cb(a, "a_final", il);

            ggml_tensor * vp = ggml_permute(ctx0, v, 0, 2, 1, 3);
            if (!v_trans) {
                vp = ggml_cont(ctx0, ggml_transpose(ctx0, vp));
            }

            ggml_tensor * kqv = ggml_mul_mat(ctx0, vp, a);
            cb(kqv, "kqv", il);

            cur = ggml_permute(ctx0, kqv, 0, 2, 1, 3);
            cur = ggml_cont_2d(ctx0, cur, n_embd_head*n_head, n_tokens);

            cur = build_lora_mm(model.layers[il].wo, cur);
            cur = ggml_add(ctx0, cur, model.layers[il].wo_b);
            cb(cur, "attn_out", il);
        }

        if (il == n_layer - 1 && inp_out_ids) {
            cur  = ggml_get_rows(ctx0,  cur, inp_out_ids);
            inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
        }

        ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL);
        cb(ffn_inp, "ffn_inp", il);

        {
            cur = build_norm(ffn_inp,
                    model.layers[il].ffn_norm,
                    model.layers[il].ffn_norm_b,
                    LLM_NORM, il);
            cb(cur, "ffn_norm", il);

            // built out rather than via build_ffn: LLM_FFN_GELU maps to ggml_gelu, the tanh
            // approximation, while the reference weights were trained under torch nn.GELU,
            // which is erf-exact. The approximation costs ~1e-3 per activation.
            cur = build_lora_mm(model.layers[il].ffn_up, cur);
            cur = ggml_add(ctx0, cur, model.layers[il].ffn_up_b);
            cur = ggml_gelu_erf(ctx0, cur);
            cur = build_lora_mm(model.layers[il].ffn_down, cur);
            cur = ggml_add(ctx0, cur, model.layers[il].ffn_down_b);
            cb(cur, "ffn_out", il);
        }

        cur = ggml_add(ctx0, cur, ffn_inp);

        cur = build_cvec(cur, il);
        cb(cur, "l_out", il);

        inpL = cur;
    }

    cur = build_norm(inpL,
            model.output_norm,
            model.output_norm_b,
            LLM_NORM, -1);
    cb(cur, "result_norm", -1);
    res->t_embd = cur;

    cur = build_lora_mm(model.output, cur);
    cb(cur, "result_output", -1);
    res->t_logits = cur;

    ggml_build_forward_expand(gf, cur);
}