// Author: Simon-Pierre Boucher — contact@spboucher.ai
#include "ops/ops.h"

#include "ops/cpu/cpu_ops.h"
#include "ops/metal/metal_ops.h"

#include <cmath>

namespace forge::ops {

namespace {

Backend g_backend = Backend::CPU;

bool gpu() { return g_backend == Backend::Metal; }

bool grad_needed(std::initializer_list<const Var*> inputs) {
    if (!Tape::get().enabled()) return false;
    for (const Var* v : inputs)
        if (v->requires_grad()) return true;
    return false;
}

// dst += src (backend-routed)
void accumulate(Tensor& dst, const Tensor& src) {
    if (gpu()) {
        metal::accumulate(dst, src);
        return;
    }
    float* pd = dst.data<float>();
    const float* ps = src.data<float>();
    for (int64_t i = 0; i < dst.numel(); ++i) pd[i] += ps[i];
}

// dst += src * s, s known on CPU
void axpy_const(Tensor& dst, const Tensor& src, float s) {
    if (gpu()) {
        Tensor tmp = Tensor::empty(src.shape());
        metal::scale(src, s, tmp);
        metal::accumulate(dst, tmp);
        return;
    }
    float* pd = dst.data<float>();
    const float* ps = src.data<float>();
    for (int64_t i = 0; i < dst.numel(); ++i) pd[i] += ps[i] * s;
}

// dst += src * s[0], s produced on-GPU (never read on CPU)
void axpy_tensor(Tensor& dst, const Tensor& src, const Tensor& s) {
    if (gpu()) {
        metal::axpy(dst, src, s);
        return;
    }
    const float sv = s.data<float>()[0];
    float* pd = dst.data<float>();
    const float* ps = src.data<float>();
    for (int64_t i = 0; i < dst.numel(); ++i) pd[i] += ps[i] * sv;
}

} // namespace

void set_backend(Backend b) { g_backend = b; }
Backend backend() { return g_backend; }

Var matmul(const Var& a, const Var& b, bool ta, bool tb) {
    const int64_t M = ta ? a.value().size(1) : a.value().size(0);
    const int64_t N = tb ? b.value().size(0) : b.value().size(1);
    Tensor out = Tensor::empty({M, N});
    if (gpu()) metal::matmul(a.value(), b.value(), out, ta, tb);
    else cpu::matmul(a.value(), b.value(), out, ta, tb);

    const bool needs = grad_needed({&a, &b});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([a, b, result, ta, tb]() {
            const Tensor& dc = result.grad();
            if (a.requires_grad()) {
                Tensor& da = a.grad();
                if (gpu()) {
                    if (!ta) metal::matmul(dc, b.value(), da, false, !tb, true);
                    else     metal::matmul(b.value(), dc, da, tb, true, true);
                } else {
                    if (!ta) cpu::matmul(dc, b.value(), da, false, !tb, true);
                    else     cpu::matmul(b.value(), dc, da, tb, true, true);
                }
            }
            if (b.requires_grad()) {
                Tensor& db = b.grad();
                if (gpu()) {
                    if (!tb) metal::matmul(a.value(), dc, db, !ta, false, true);
                    else     metal::matmul(dc, a.value(), db, true, ta, true);
                } else {
                    if (!tb) cpu::matmul(a.value(), dc, db, !ta, false, true);
                    else     cpu::matmul(dc, a.value(), db, true, ta, true);
                }
            }
        });
    }
    return result;
}

Var add(const Var& a, const Var& b) {
    Tensor out = Tensor::empty(a.value().shape());
    if (gpu()) metal::add(a.value(), b.value(), out);
    else cpu::add(a.value(), b.value(), out);

    const bool needs = grad_needed({&a, &b});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([a, b, result]() {
            if (a.requires_grad()) accumulate(a.grad(), result.grad());
            if (b.requires_grad()) accumulate(b.grad(), result.grad());
        });
    }
    return result;
}

Var add_bias(const Var& x, const Var& bias) {
    Tensor out = Tensor::empty(x.value().shape());
    if (gpu()) metal::add_bias(x.value(), bias.value(), out);
    else cpu::add_bias(x.value(), bias.value(), out);

    const bool needs = grad_needed({&x, &bias});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([x, bias, result]() {
            if (x.requires_grad()) accumulate(x.grad(), result.grad());
            if (bias.requires_grad()) {
                if (gpu()) metal::add_bias_backward(result.grad(), bias.grad());
                else cpu::add_bias_backward(result.grad(), bias.grad());
            }
        });
    }
    return result;
}

Var mul(const Var& a, const Var& b) {
    Tensor out = Tensor::empty(a.value().shape());
    if (gpu()) metal::mul(a.value(), b.value(), out);
    else cpu::mul(a.value(), b.value(), out);

    const bool needs = grad_needed({&a, &b});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([a, b, result]() {
            const Tensor& dout = result.grad();
            if (a.requires_grad()) {
                Tensor tmp = Tensor::empty(dout.shape());
                if (gpu()) metal::mul(dout, b.value(), tmp);
                else cpu::mul(dout, b.value(), tmp);
                accumulate(a.grad(), tmp);
            }
            if (b.requires_grad()) {
                Tensor tmp = Tensor::empty(dout.shape());
                if (gpu()) metal::mul(dout, a.value(), tmp);
                else cpu::mul(dout, a.value(), tmp);
                accumulate(b.grad(), tmp);
            }
        });
    }
    return result;
}

Var scale(const Var& a, float s) {
    Tensor out = Tensor::empty(a.value().shape());
    if (gpu()) metal::scale(a.value(), s, out);
    else cpu::scale(a.value(), s, out);

    const bool needs = grad_needed({&a});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([a, result, s]() {
            if (a.requires_grad()) axpy_const(a.grad(), result.grad(), s);
        });
    }
    return result;
}

Var silu(const Var& x) {
    Tensor out = Tensor::empty(x.value().shape());
    if (gpu()) metal::silu(x.value(), out);
    else cpu::silu(x.value(), out);

    const bool needs = grad_needed({&x});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([x, result]() {
            if (!x.requires_grad()) return;
            if (gpu()) metal::silu_backward(x.value(), result.grad(), x.grad());
            else cpu::silu_backward(x.value(), result.grad(), x.grad());
        });
    }
    return result;
}

Var gelu(const Var& x) {
    Tensor out = Tensor::empty(x.value().shape());
    if (gpu()) metal::gelu(x.value(), out);
    else cpu::gelu(x.value(), out);

    const bool needs = grad_needed({&x});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([x, result]() {
            if (!x.requires_grad()) return;
            if (gpu()) metal::gelu_backward(x.value(), result.grad(), x.grad());
            else cpu::gelu_backward(x.value(), result.grad(), x.grad());
        });
    }
    return result;
}

Var relu2(const Var& x) {
    Tensor out = Tensor::empty(x.value().shape());
    if (gpu()) metal::relu2(x.value(), out);
    else cpu::relu2(x.value(), out);

    const bool needs = grad_needed({&x});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([x, result]() {
            if (!x.requires_grad()) return;
            if (gpu()) metal::relu2_backward(x.value(), result.grad(), x.grad());
            else cpu::relu2_backward(x.value(), result.grad(), x.grad());
        });
    }
    return result;
}

Var softcap(const Var& x, float cap) {
    Tensor out = Tensor::empty(x.value().shape());
    if (gpu()) metal::softcap(x.value(), cap, out);
    else cpu::softcap(x.value(), cap, out);

    const bool needs = grad_needed({&x});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([x, cap, result]() {
            if (!x.requires_grad()) return;
            if (gpu()) metal::softcap_backward(result.value(), result.grad(), cap, x.grad());
            else cpu::softcap_backward(result.value(), result.grad(), cap, x.grad());
        });
    }
    return result;
}

Var softmax(const Var& x) {
    Tensor out = Tensor::empty(x.value().shape());
    if (gpu()) metal::softmax(x.value(), out);
    else cpu::softmax(x.value(), out);

    const bool needs = grad_needed({&x});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([x, result]() {
            if (!x.requires_grad()) return;
            if (gpu()) metal::softmax_backward(result.value(), result.grad(), x.grad());
            else cpu::softmax_backward(result.value(), result.grad(), x.grad());
        });
    }
    return result;
}

Var fake_quant(const Var& w, QuantMode mode) {
    if (mode == QuantMode::None) return w;
    const int m = mode == QuantMode::Int8 ? 0 : 1;
    Tensor out = Tensor::empty(w.value().shape());
    if (gpu()) metal::fake_quant(w.value(), m, out);
    else cpu::fake_quant(w.value(), m, out);

    const bool needs = grad_needed({&w});
    Var result(std::move(out), needs);
    if (needs) {
        // Straight-through estimator: d(quant(w))/dw ≈ I.
        Tape::get().record([w, result]() {
            if (w.requires_grad()) accumulate(w.grad(), result.grad());
        });
    }
    return result;
}

Var sigmoid(const Var& x) {
    Tensor out = Tensor::empty(x.value().shape());
    if (gpu()) metal::sigmoid(x.value(), out);
    else cpu::sigmoid(x.value(), out);

    const bool needs = grad_needed({&x});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([x, result]() {
            if (!x.requires_grad()) return;
            if (gpu()) metal::sigmoid_backward(result.value(), result.grad(), x.grad());
            else cpu::sigmoid_backward(result.value(), result.grad(), x.grad());
        });
    }
    return result;
}

Var topk_renorm(const Var& probs, const Tensor& bias, int64_t k, bool norm) {
    Tensor out = Tensor::empty(probs.value().shape());
    if (gpu()) metal::topk_renorm(probs.value(), bias, k, norm, out);
    else cpu::topk_renorm(probs.value(), bias, k, norm, out);

    const bool needs = grad_needed({&probs});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([probs, bias, k, norm, result]() {
            if (!probs.requires_grad()) return;
            if (gpu())
                metal::topk_renorm_backward(probs.value(), bias, result.grad(), k,
                                            norm, probs.grad());
            else
                cpu::topk_renorm_backward(probs.value(), bias, result.grad(), k,
                                          norm, probs.grad());
        });
    }
    return result;
}

Var row_scale(const Var& x, const Var& gates, int64_t e) {
    Tensor out = Tensor::empty(x.value().shape());
    if (gpu()) metal::row_scale(x.value(), gates.value(), e, out);
    else cpu::row_scale(x.value(), gates.value(), e, out);

    const bool needs = grad_needed({&x, &gates});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([x, gates, e, result]() {
            const Tensor& dout = result.grad();
            if (x.requires_grad()) {
                if (gpu()) metal::row_scale_accumulate(dout, gates.value(), e, x.grad());
                else cpu::row_scale_accumulate(dout, gates.value(), e, x.grad());
            }
            if (gates.requires_grad()) {
                if (gpu())
                    metal::row_scale_gate_backward(dout, x.value(), e, gates.grad());
                else
                    cpu::row_scale_gate_backward(dout, x.value(), e, gates.grad());
            }
        });
    }
    return result;
}

Var rmsnorm(const Var& x, const Var& w, float eps) {
    Tensor out = Tensor::empty(x.value().shape());
    if (gpu()) metal::rmsnorm(x.value(), w.value(), eps, out);
    else cpu::rmsnorm(x.value(), w.value(), eps, out);

    const bool needs = grad_needed({&x, &w});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([x, w, eps, result]() {
            Tensor dx_scratch, dw_scratch;
            Tensor& dx = x.requires_grad() ? x.grad()
                       : (dx_scratch = Tensor::zeros(x.value().shape()));
            Tensor& dw = w.requires_grad() ? w.grad()
                       : (dw_scratch = Tensor::zeros(w.value().shape()));
            if (gpu()) metal::rmsnorm_backward(x.value(), w.value(), eps, result.grad(), dx, dw);
            else cpu::rmsnorm_backward(x.value(), w.value(), eps, result.grad(), dx, dw);
        });
    }
    return result;
}

Var layernorm(const Var& x, const Var& w, const Var& b, float eps) {
    Tensor out = Tensor::empty(x.value().shape());
    if (gpu()) metal::layernorm(x.value(), w.value(), b.value(), eps, out);
    else cpu::layernorm(x.value(), w.value(), b.value(), eps, out);

    const bool needs = grad_needed({&x, &w, &b});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([x, w, b, eps, result]() {
            Tensor dx_scratch, dw_scratch, db_scratch;
            Tensor& dx = x.requires_grad() ? x.grad()
                       : (dx_scratch = Tensor::zeros(x.value().shape()));
            Tensor& dw = w.requires_grad() ? w.grad()
                       : (dw_scratch = Tensor::zeros(w.value().shape()));
            Tensor& db = b.requires_grad() ? b.grad()
                       : (db_scratch = Tensor::zeros(b.value().shape()));
            if (gpu())
                metal::layernorm_backward(x.value(), w.value(), eps, result.grad(), dx, dw, db);
            else
                cpu::layernorm_backward(x.value(), w.value(), eps, result.grad(), dx, dw, db);
        });
    }
    return result;
}

Var embedding(const Var& weight, const Tensor& ids) {
    const int64_t C = weight.value().size(1);
    std::vector<int64_t> out_shape = ids.shape();
    out_shape.push_back(C);
    Tensor out = Tensor::empty(std::move(out_shape));
    if (gpu()) metal::embedding(weight.value(), ids, out);
    else cpu::embedding(weight.value(), ids, out);

    const bool needs = grad_needed({&weight});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([weight, ids, result]() {
            if (!weight.requires_grad()) return;
            if (gpu()) metal::embedding_backward(ids, result.grad(), weight.grad());
            else cpu::embedding_backward(ids, result.grad(), weight.grad());
        });
    }
    return result;
}

Tensor rope_freqs(int64_t head_dim, float theta, float scale_factor,
                  float low_freq_factor, float high_freq_factor, int64_t original_ctx) {
    Tensor freqs = Tensor::empty({head_dim / 2});
    float* f = freqs.data<float>();
    constexpr float kTwoPi = 6.28318530717958647692f;
    for (int64_t k = 0; k < head_dim / 2; ++k) {
        float inv = std::pow(theta, -2.0f * float(k) / float(head_dim));
        if (scale_factor > 0.0f) {
            // HF "llama3" rope scaling (modeling_rope_utils.py).
            const float wavelen = kTwoPi / inv;
            const float lo = float(original_ctx) / low_freq_factor;   // long waves
            const float hi = float(original_ctx) / high_freq_factor;  // short waves
            if (wavelen > lo) {
                inv /= scale_factor;
            } else if (wavelen > hi) {
                const float s = (float(original_ctx) / wavelen - low_freq_factor) /
                                (high_freq_factor - low_freq_factor);
                inv = (1.0f - s) * inv / scale_factor + s * inv;
            }
        }
        f[k] = inv;
    }
    return freqs;
}

Var rope(const Var& x, int64_t n_heads, const Tensor& freqs, int64_t pos_offset) {
    Tensor out = Tensor::empty(x.value().shape());
    if (gpu()) metal::rope(x.value(), n_heads, freqs, pos_offset, out);
    else cpu::rope(x.value(), n_heads, freqs, pos_offset, out);

    const bool needs = grad_needed({&x});
    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record([x, n_heads, freqs, pos_offset, result]() {
            if (!x.requires_grad()) return;
            if (gpu()) metal::rope_backward(result.grad(), n_heads, freqs, pos_offset, x.grad());
            else cpu::rope_backward(result.grad(), n_heads, freqs, pos_offset, x.grad());
        });
    }
    return result;
}

Var rope(const Var& x, int64_t n_heads, float theta, int64_t pos_offset) {
    const int64_t hd = x.value().size(2) / n_heads;
    return rope(x, n_heads, rope_freqs(hd, theta), pos_offset);
}

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, float attn_softcap) {
    const int64_t B = q.value().size(0), T = q.value().size(1);
    const int64_t head_dim = q.value().size(2) / n_heads;
    Tensor out = Tensor::empty(q.value().shape());
    const bool needs = grad_needed({&q, &k, &v});

    // Fused path: keeps only the per-row logsumexp instead of a [B,H,T,T]
    // probability tensor, so memory is linear in T rather than quadratic.
    // Softcap needs the materialized-probs path (the fused kernels' online
    // softmax has no cap hook yet).
    if (gpu() && metal::flash_supported(head_dim) && attn_softcap <= 0.0f) {
        Tensor lse = Tensor::empty({B, n_heads, T});
        metal::flash_attention(q.value(), k.value(), v.value(), n_heads, n_kv_heads,
                               causal, scale, out, lse, metal::FlashKernel::Auto,
                               window);
        Var result(std::move(out), needs);
        if (needs) {
            Tape::get().record(
                [q, k, v, lse, n_heads, n_kv_heads, causal, scale, window, result]() {
                    Tensor dq_scratch, dk_scratch, dv_scratch;
                    Tensor& dq = q.requires_grad() ? q.grad()
                               : (dq_scratch = Tensor::zeros(q.value().shape()));
                    Tensor& dk = k.requires_grad() ? k.grad()
                               : (dk_scratch = Tensor::zeros(k.value().shape()));
                    Tensor& dv = v.requires_grad() ? v.grad()
                               : (dv_scratch = Tensor::zeros(v.value().shape()));
                    metal::flash_attention_backward(q.value(), k.value(), v.value(),
                                                    result.value(), lse, result.grad(),
                                                    n_heads, n_kv_heads, causal, scale,
                                                    dq, dk, dv,
                                                    metal::FlashKernel::Auto, window);
                });
        }
        return result;
    }

    Tensor probs = Tensor::empty({B, n_heads, T, T});
    if (gpu())
        metal::attention(q.value(), k.value(), v.value(), n_heads, n_kv_heads, causal,
                         scale, out, &probs, window, attn_softcap);
    else
        cpu::attention(q.value(), k.value(), v.value(), n_heads, n_kv_heads, causal,
                       scale, out, &probs, window, attn_softcap);

    Var result(std::move(out), needs);
    if (needs) {
        Tape::get().record(
            [q, k, v, probs, n_heads, n_kv_heads, scale, attn_softcap, result]() {
                Tensor dq_scratch, dk_scratch, dv_scratch;
                Tensor& dq = q.requires_grad() ? q.grad()
                           : (dq_scratch = Tensor::zeros(q.value().shape()));
                Tensor& dk = k.requires_grad() ? k.grad()
                           : (dk_scratch = Tensor::zeros(k.value().shape()));
                Tensor& dv = v.requires_grad() ? v.grad()
                           : (dv_scratch = Tensor::zeros(v.value().shape()));
                if (gpu())
                    metal::attention_backward(q.value(), k.value(), v.value(), probs,
                                              result.value(), result.grad(), n_heads,
                                              n_kv_heads, scale, dq, dk, dv,
                                              attn_softcap);
                else
                    cpu::attention_backward(q.value(), k.value(), v.value(), probs,
                                            result.value(), result.grad(), n_heads,
                                            n_kv_heads, scale, dq, dk, dv,
                                            attn_softcap);
            });
    }
    return result;
}

Var cross_entropy(const Var& logits, const Tensor& targets) {
    const bool needs = grad_needed({&logits});
    const int64_t N = logits.value().size(0);

    // n_valid comes from the CPU-resident targets (written by the data
    // loader, never touched by the GPU).
    int64_t n_valid = 0;
    for (int64_t i = 0; i < N; ++i) {
        const int32_t t = targets.dtype() == DType::I32
                              ? targets.data<int32_t>()[i]
                              : int32_t(targets.data<uint16_t>()[i]);
        if (t >= 0) ++n_valid;
    }

    // llm.c pattern: the logit gradient is a byproduct of the forward pass;
    // save it and scale by d(loss) at backward time.
    Tensor dlogits;
    if (needs) dlogits = Tensor::zeros(logits.value().shape());

    Tensor loss_out;
    if (gpu()) {
        Tensor losses = Tensor::empty({N});
        loss_out = Tensor::empty({1});
        metal::cross_entropy(logits.value(), targets, n_valid, losses, loss_out,
                             needs ? &dlogits : nullptr);
    } else {
        const float loss =
            cpu::cross_entropy(logits.value(), targets, needs ? &dlogits : nullptr);
        loss_out = Tensor::full({1}, loss);
    }

    Var result(std::move(loss_out), needs);
    if (needs) {
        Tape::get().record([logits, dlogits, result]() {
            if (logits.requires_grad())
                axpy_tensor(logits.grad(), dlogits, result.grad());
        });
    }
    return result;
}

Var reshape(const Var& x, std::vector<int64_t> shape) {
    return x.reshaped(std::move(shape));
}

} // namespace forge::ops