/*
 * kcpp / ICPP - the Internet C++ compiler and interpreter
 *
 * Copyright (c) 2026 Muhammad Anisur Rahman
 * Free to use and distribute under the MIT License.
 *
 * SPDX-License-Identifier: MIT
 * See the LICENSE file at the repository root for full terms.
 */
/* -llmtrace educational tracing (run by the harness with KCPP_LLMTRACE=1,
   equivalent to the kcpp -llmtrace flag): exercises the whole traced
   pipeline end to end at tiny sizes - tokenize, embedding lookup,
   positional encoding, add, linear, attention (causal), loss, backward,
   SGD update - and the harness greps for the SGD step's banner, which only
   prints when tracing is live all the way through. Also proves the traced
   builtins still COMPUTE correctly while tracing (loss for near-uniform
   tiny weights over a 256-token vocab must be ~ln(256)=5.5). */
int main(void)
{
    int ids; int emb; int x; int pe;
    int wq; int bq; int wk; int bk; int wv; int bv;
    int q; int k; int v; int attn;
    int wlm; int blm; int logits; int targets; int dlogits; int zgrad;
    double loss;
    int i; int j;

    ids = tensor_tokenize("hi");
    emb = tensor_new(256, 4);
    for (i = 0; i < 256; i = i + 1)
    {
        for (j = 0; j < 4; j = j + 1)
        {
            tensor_set(emb, i, j, 0.01 * (i % 10) + 0.1 * j);
        }
    }
    x = tensor_embedding_lookup(emb, ids);
    pe = tensor_positional_encoding(2, 4);
    x = tensor_add(x, pe);
    wq = tensor_new(4, 4); bq = tensor_new(1, 4);
    wk = tensor_new(4, 4); bk = tensor_new(1, 4);
    wv = tensor_new(4, 4); bv = tensor_new(1, 4);
    for (i = 0; i < 4; i = i + 1)
    {
        for (j = 0; j < 4; j = j + 1)
        {
            tensor_set(wq, i, j, 0.1 * (i + 1));
            tensor_set(wk, i, j, 0.1 * (j + 1));
            tensor_set(wv, i, j, 0.05 * (i + j));
        }
    }
    q = tensor_linear(x, wq, bq);
    k = tensor_linear(x, wk, bk);
    v = tensor_linear(x, wv, bv);
    attn = tensor_multihead_attention(q, k, v, 2, 1);
    wlm = tensor_new(4, 256);
    blm = tensor_new(1, 256);
    logits = tensor_linear(attn, wlm, blm);
    targets = tensor_new(1, 2);
    tensor_set(targets, 0, 0, 105.0);
    tensor_set(targets, 0, 1, 33.0);
    loss = tensor_cross_entropy_loss(logits, targets);
    iprint("llmtrace loss=", loss);
    dlogits = tensor_cross_entropy_loss_backward(logits, targets);
    zgrad = tensor_new(4, 256);
    tensor_sgd_step(wlm, zgrad, 0.01);
    return 0;
}
