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

#include "nn/attention.h"
#include "nn/config.h"
#include "nn/embedding.h"
#include "nn/mlp.h"

#include <cmath>
#include <memory>
#include <vector>

namespace forge::nn {

// Norm wrapper so blocks don't branch on the norm flavour.
class Norm : public Module {
public:
    Norm(const ModelConfig& cfg) : layernorm_(cfg.norm == "layernorm"), eps_(cfg.norm_eps) {
        w_ = register_param("weight", Tensor::full({cfg.d_model}, 1.0f));
        if (layernorm_) b_ = register_param("bias", Tensor::zeros({cfg.d_model}));
    }

    Var forward(const Var& x) const {
        return layernorm_ ? ops::layernorm(x, w_, b_, eps_) : ops::rmsnorm(x, w_, eps_);
    }

private:
    bool layernorm_;
    float eps_;
    Var w_, b_;
};

// Residual block; norm placement is config-selected:
//   pre      (default): x += attn(norm1(x));        x += mlp(norm2(x))
//   post     (OLMo 2):  x += norm1(attn(x));        x += norm2(mlp(x))
//   sandwich (Gemma):   x += norm1p(attn(norm1(x))); x += norm2p(mlp(norm2(x)))
class TransformerBlock : public Module {
public:
    TransformerBlock(const ModelConfig& cfg, float proj_std, std::mt19937_64& rng,
                     int64_t layer_idx = 0)
        : placement_(cfg.norm_placement == "pre" ? Placement::Pre
                   : cfg.norm_placement == "post" ? Placement::Post
                                                  : Placement::Sandwich),
          norm1_(std::make_unique<Norm>(cfg)),
          attn_(std::make_unique<CausalSelfAttention>(cfg, proj_std, rng, layer_idx)),
          norm2_(std::make_unique<Norm>(cfg)) {
        if (cfg.n_experts > 0 && layer_idx >= cfg.first_k_dense)
            mlp_ = std::make_unique<MoEMLP>(cfg, proj_std, rng);
        else
            mlp_ = std::make_unique<MLP>(cfg, proj_std, rng);
        absorb("norm1", *norm1_);
        absorb("attn", *attn_);
        absorb("norm2", *norm2_);
        absorb("mlp", *mlp_);
        if (placement_ == Placement::Sandwich) {
            norm1_post_ = std::make_unique<Norm>(cfg);
            norm2_post_ = std::make_unique<Norm>(cfg);
            absorb("norm1_post", *norm1_post_);
            absorb("norm2_post", *norm2_post_);
        }
    }

    Var forward(const Var& x) const {
        const int64_t B = x.value().size(0), T = x.value().size(1), C = x.value().size(2);
        Var a;
        switch (placement_) {
            case Placement::Pre:      a = attn_->forward(norm1_->forward(x)); break;
            case Placement::Post:     a = norm1_->forward(attn_->forward(x)); break;
            case Placement::Sandwich:
                a = norm1_post_->forward(attn_->forward(norm1_->forward(x)));
                break;
        }
        Var h = ops::add(x, a);
        Var m;
        switch (placement_) {
            case Placement::Pre:
                m = mlp_->forward(norm2_->forward(h).reshaped({B * T, C}))
                        .reshaped({B, T, C});
                break;
            case Placement::Post:
                m = norm2_->forward(
                        mlp_->forward(h.reshaped({B * T, C})).reshaped({B, T, C}));
                break;
            case Placement::Sandwich:
                m = norm2_post_->forward(
                        mlp_->forward(norm2_->forward(h).reshaped({B * T, C}))
                            .reshaped({B, T, C}));
                break;
        }
        return ops::add(h, m);
    }

    // MoE load-balance loss of the last forward (undefined for dense MLP).
    Var moe_aux() const { return mlp_->aux(); }
    void moe_bias_update(float gamma) { mlp_->bias_update(gamma); }

private:
    enum class Placement { Pre, Post, Sandwich };
    Placement placement_;
    std::unique_ptr<Norm> norm1_;
    std::unique_ptr<AttentionBase> attn_;
    std::unique_ptr<Norm> norm2_;
    std::unique_ptr<MLPBase> mlp_;
    std::unique_ptr<Norm> norm1_post_, norm2_post_; // sandwich only
};

// Decoder-only transformer, entirely shaped by ModelConfig.
class Transformer : public Module {
public:
    Transformer(const ModelConfig& cfg, uint64_t seed) : cfg_(cfg) {
        std::mt19937_64 rng(seed);
        const float std = 0.02f;
        // Residual projections scaled by 1/sqrt(2L): one attn + one mlp
        // residual add per layer (nanoGPT).
        const float proj_std = std / std::sqrt(2.0f * float(cfg.n_layers));

        tok_emb_ = std::make_unique<Embedding>(cfg.vocab_size, cfg.d_model, std, rng);
        absorb("tok_emb", *tok_emb_);
        if (!cfg.use_rope) {
            pos_emb_ = std::make_unique<Embedding>(cfg.context_length, cfg.d_model, std, rng);
            absorb("pos_emb", *pos_emb_);
        }
        for (int64_t i = 0; i < cfg.n_layers; ++i) {
            blocks_.push_back(std::make_unique<TransformerBlock>(cfg, proj_std, rng, i));
            absorb("blocks." + std::to_string(i), *blocks_.back());
        }
        final_norm_ = std::make_unique<Norm>(cfg);
        absorb("final_norm", *final_norm_);

        if (cfg.tied_embeddings) {
            lm_head_weight_ = register_param("lm_head.weight", tok_emb_->weight());
        } else {
            lm_head_weight_ = register_param(
                "lm_head.weight", normal_init({cfg.vocab_size, cfg.d_model}, std, rng));
        }
    }

    // ids: [B, T] u16/i32 → logits [B*T, V]
    Var forward(const Tensor& ids) const {
        const int64_t B = ids.shape()[0], T = ids.shape()[1];
        Var x = tok_emb_->forward(ids);
        if (cfg_.scale_embeddings)
            x = ops::scale(x, std::sqrt(float(cfg_.d_model)));
        if (pos_emb_) {
            Tensor pos = Tensor::empty({1, T}, DType::I32);
            for (int64_t t = 0; t < T; ++t) pos.data<int32_t>()[t] = int32_t(t);
            Var p = pos_emb_->forward(pos); // [1, T, C]
            // Broadcast over batch: viewed as [B, T*C] rows + a [T*C] "bias",
            // add_bias sums the positional grad over the batch — exactly right.
            Var xb = x.reshaped({B, T * cfg_.d_model});
            Var pflat = p.reshaped({T * cfg_.d_model});
            x = ops::add_bias(xb, pflat).reshaped({B, T, cfg_.d_model});
        }
        for (const auto& blk : blocks_) x = blk->forward(x);
        x = final_norm_->forward(x);
        Var x2d = x.reshaped({B * T, cfg_.d_model});
        Var logits = ops::matmul(x2d, lm_head_weight_, false, true); // [B*T, V]
        if (cfg_.final_softcap > 0.0f)
            logits = ops::softcap(logits, cfg_.final_softcap);
        return logits;
    }

    // targets: [B, T] with ignore_index=-1 → scalar mean CE loss. With MoE,
    // each block's load-balance term is added (also during eval — its scale
    // is ~moe_aux_weight, negligible next to CE but keeps train/val
    // comparable).
    Var loss(const Tensor& ids, const Tensor& targets) const {
        Var logits = forward(ids);
        Tensor tflat = targets;
        Var l = ops::cross_entropy(logits, tflat.view({targets.numel()}));
        if (cfg_.n_experts > 0 && cfg_.moe_aux_weight > 0.0f) {
            for (const auto& blk : blocks_) {
                Var a = blk->moe_aux();
                if (a.defined()) l = ops::add(l, a);
            }
        }
        return l;
    }

    const ModelConfig& config() const { return cfg_; }

    // noaux balancing step (V3): call after the optimizer sync, CPU-side.
    void update_moe_bias(float gamma) {
        for (auto& blk : blocks_) blk->moe_bias_update(gamma);
    }

private:
    ModelConfig cfg_;
    std::unique_ptr<Embedding> tok_emb_;
    std::unique_ptr<Embedding> pos_emb_; // null when RoPE
    std::vector<std::unique_ptr<TransformerBlock>> blocks_;
    std::unique_ptr<Norm> final_norm_;
    Var lm_head_weight_;
};

} // namespace forge::nn