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

#include "core/autograd.h"
#include "core/tensor.h"

#include <cstdint>

// Autograd-aware ops. Forward runs on the active backend (CPU reference or
// Metal stream) and, when the tape is enabled and an input requires grad,
// records a backward lambda that accumulates into input .grad tensors on
// the same backend.
//
// Metal-mode contract: values live in unified memory but are produced
// asynchronously — CPU code must not READ op outputs (incl. the loss
// scalar) before metal::sync(), and must not WRITE tensors that pending
// dispatches read. The trainer's step structure guarantees both.
namespace forge::ops {

enum class Backend { CPU, Metal };
void set_backend(Backend b);
Backend backend();

// c = a ⋅ b (2-D, optional transposes on the stored operands)
Var matmul(const Var& a, const Var& b, bool transpose_a = false, bool transpose_b = false);

Var add(const Var& a, const Var& b);              // same shape
Var add_bias(const Var& x, const Var& bias);      // x [N,C] + bias [C]
Var mul(const Var& a, const Var& b);              // elementwise
Var scale(const Var& a, float s);

Var silu(const Var& x);
Var gelu(const Var& x);
Var relu2(const Var& x); // max(x,0)^2 — nanoGPT-speedrun activation
Var sigmoid(const Var& x); // σ(x) — sigmoid MoE routing (DeepSeek V3)

// Gemma-style soft capping: cap * tanh(x / cap). Bounds logits smoothly.
Var softcap(const Var& x, float cap);

// Row softmax over the last dim (autograd; the attention op has its own
// fused softmax — this one is for the MoE router's small rows).
Var softmax(const Var& x);

// QAT: per-row fake quantization of a 2-D weight. Forward quantizes
// (int8 absmax / ternary absmean); backward is the straight-through
// estimator — gradients flow unchanged into the f32 master weight.
enum class QuantMode { None, Int8, Ternary };
Var fake_quant(const Var& w, QuantMode mode);

// MoE gating: keep top-k per row of probs, zero the rest. Selection ranks
// probs + bias (bias [E], requires_grad=false — the DeepSeek-V3 noaux
// balance bias; pass zeros for classic routing) but gate values are the
// biasless probs; norm renormalizes kept gates to sum 1. Gradients flow
// through kept entries only.
Var topk_renorm(const Var& probs, const Tensor& bias, int64_t k, bool norm = true);
// y[i,:] = x[i,:] * gates[i,e] — weight expert e's output by its gate column.
Var row_scale(const Var& x, const Var& gates, int64_t e);

Var rmsnorm(const Var& x, const Var& w, float eps);
Var layernorm(const Var& x, const Var& w, const Var& b, float eps);

// weight [V,C], ids [B,T] (u16/i32) → [B,T,C]
Var embedding(const Var& weight, const Tensor& ids);

// Build the [head_dim/2] inverse-frequency table: theta^(-2k/hd), optionally
// rescaled HF-"llama3" style (scale_factor > 0): high-frequency components
// untouched, low-frequency divided by factor, smooth blend between.
Tensor rope_freqs(int64_t head_dim, float theta, float scale_factor = 0.0f,
                  float low_freq_factor = 1.0f, float high_freq_factor = 4.0f,
                  int64_t original_ctx = 8192);

// x [B,T,H*hd], interleaved-pairs RoPE per head. freqs from rope_freqs()
// (callers cache it — one table per layer flavour).
Var rope(const Var& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset = 0);
// Convenience: plain-theta table built per call (tests / one-offs).
Var rope(const Var& x, int64_t n_heads, float theta, int64_t pos_offset = 0);

// q [B,T,H*hd], k/v [B,T,Hkv*hd] → [B,T,H*hd]; causal + GQA.
// window > 0: sliding-window attention (Mistral/Gemma3) — fused scalar path.
// attn_softcap > 0: cap·tanh(score/cap) pre-softmax (Gemma2) — routes to the
// unfused path (probs materialized), so reserve it for short contexts.
Var attention(const Var& q, const Var& k, const Var& v,
              int64_t n_heads, int64_t n_kv_heads, bool causal, float scale,
              int64_t window = 0, float attn_softcap = 0.0f);

// logits [N,V], targets [N] (i32/u16, ignore_index=-1) → scalar mean loss
Var cross_entropy(const Var& logits, const Tensor& targets);

// Shape helper: shares storage, no tape node needed (grad shapes follow value shapes).
Var reshape(const Var& x, std::vector<int64_t> shape);

} // namespace forge::ops