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

#include "nn/config.h"
#include "nn/linear.h"
#include "ops/cpu/cpu_ops.h"
#include "ops/metal/metal_ops.h"
#include "ops/ops.h"

#include <memory>
#include <string>
#include <vector>

namespace forge::nn {

// Interface so the MoE variant slots into TransformerBlock without branching
// at every call site. aux() is the block's auxiliary loss ([1] Var, undefined
// when the variant has none).
class MLPBase : public Module {
public:
    virtual ~MLPBase() = default;
    // x: [N, C] 2-D
    virtual Var forward(const Var& x) const = 0;
    virtual Var aux() const { return Var(); }
    // noaux load-balance bias step (V3); no-op for dense/softmax variants.
    virtual void bias_update(float /*gamma*/) {}
};

// SwiGLU: w2( silu(x w1) ⊙ (x w3) )  |  GELU/ReLU²: w2( act(x w1) )
class MLP : public MLPBase {
public:
    MLP(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng)
        : swiglu_(cfg.activation == "swiglu"), relu2_(cfg.activation == "relu2") {
        const float std = 0.02f;
        const int64_t C = cfg.d_model, F = cfg.d_ff;
        w1_ = std::make_unique<Linear>(C, F, false, std, rng);
        if (swiglu_) {
            w3_ = std::make_unique<Linear>(C, F, false, std, rng);
        }
        w2_ = std::make_unique<Linear>(F, C, false, proj_std, rng); // residual projection
        const ops::QuantMode qm = quant_mode_from(cfg.quant);
        w1_->set_quant(qm);
        if (w3_) w3_->set_quant(qm);
        w2_->set_quant(qm);
        absorb("w1", *w1_);
        if (w3_) absorb("w3", *w3_);
        absorb("w2", *w2_);
    }

    Var forward(const Var& x) const override {
        if (swiglu_) {
            Var gate = ops::silu(w1_->forward(x));
            Var up = w3_->forward(x);
            return w2_->forward(ops::mul(gate, up));
        }
        Var h = w1_->forward(x);
        return w2_->forward(relu2_ ? ops::relu2(h) : ops::gelu(h));
    }

private:
    bool swiglu_, relu2_;
    std::unique_ptr<Linear> w1_, w2_, w3_;
};

// Mixture-of-experts MLP. Router scoring is softmax (classic) or sigmoid
// (DeepSeek V3/K2); top-k selection may be steered by a per-expert balance
// bias (aux-loss-free balancing: bias_update() nudges it ±gamma by observed
// load after each optimizer step), gates are renormalized or raw, and the
// routed sum can be rescaled (routed_scaling_factor). Experts may use their
// own d_ff (moe_d_ff). v1 computes EVERY expert densely and weights by the
// (mostly zero) gates — correctness first; token gather/scatter sparsity is
// a later optimization. The optional differentiable aux loss
// E · Σ_e (mean_i score)² remains available (moe_aux_weight).
class MoEMLP : public MLPBase {
public:
    MoEMLP(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng)
        : n_experts_(cfg.n_experts), top_k_(cfg.moe_top_k),
          aux_weight_(cfg.moe_aux_weight),
          sigmoid_(cfg.moe_scoring == "sigmoid"),
          norm_topk_(cfg.moe_norm_topk),
          routed_scale_(cfg.routed_scaling_factor),
          track_load_(cfg.moe_bias_gamma > 0.0f) {
        router_ = std::make_unique<Linear>(cfg.d_model, n_experts_, false, 0.02f, rng);
        absorb("router", *router_);
        ModelConfig ecfg = cfg;
        if (cfg.moe_d_ff > 0) ecfg.d_ff = cfg.moe_d_ff;
        experts_.reserve(size_t(n_experts_));
        for (int64_t e = 0; e < n_experts_; ++e) {
            experts_.push_back(std::make_unique<MLP>(ecfg, proj_std, rng));
            absorb("experts." + std::to_string(e), *experts_.back());
        }
        // Always-active experts, added ungated (DeepSeek/Kimi style): they
        // absorb common knowledge so routed experts can specialize.
        for (int64_t s = 0; s < cfg.n_shared_experts; ++s) {
            shared_.push_back(std::make_unique<MLP>(ecfg, proj_std, rng));
            absorb("shared." + std::to_string(s), *shared_.back());
        }
        bias_ = Tensor::zeros({n_experts_});
        counts_ = Tensor::zeros({n_experts_});
        // Registered (grad-free) so checkpoints and .forge snapshots carry
        // the balance state; the optimizer skips non-grad params.
        register_param("router.balance_bias", Var(bias_, /*requires_grad=*/false));
    }

    Var forward(const Var& x) const override {
        const int64_t N = x.value().size(0);
        Var logits = router_->forward(x); // [N, E]
        Var probs = sigmoid_ ? ops::sigmoid(logits) : ops::softmax(logits);
        Var gates = ops::topk_renorm(probs, bias_, top_k_, norm_topk_);

        Var y = ops::row_scale(experts_[0]->forward(x), gates, 0);
        for (int64_t e = 1; e < n_experts_; ++e)
            y = ops::add(y, ops::row_scale(experts_[size_t(e)]->forward(x), gates, e));
        if (routed_scale_ != 1.0f) y = ops::scale(y, routed_scale_);
        for (const auto& s : shared_) y = ops::add(y, s->forward(x));

        if (track_load_) {
            if (ops::backend() == ops::Backend::Metal)
                metal::expert_counts(gates.value(), counts_);
            else
                cpu::expert_counts(gates.value(), counts_);
        }

        if (aux_weight_ > 0.0f) {
            // Column means of probs via a constant 1/N row — every step is an
            // existing autograd op, so the router gets balance gradients.
            Var mean_row(Tensor::full({1, N}, 1.0f / float(N)), /*requires_grad=*/false);
            Var col_mean = ops::matmul(mean_row, probs);            // [1, E]
            Var sq = ops::mul(col_mean, col_mean);                  // [1, E]
            Var ones_e(Tensor::full({1, n_experts_}, 1.0f), /*requires_grad=*/false);
            Var sum = ops::matmul(sq, ones_e, false, true);         // [1, 1]
            aux_ = ops::scale(sum.reshaped({1}), aux_weight_ * float(n_experts_));
        } else {
            aux_ = Var();
        }
        return y;
    }

    Var aux() const override { return aux_; }

    // V3 noaux rule: overloaded experts (above mean load) lose selection
    // bias, underloaded gain. Called by the trainer AFTER the optimizer's
    // final sync — CPU reads/writes are safe there.
    void bias_update(float gamma) override {
        float* c = counts_.data<float>();
        float* b = bias_.data<float>();
        double total = 0.0;
        for (int64_t e = 0; e < n_experts_; ++e) total += c[e];
        const float mean = float(total / double(n_experts_));
        for (int64_t e = 0; e < n_experts_; ++e) {
            if (c[e] > mean) b[e] -= gamma;
            else if (c[e] < mean) b[e] += gamma;
            c[e] = 0.0f;
        }
    }

private:
    int64_t n_experts_, top_k_;
    float aux_weight_;
    bool sigmoid_, norm_topk_;
    float routed_scale_;
    bool track_load_;
    std::unique_ptr<Linear> router_;
    std::vector<std::unique_ptr<MLP>> experts_;
    std::vector<std::unique_ptr<MLP>> shared_;
    Tensor bias_;           // [E] selection bias (noaux balancing state)
    mutable Tensor counts_; // [E] load accumulator between bias updates
    mutable Var aux_; // set by the last forward; consumed by Transformer::loss
};

} // namespace forge::nn