// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Numerical gradient checking (central differences) for every op with
// parameters, plus a full tiny-model loss. Forward math is f32, so the
// difference quotient carries ~1e-4-relative noise; eps is chosen per the
// usual sqrt(machine-eps) rule for O(1) inputs and tolerances follow
// CLAUDE.md (rel error <= 1e-3 with an absolute floor).
#include "core/autograd.h"
#include "nn/transformer.h"
#include "ops/ops.h"

#include <cmath>
#include <cstdio>
#include <functional>
#include <random>
#include <vector>

using namespace forge;

namespace {

int g_failures = 0;
constexpr float kEps = 1e-2f;
constexpr float kRelTol = 1e-3f;
constexpr float kAbsTol = 1e-4f;

// loss_fn must run forward and return the scalar loss Var. Gradcheck
// perturbs every element of every checked Var and compares the analytic
// gradient against (f(x+e) - f(x-e)) / 2e.
void gradcheck(const char* name, const std::function<Var()>& loss_fn,
               const std::vector<Var>& checked, float eps = kEps, float tol = kRelTol) {
    // analytic
    for (const Var& p : checked) p.zero_grad();
    Tape::get().clear();
    Var loss = loss_fn();
    Tape::get().backward(loss);

    float worst = 0.0f;
    for (const Var& p : checked) {
        const int64_t n = p.value().numel();
        float* x = p.value().data<float>();
        const float* g = p.grad().data<float>();
        for (int64_t i = 0; i < n; ++i) {
            const float saved = x[i];
            float fp, fm;
            {
                NoGrad ng;
                x[i] = saved + eps;
                fp = loss_fn().value().data<float>()[0];
                x[i] = saved - eps;
                fm = loss_fn().value().data<float>()[0];
            }
            x[i] = saved;
            const float num = (fp - fm) / (2.0f * eps);
            const float ana = g[i];
            const float err = std::fabs(num - ana);
            const float rel = err / std::max(kAbsTol / kRelTol, std::fabs(num) + std::fabs(ana));
            worst = std::max(worst, rel);
        }
    }
    if (worst <= tol) {
        std::printf("  ok: %s (max rel err %.2e)\n", name, double(worst));
    } else {
        std::printf("  FAIL: %s (max rel err %.2e > %.0e)\n", name, double(worst),
                    double(tol));
        ++g_failures;
    }
}

Var make_var(std::vector<int64_t> shape, std::mt19937_64& rng, float std = 1.0f) {
    Tensor t = Tensor::empty(std::move(shape));
    cpu::fill_normal(t, 0.0f, std, rng);
    return Var(std::move(t), /*requires_grad=*/true);
}

// Reduce an op output to a scalar via a fixed random projection so every
// output element influences the loss.
Var project(const Var& y, const Tensor& r) {
    Var rv(r, /*requires_grad=*/false);
    const int64_t n = y.value().numel();
    Var y2 = y.reshaped({1, n});
    Var r2 = rv.reshaped({1, n});
    return ops::matmul(y2, r2, false, true); // [1,1]
}

Tensor random_proj(int64_t n, std::mt19937_64& rng) {
    Tensor r = Tensor::empty({n});
    cpu::fill_normal(r, 0.0f, 1.0f, rng);
    return r;
}

} // namespace

int main() {
    std::mt19937_64 rng(42);

    // matmul, all transpose variants
    for (int ta = 0; ta <= 1; ++ta) {
        for (int tb = 0; tb <= 1; ++tb) {
            const int64_t M = 3, K = 4, N = 5;
            Var a = ta ? make_var({K, M}, rng) : make_var({M, K}, rng);
            Var b = tb ? make_var({N, K}, rng) : make_var({K, N}, rng);
            Tensor r = random_proj(M * N, rng);
            char name[64];
            std::snprintf(name, sizeof(name), "matmul ta=%d tb=%d", ta, tb);
            gradcheck(name, [&]() { return project(ops::matmul(a, b, ta, tb), r); },
                      {a, b});
        }
    }

    // add_bias
    {
        Var x = make_var({4, 6}, rng);
        Var b = make_var({6}, rng);
        Tensor r = random_proj(24, rng);
        gradcheck("add_bias", [&]() { return project(ops::add_bias(x, b), r); }, {x, b});
    }

    // mul
    {
        Var a = make_var({3, 5}, rng);
        Var b = make_var({3, 5}, rng);
        Tensor r = random_proj(15, rng);
        gradcheck("mul", [&]() { return project(ops::mul(a, b), r); }, {a, b});
    }

    // silu / gelu
    {
        Var x = make_var({4, 5}, rng);
        Tensor r = random_proj(20, rng);
        gradcheck("silu", [&]() { return project(ops::silu(x), r); }, {x});
        gradcheck("gelu", [&]() { return project(ops::gelu(x), r); }, {x});
    }

    // rmsnorm / layernorm
    {
        Var x = make_var({3, 8}, rng);
        Var w = make_var({8}, rng);
        Var b = make_var({8}, rng);
        Tensor r = random_proj(24, rng);
        gradcheck("rmsnorm", [&]() { return project(ops::rmsnorm(x, w, 1e-6f), r); }, {x, w});
        gradcheck("layernorm",
                  [&]() { return project(ops::layernorm(x, w, b, 1e-6f), r); }, {x, w, b});
    }

    // embedding
    {
        Var w = make_var({7, 4}, rng);
        Tensor ids = Tensor::empty({2, 3}, DType::I32);
        std::mt19937_64 idrng(7);
        cpu::fill_uniform_int(ids, 0, 7, idrng);
        Tensor r = random_proj(2 * 3 * 4, rng);
        gradcheck("embedding", [&]() { return project(ops::embedding(w, ids), r); }, {w});
    }

    // rope
    {
        Var x = make_var({2, 5, 2 * 6}, rng); // B=2, T=5, H=2, hd=6
        Tensor r = random_proj(2 * 5 * 12, rng);
        gradcheck("rope", [&]() { return project(ops::rope(x, 2, 10000.0f, 0), r); }, {x});
    }

    // attention: MHA causal, then GQA
    {
        const int64_t B = 2, T = 4, H = 2, hd = 4;
        Var q = make_var({B, T, H * hd}, rng, 0.5f);
        Var k = make_var({B, T, H * hd}, rng, 0.5f);
        Var v = make_var({B, T, H * hd}, rng, 0.5f);
        Tensor r = random_proj(B * T * H * hd, rng);
        const float scale = 1.0f / std::sqrt(float(hd));
        gradcheck("attention causal MHA",
                  [&]() { return project(ops::attention(q, k, v, H, H, true, scale), r); },
                  {q, k, v});

        Var kg = make_var({B, T, 1 * hd}, rng, 0.5f);
        Var vg = make_var({B, T, 1 * hd}, rng, 0.5f);
        gradcheck("attention causal GQA (2q/1kv)",
                  [&]() { return project(ops::attention(q, kg, vg, H, 1, true, scale), r); },
                  {q, kg, vg});
    }

    // cross entropy (with an ignored index)
    {
        Var logits = make_var({6, 9}, rng);
        Tensor targets = Tensor::empty({6}, DType::I32);
        std::mt19937_64 idrng(11);
        cpu::fill_uniform_int(targets, 0, 9, idrng);
        targets.data<int32_t>()[3] = -1; // ignore_index
        gradcheck("cross_entropy", [&]() { return ops::cross_entropy(logits, targets); },
                  {logits});
    }

    // full tiny model: every parameter of a 2-layer transformer
    {
        ModelConfig cfg;
        cfg.n_layers = 2;
        cfg.d_model = 16;
        cfg.n_heads = 2;
        cfg.n_kv_heads = 1; // exercise GQA
        cfg.d_ff = 24;
        cfg.vocab_size = 11;
        cfg.context_length = 8;
        cfg.tied_embeddings = true;
        nn::Transformer model(cfg, 123);

        Tensor ids = Tensor::empty({2, 5}, DType::I32);
        Tensor targets = Tensor::empty({2 * 5}, DType::I32);
        std::mt19937_64 idrng(13);
        cpu::fill_uniform_int(ids, 0, cfg.vocab_size, idrng);
        cpu::fill_uniform_int(targets, 0, cfg.vocab_size, idrng);

        std::vector<Var> params;
        std::unordered_map<const void*, bool> seen;
        for (const auto& [name, p] : model.named_parameters()) {
            if (!seen.emplace(p.id(), true).second) continue;
            params.push_back(p);
        }
        // Composite check: the analytic gradient is exact per the per-op
        // checks above; the f32 forward pass through ~100 ops limits the
        // central-difference quotient itself (error scales as eps^2 —
        // verified 2.6e-2 @ eps=1e-2 -> 2.5e-3 @ eps=3e-3), so this runs
        // with a documented noise-bound tolerance, not the per-op 1e-3.
        gradcheck("tiny transformer (all params)",
                  [&]() {
                      Var logits = model.forward(ids);
                      return ops::cross_entropy(logits, targets);
                  },
                  params, /*eps=*/3e-3f, /*tol=*/5e-3f);
    }

    if (g_failures) {
        std::printf("\n%d gradcheck(s) FAILED\n", g_failures);
        return 1;
    }
    std::printf("\nall gradchecks passed\n");
    return 0;
}