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

#include <algorithm>
#include <cstdint>
#include <fstream>
#include <stdexcept>
#include <string>

#include <nlohmann/json.hpp>

namespace forge {

// The entire architecture comes from here — no hardcoded model sizes
// anywhere else in the codebase.
struct ModelConfig {
    std::string name = "model";
    int64_t n_layers = 6;
    int64_t d_model = 384;
    int64_t n_heads = 6;
    int64_t n_kv_heads = 6;      // < n_heads => GQA
    int64_t d_ff = 1024;
    int64_t vocab_size = 4096;
    int64_t context_length = 512;
    bool tied_embeddings = true;
    bool use_rope = true;        // false => learned positional embeddings
    float rope_theta = 10000.0f;
    std::string norm = "rmsnorm";     // "rmsnorm" | "layernorm"
    float norm_eps = 1e-6f;
    std::string activation = "swiglu"; // "swiglu" | "gelu"
    float dropout = 0.0f;
    // Architecture-variant knobs (all off by default = the classic recipe):
    bool qk_norm = false;         // RMSNorm on q/k per head before RoPE (Qwen3/Gemma3)
    float final_softcap = 0.0f;   // logits = cap*tanh(logits/cap) (Gemma2: 30)
    bool scale_embeddings = false;// x *= sqrt(d_model) after embedding (Gemma)
    bool attention_bias = false;  // bias on q/k/v projections only (Qwen2.5)
    int64_t head_dim_override = 0;// decouple head_dim from d_model/n_heads (Qwen3)
    int64_t nope_every = 0;       // skip RoPE every Nth layer (SmolLM3: 4)
    std::string norm_placement = "pre"; // "pre" | "post" (OLMo2) | "sandwich" (Gemma)
    // RoPE scaling, HF "llama3" style (0 = off). Applied to the inv-freq table.
    float rope_scale_factor = 0.0f;      // llama3: 32
    float rope_scale_low = 1.0f;         // low_freq_factor
    float rope_scale_high = 4.0f;        // high_freq_factor
    int64_t rope_scale_orig_ctx = 8192;  // original_max_position_embeddings
    // Sliding-window attention (Mistral/Gemma3). 0 = full attention everywhere.
    int64_t sliding_window = 0;
    int64_t sliding_global_every = 0; // every Nth layer is global (Gemma3: 6); 0 = none
    float rope_theta_global = 0.0f;   // theta for global layers (Gemma3: 1e6); 0 = theta
    float attn_softcap = 0.0f;        // cap on attention scores (Gemma2: 50)
    // Quantization-aware training: linear weights are fake-quantized each
    // forward (per-row scales, straight-through estimator in backward).
    // Master weights and the optimizer stay f32.
    std::string quant = "none";       // "none" | "int8" | "ternary"
    // Mixture-of-experts MLP: n_experts > 0 replaces each block's MLP with
    // `n_experts` expert MLPs (d_ff each) + a top-k softmax router. v1
    // computes every expert densely and masks — correctness first.
    int64_t n_experts = 0;            // 0 => dense MLP
    int64_t moe_top_k = 2;
    float moe_aux_weight = 0.01f;     // load-balance loss: E * sum_e mean_gate_e^2
    int64_t n_shared_experts = 0;     // always-active experts (DeepSeek/Kimi style)
    // DeepSeek-V3-style routing refinements:
    std::string moe_scoring = "softmax"; // "softmax" | "sigmoid" (V3/K2)
    bool moe_norm_topk = true;        // renormalize kept gates to sum 1
    float routed_scaling_factor = 1.0f; // y = shared + factor * sum(gated) (V3: 2.5)
    int64_t moe_d_ff = 0;             // per-expert d_ff; 0 = d_ff (V3: 2048 vs 18432)
    int64_t first_k_dense = 0;        // first k layers keep a dense MLP (V3: 3, K2: 1)
    // Aux-loss-FREE balancing (V3 "noaux"): select top-k on s+b but gate on s;
    // after each step b += gamma for underloaded experts, -= gamma for
    // overloaded. 0 disables (bias stays zero).
    float moe_bias_gamma = 0.0f;      // V3: 0.001

    int64_t head_dim() const {
        return head_dim_override > 0 ? head_dim_override : d_model / n_heads;
    }

    // True when layer i (0-based) attends globally; only meaningful with
    // sliding_window > 0.
    bool layer_is_global(int64_t i) const {
        if (sliding_window <= 0) return true;
        return sliding_global_every > 0 && (i + 1) % sliding_global_every == 0;
    }
    bool layer_uses_rope(int64_t i) const {
        return use_rope && !(nope_every > 0 && (i + 1) % nope_every == 0);
    }

    // Total parameter count (SwiGLU MLP has three mats, GELU/ReLU² have two).
    int64_t num_params() const {
        const int64_t hd = head_dim();
        int64_t attn = d_model * n_heads * hd             // wq
                     + 2 * d_model * n_kv_heads * hd      // wk, wv
                     + n_heads * hd * d_model;            // wo
        if (attention_bias) attn += (n_heads + 2 * n_kv_heads) * hd;
        const int64_t act_mats = (activation == "swiglu") ? 3 : 2;
        const int64_t mlp_dense = act_mats * d_model * d_ff;
        const int64_t expert_dff = (n_experts > 0 && moe_d_ff > 0) ? moe_d_ff : d_ff;
        const int64_t mlp_moe = (n_experts + n_shared_experts) * act_mats * d_model *
                                    expert_dff +
                                n_experts * d_model;       // experts + router
        // first_k_dense layers keep the dense MLP (V3-style).
        const int64_t n_moe_layers =
            n_experts > 0 ? n_layers - std::min(first_k_dense, n_layers) : 0;
        const int64_t mlp_total =
            n_moe_layers * mlp_moe + (n_layers - n_moe_layers) * mlp_dense;
        const int64_t norms_per_layer = norm_placement == "sandwich" ? 4 : 2;
        int64_t norms = (norm == "layernorm" ? 2 : 1) * d_model *
                        (norms_per_layer * n_layers + 1);
        if (qk_norm) norms += 2 * head_dim() * n_layers;
        int64_t total = n_layers * attn + mlp_total + norms + vocab_size * d_model;
        if (!tied_embeddings) total += vocab_size * d_model;
        if (!use_rope) total += context_length * d_model;
        return total;
    }
};

struct TrainConfig {
    float lr = 6e-4f;
    float min_lr_ratio = 0.1f;       // min_lr = lr * ratio
    int64_t warmup_steps = 2000;
    int64_t max_steps = 100000;
    std::string schedule = "cosine"; // "cosine" | "wsd" (warmup-stable-decay)
    float wsd_decay_frac = 0.15f;    // WSD: fraction of max_steps in the cooldown
    // "muon": Newton-Schulz orthogonalized momentum on 2-D hidden matrices,
    // AdamW on embeddings/head/norms (Keller Jordan's Muon). muon_lr is the
    // peak LR of the Muon group; it follows the same schedule shape as `lr`.
    std::string optimizer = "adamw"; // "adamw" | "muon"
    float muon_lr = 0.02f;
    float muon_momentum = 0.95f;
    float beta1 = 0.9f;
    float beta2 = 0.95f;
    float eps = 1e-8f;
    float weight_decay = 0.1f;       // dim>=2 params only
    float grad_clip = 1.0f;          // 0 disables
    int64_t batch_size = 32;         // sequences per micro-batch
    int64_t grad_accum_steps = 1;
    std::string precision = "f32";   // "f32" | "f16" | "bf16" (compute dtype)
    int64_t checkpoint_every = 1000;
    // Also commit each checkpoint's weights into <out>/model.forge — the
    // git-style zero-copy format (src/core/fmodel.h). The .bin checkpoint
    // stays the resume source of truth (it carries optimizer state).
    bool forge_save = true;
    std::string forge_dtype = "f32"; // "f32" | "f16" | "bf16"
    int64_t eval_every = 500;
    int64_t eval_batches = 20;
    uint64_t seed = 1337;
    bool deterministic = false;
};

struct Config {
    ModelConfig model;
    TrainConfig train;
};

inline void from_json(const nlohmann::json& j, ModelConfig& c) {
    c.name = j.value("name", c.name);
    c.n_layers = j.value("n_layers", c.n_layers);
    c.d_model = j.value("d_model", c.d_model);
    c.n_heads = j.value("n_heads", c.n_heads);
    c.n_kv_heads = j.value("n_kv_heads", c.n_heads);
    c.d_ff = j.value("d_ff", c.d_ff);
    c.vocab_size = j.value("vocab_size", c.vocab_size);
    c.context_length = j.value("context_length", c.context_length);
    c.tied_embeddings = j.value("tied_embeddings", c.tied_embeddings);
    c.use_rope = j.value("use_rope", c.use_rope);
    c.rope_theta = j.value("rope_theta", c.rope_theta);
    c.norm = j.value("norm", c.norm);
    c.norm_eps = j.value("norm_eps", c.norm_eps);
    c.activation = j.value("activation", c.activation);
    c.dropout = j.value("dropout", c.dropout);
    c.quant = j.value("quant", c.quant);
    c.n_experts = j.value("n_experts", c.n_experts);
    c.moe_top_k = j.value("moe_top_k", c.moe_top_k);
    c.moe_aux_weight = j.value("moe_aux_weight", c.moe_aux_weight);
    c.n_shared_experts = j.value("n_shared_experts", c.n_shared_experts);
    c.moe_scoring = j.value("moe_scoring", c.moe_scoring);
    c.moe_norm_topk = j.value("moe_norm_topk", c.moe_norm_topk);
    c.routed_scaling_factor = j.value("routed_scaling_factor", c.routed_scaling_factor);
    c.moe_d_ff = j.value("moe_d_ff", c.moe_d_ff);
    c.first_k_dense = j.value("first_k_dense", c.first_k_dense);
    c.moe_bias_gamma = j.value("moe_bias_gamma", c.moe_bias_gamma);
    c.qk_norm = j.value("qk_norm", c.qk_norm);
    c.final_softcap = j.value("final_softcap", c.final_softcap);
    c.scale_embeddings = j.value("scale_embeddings", c.scale_embeddings);
    c.attention_bias = j.value("attention_bias", c.attention_bias);
    c.head_dim_override = j.value("head_dim", c.head_dim_override);
    c.nope_every = j.value("nope_every", c.nope_every);
    c.norm_placement = j.value("norm_placement", c.norm_placement);
    c.rope_scale_factor = j.value("rope_scale_factor", c.rope_scale_factor);
    c.rope_scale_low = j.value("rope_scale_low", c.rope_scale_low);
    c.rope_scale_high = j.value("rope_scale_high", c.rope_scale_high);
    c.rope_scale_orig_ctx = j.value("rope_scale_orig_ctx", c.rope_scale_orig_ctx);
    c.sliding_window = j.value("sliding_window", c.sliding_window);
    c.sliding_global_every = j.value("sliding_global_every", c.sliding_global_every);
    c.rope_theta_global = j.value("rope_theta_global", c.rope_theta_global);
    c.attn_softcap = j.value("attn_softcap", c.attn_softcap);

    if (c.head_dim_override == 0 && c.d_model % c.n_heads != 0)
        throw std::runtime_error("config: d_model must be divisible by n_heads "
                                 "(or set head_dim explicitly)");
    if (c.n_heads % c.n_kv_heads != 0)
        throw std::runtime_error("config: n_heads must be divisible by n_kv_heads");
    if (c.head_dim() % 2 != 0)
        throw std::runtime_error("config: head_dim must be even (RoPE pairs)");
    if (c.norm != "rmsnorm" && c.norm != "layernorm")
        throw std::runtime_error("config: norm must be rmsnorm or layernorm");
    if (c.activation != "swiglu" && c.activation != "gelu" && c.activation != "relu2")
        throw std::runtime_error("config: activation must be swiglu, gelu or relu2");
    if (c.norm_placement != "pre" && c.norm_placement != "post" &&
        c.norm_placement != "sandwich")
        throw std::runtime_error("config: norm_placement must be pre, post or sandwich");
    if (c.sliding_window < 0 || c.sliding_global_every < 0 || c.nope_every < 0)
        throw std::runtime_error("config: window/pattern values must be >= 0");
    if (c.attn_softcap < 0.0f)
        throw std::runtime_error("config: attn_softcap must be >= 0");
    if (c.quant != "none" && c.quant != "int8" && c.quant != "ternary")
        throw std::runtime_error("config: quant must be none, int8 or ternary");
    if (c.n_experts < 0)
        throw std::runtime_error("config: n_experts must be >= 0");
    if (c.n_experts > 0 && (c.moe_top_k < 1 || c.moe_top_k > c.n_experts))
        throw std::runtime_error("config: moe_top_k must be in [1, n_experts]");
    if (c.n_shared_experts < 0 || (c.n_shared_experts > 0 && c.n_experts == 0))
        throw std::runtime_error("config: n_shared_experts requires n_experts > 0");
    if (c.moe_scoring != "softmax" && c.moe_scoring != "sigmoid")
        throw std::runtime_error("config: moe_scoring must be softmax or sigmoid");
    if (c.first_k_dense < 0 || c.first_k_dense > c.n_layers)
        throw std::runtime_error("config: first_k_dense must be in [0, n_layers]");
    if (c.moe_d_ff < 0 || c.moe_bias_gamma < 0.0f)
        throw std::runtime_error("config: moe_d_ff and moe_bias_gamma must be >= 0");
    if (c.final_softcap < 0.0f)
        throw std::runtime_error("config: final_softcap must be >= 0");
}

inline void from_json(const nlohmann::json& j, TrainConfig& c) {
    c.lr = j.value("lr", c.lr);
    c.min_lr_ratio = j.value("min_lr_ratio", c.min_lr_ratio);
    c.warmup_steps = j.value("warmup_steps", c.warmup_steps);
    c.max_steps = j.value("max_steps", c.max_steps);
    c.schedule = j.value("schedule", c.schedule);
    c.wsd_decay_frac = j.value("wsd_decay_frac", c.wsd_decay_frac);
    c.optimizer = j.value("optimizer", c.optimizer);
    c.muon_lr = j.value("muon_lr", c.muon_lr);
    c.muon_momentum = j.value("muon_momentum", c.muon_momentum);
    c.beta1 = j.value("beta1", c.beta1);
    c.beta2 = j.value("beta2", c.beta2);
    c.eps = j.value("eps", c.eps);
    c.weight_decay = j.value("weight_decay", c.weight_decay);
    c.grad_clip = j.value("grad_clip", c.grad_clip);
    c.batch_size = j.value("batch_size", c.batch_size);
    c.grad_accum_steps = j.value("grad_accum_steps", c.grad_accum_steps);
    c.precision = j.value("precision", c.precision);
    c.checkpoint_every = j.value("checkpoint_every", c.checkpoint_every);
    c.forge_save = j.value("forge_save", c.forge_save);
    c.forge_dtype = j.value("forge_dtype", c.forge_dtype);
    c.eval_every = j.value("eval_every", c.eval_every);
    c.eval_batches = j.value("eval_batches", c.eval_batches);
    c.seed = j.value("seed", c.seed);
    c.deterministic = j.value("deterministic", c.deterministic);

    if (c.schedule != "cosine" && c.schedule != "wsd")
        throw std::runtime_error("config: schedule must be cosine or wsd");
    if (c.wsd_decay_frac <= 0.0f || c.wsd_decay_frac >= 1.0f)
        throw std::runtime_error("config: wsd_decay_frac must be in (0, 1)");
    if (c.optimizer != "adamw" && c.optimizer != "muon")
        throw std::runtime_error("config: optimizer must be adamw or muon");
    if (c.forge_dtype != "f32" && c.forge_dtype != "f16" && c.forge_dtype != "bf16")
        throw std::runtime_error("config: forge_dtype must be f32, f16 or bf16");
}

inline Config load_config(const std::string& path) {
    std::ifstream in(path);
    if (!in) throw std::runtime_error("config: cannot open " + path);
    nlohmann::json j = nlohmann::json::parse(in);
    Config cfg;
    if (j.contains("model")) j.at("model").get_to(cfg.model);
    if (j.contains("train")) j.at("train").get_to(cfg.train);
    return cfg;
}

} // namespace forge