// Author: Simon-Pierre Boucher — contact@spboucher.ai
#pragma once

#include "core/tensor.h"

#include <cstdint>
#include <random>

// CPU reference implementations (forward + backward). These are the ground
// truth every Metal kernel is validated against — clarity beats speed here.
// All tensors f32 and contiguous unless stated otherwise. Backward
// functions ACCUMULATE into their d* outputs (autograd semantics).
namespace forge::cpu {

// ---- init -----------------------------------------------------------------
void fill_normal(Tensor& t, float mean, float stddev, std::mt19937_64& rng);
void fill_uniform_int(Tensor& t, int64_t low, int64_t high, std::mt19937_64& rng);

// ---- matmul ---------------------------------------------------------------
// C[M,N] = A ⋅ B with optional transposes (A stored [M,K] or [K,M], B stored
// [K,N] or [N,K]). accumulate=true adds into C instead of overwriting
// (backward passes accumulate into gradients).
void matmul(const Tensor& a, const Tensor& b, Tensor& c,
            bool transpose_a = false, bool transpose_b = false,
            bool accumulate = false);

// ---- elementwise ----------------------------------------------------------
void add(const Tensor& a, const Tensor& b, Tensor& out);
void mul(const Tensor& a, const Tensor& b, Tensor& out);
void scale(const Tensor& a, float s, Tensor& out);
// x: [N, C], bias: [C], out[i,j] = x[i,j] + bias[j]
void add_bias(const Tensor& x, const Tensor& bias, Tensor& out);
void add_bias_backward(const Tensor& dout, Tensor& dbias);

void silu(const Tensor& x, Tensor& out);
void silu_backward(const Tensor& x, const Tensor& dout, Tensor& dx);
void gelu(const Tensor& x, Tensor& out); // tanh approximation (GPT-2)
void gelu_backward(const Tensor& x, const Tensor& dout, Tensor& dx);
void relu2(const Tensor& x, Tensor& out); // max(x,0)^2
void relu2_backward(const Tensor& x, const Tensor& dout, Tensor& dx);
void softcap(const Tensor& x, float cap, Tensor& out); // cap*tanh(x/cap)
// dx += dout * (1 - (y/cap)^2) — takes the forward OUTPUT y
void softcap_backward(const Tensor& y, const Tensor& dout, float cap, Tensor& dx);

// ---- norms (row-wise over the last dim) -----------------------------------
// x: [N, C], w: [C] (b: [C] for layernorm), out: [N, C]
void rmsnorm(const Tensor& x, const Tensor& w, float eps, Tensor& out);
void rmsnorm_backward(const Tensor& x, const Tensor& w, float eps,
                      const Tensor& dout, Tensor& dx, Tensor& dw);
void layernorm(const Tensor& x, const Tensor& w, const Tensor& b, float eps, Tensor& out);
void layernorm_backward(const Tensor& x, const Tensor& w, float eps,
                        const Tensor& dout, Tensor& dx, Tensor& dw, Tensor& db);

// ---- softmax (row-wise over the last dim, max-subtracted) ------------------
void softmax(const Tensor& x, Tensor& out);
// dx = P ∘ (dout − rowsum(dout ∘ P))
void softmax_backward(const Tensor& p, const Tensor& dout, Tensor& dx);

// ---- embedding -------------------------------------------------------------
// weight: [V, C] f32; ids: [N] i32; out: [N, C]
void embedding(const Tensor& weight, const Tensor& ids, Tensor& out);
void embedding_backward(const Tensor& ids, const Tensor& dout, Tensor& dweight);

// ---- RoPE (interleaved-pairs / GPT-J convention; see RESEARCH.md §7) -------
// x: [B, T, H*hd]; rotates pairs (2k, 2k+1) inside each head. pos_offset
// shifts absolute positions (KV-cache generation).
// freqs: [head_dim/2] per-pair inverse frequencies, precomputed by the caller.
void rope(const Tensor& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset,
          Tensor& out);
// Backward of a rotation is the inverse rotation applied to dout.
void rope_backward(const Tensor& dout, int64_t n_heads, const Tensor& freqs,
                   int64_t pos_offset, Tensor& dx);

// ---- attention (composed reference: scores → mask → softmax → PV) ----------
// q: [B, T, H*hd], k/v: [B, T, Hkv*hd], out: [B, T, H*hd]. GQA via
// kv_head = h / (H / Hkv). If probs_out is non-null it receives the softmax
// probabilities [B, H, T, T] (needed for backward).
// window > 0: sliding-window attention (attend to the last `window` keys,
// self included). attn_softcap > 0: cap·tanh(score/cap) pre-softmax (Gemma 2).
void attention(const Tensor& q, const Tensor& k, const Tensor& v,
               int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,
               Tensor& out, Tensor* probs_out, int64_t window = 0,
               float attn_softcap = 0.0f);
// Uses the FlashAttention-2 identity rowsum(dP ∘ P) == dO_i · O_i, so `out`
// (the forward result) is required. Masked positions have prob 0, so the
// window needs no explicit handling here.
void attention_backward(const Tensor& q, const Tensor& k, const Tensor& v,
                        const Tensor& probs, const Tensor& out, const Tensor& dout,
                        int64_t n_heads, int64_t n_kv_heads, float scale,
                        Tensor& dq, Tensor& dk, Tensor& dv, float attn_softcap = 0.0f);

// ---- QAT / MoE --------------------------------------------------------------
// Per-row fake quantization of a 2-D weight; mode 0 = int8 (absmax/127),
// mode 1 = ternary (absmean, BitNet-style). Backward is STE — identity.
void fake_quant(const Tensor& w, int mode, Tensor& out);
void sigmoid(const Tensor& x, Tensor& out);
// dx += dout * y * (1-y) — takes the forward OUTPUT y
void sigmoid_backward(const Tensor& y, const Tensor& dout, Tensor& dx);
// Keep top-k per row of p [N,E] — selected by p+bias, gated by p alone
// (bias [E], all-zero for classic routing); norm renormalizes kept gates to
// sum 1. Ties broken by lower index, same as the Metal kernel.
void topk_renorm(const Tensor& p, const Tensor& bias, int64_t k, bool norm,
                 Tensor& out);
void topk_renorm_backward(const Tensor& p, const Tensor& bias, const Tensor& dout,
                          int64_t k, bool norm, Tensor& dp);
// counts[e] += #rows with nonzero gate for expert e (noaux load statistic)
void expert_counts(const Tensor& gates, Tensor& counts);
// out[i,:] = x[i,:] * gates[i,e]; _accumulate does dst += (dx path);
// gate_backward does dgates[i,e] += dot(dout[i,:], x[i,:]).
void row_scale(const Tensor& x, const Tensor& gates, int64_t e, Tensor& out);
void row_scale_accumulate(const Tensor& x, const Tensor& gates, int64_t e, Tensor& dst);
void row_scale_gate_backward(const Tensor& dout, const Tensor& x, int64_t e,
                             Tensor& dgates);

// ---- fused softmax cross-entropy -------------------------------------------
// logits: [N, V]; targets: [N] i32 (ignore_index = -1). Returns mean loss
// over valid rows. dlogits (if non-null) receives (softmax − onehot)/n_valid
// — the gradient for dloss = 1, ACCUMULATED.
float cross_entropy(const Tensor& logits, const Tensor& targets, Tensor* dlogits);

} // namespace forge::cpu