// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// End-to-end sanity (CLAUDE.md testing protocol #3): overfit a single batch
// of 64 sequences to loss < 0.05 within 500 steps. Tiny model so the CPU
// reference path finishes in reasonable time; the same test re-runs on the
// GPU path at M4.
#include "nn/transformer.h"
#include "ops/ops.h"
#include "train/optimizer.h"

#include <cstdio>
#include <random>

using namespace forge;

int main() {
    ModelConfig cfg;
    cfg.n_layers = 2;
    cfg.d_model = 64;
    cfg.n_heads = 2;
    cfg.n_kv_heads = 2;
    cfg.d_ff = 172; // ~8/3 * d
    cfg.vocab_size = 128;
    cfg.context_length = 32;
    cfg.tied_embeddings = true;

    const int64_t B = 64, T = 32;

    nn::Transformer model(cfg, 1337);

    std::mt19937_64 rng(999);
    Tensor ids = Tensor::empty({B, T}, DType::I32);
    cpu::fill_uniform_int(ids, 0, cfg.vocab_size, rng);
    // next-token targets within the fixed batch
    Tensor targets = Tensor::empty({B * T}, DType::I32);
    for (int64_t b = 0; b < B; ++b) {
        for (int64_t t = 0; t < T; ++t) {
            targets.data<int32_t>()[b * T + t] =
                (t + 1 < T) ? ids.data<int32_t>()[b * T + t + 1] : -1;
        }
    }

    train::AdamW::Options opts;
    opts.weight_decay = 0.0f; // pure memorization task
    train::AdamW opt(model.named_parameters(), opts);

    const float lr = 3e-3f;
    float loss_val = -1.0f;
    for (int step = 0; step < 500; ++step) {
        opt.zero_grad();
        Var loss = model.loss(ids, targets);
        Tape::get().backward(loss);
        opt.clip_global_norm(1.0f);
        opt.step(lr);
        loss_val = loss.value().data<float>()[0];
        if (step % 50 == 0) std::printf("step %3d loss %.4f\n", step, double(loss_val));
        if (loss_val < 0.05f) {
            std::printf("step %3d loss %.4f — target reached\n", step, double(loss_val));
            break;
        }
    }

    if (loss_val < 0.05f) {
        std::printf("\noverfit ok (final loss %.4f)\n", double(loss_val));
        return 0;
    }
    std::printf("\noverfit FAILED (final loss %.4f >= 0.05)\n", double(loss_val));
    return 1;
}