// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Codable mirror of Forge's JSON config (src/nn/config.h, RESEARCH.md §2).
// Field names are byte-compatible with the C++ parser; defaults match the
// C++ defaults so a partially-specified JSON round-trips identically.
import Foundation

struct ModelConfig: Codable, Equatable {
    var name = "model"
    var nLayers = 6
    var dModel = 384
    var nHeads = 6
    var nKvHeads = 6
    var dFf = 1024
    var vocabSize = 4096
    var contextLength = 512
    var tiedEmbeddings = true
    var useRope = true
    var ropeTheta = 10000.0
    var norm = "rmsnorm"          // rmsnorm | layernorm
    var normEps = 1e-6
    var activation = "swiglu"     // swiglu | gelu | relu2
    var dropout = 0.0
    var quant = "none"            // none | int8 | ternary
    var qkNorm = false
    var finalSoftcap = 0.0
    var scaleEmbeddings = false
    var attentionBias = false
    var headDimOverride = 0       // JSON key "head_dim"; 0 = d_model/n_heads
    var nopeEvery = 0
    var normPlacement = "pre"     // pre | post | sandwich
    var ropeScaleFactor = 0.0
    var ropeScaleLow = 1.0
    var ropeScaleHigh = 4.0
    var ropeScaleOrigCtx = 8192
    var slidingWindow = 0
    var slidingGlobalEvery = 0
    var ropeThetaGlobal = 0.0
    var attnSoftcap = 0.0
    var nExperts = 0
    var moeTopK = 2
    var moeAuxWeight = 0.01
    var nSharedExperts = 0
    var moeScoring = "softmax"    // softmax | sigmoid
    var moeNormTopk = true
    var routedScalingFactor = 1.0
    var moeDFf = 0
    var firstKDense = 0
    var moeBiasGamma = 0.0

    enum CodingKeys: String, CodingKey {
        case name
        case nLayers = "n_layers", dModel = "d_model", nHeads = "n_heads"
        case nKvHeads = "n_kv_heads", dFf = "d_ff", vocabSize = "vocab_size"
        case contextLength = "context_length", tiedEmbeddings = "tied_embeddings"
        case useRope = "use_rope", ropeTheta = "rope_theta", norm
        case normEps = "norm_eps", activation, dropout, quant
        case qkNorm = "qk_norm", finalSoftcap = "final_softcap"
        case scaleEmbeddings = "scale_embeddings", attentionBias = "attention_bias"
        case headDimOverride = "head_dim", nopeEvery = "nope_every"
        case normPlacement = "norm_placement"
        case ropeScaleFactor = "rope_scale_factor", ropeScaleLow = "rope_scale_low"
        case ropeScaleHigh = "rope_scale_high", ropeScaleOrigCtx = "rope_scale_orig_ctx"
        case slidingWindow = "sliding_window"
        case slidingGlobalEvery = "sliding_global_every"
        case ropeThetaGlobal = "rope_theta_global", attnSoftcap = "attn_softcap"
        case nExperts = "n_experts", moeTopK = "moe_top_k"
        case moeAuxWeight = "moe_aux_weight", nSharedExperts = "n_shared_experts"
        case moeScoring = "moe_scoring", moeNormTopk = "moe_norm_topk"
        case routedScalingFactor = "routed_scaling_factor", moeDFf = "moe_d_ff"
        case firstKDense = "first_k_dense", moeBiasGamma = "moe_bias_gamma"
    }

    // Forge tolerates missing keys everywhere — mirror that.
    init() {}
    init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: CodingKeys.self)
        func g<T: Decodable>(_ k: CodingKeys, _ d: T) -> T {
            (try? c.decodeIfPresent(T.self, forKey: k)) as? T ?? d
        }
        name = g(.name, name); nLayers = g(.nLayers, nLayers)
        dModel = g(.dModel, dModel); nHeads = g(.nHeads, nHeads)
        nKvHeads = g(.nKvHeads, nHeads); dFf = g(.dFf, dFf)
        vocabSize = g(.vocabSize, vocabSize)
        contextLength = g(.contextLength, contextLength)
        tiedEmbeddings = g(.tiedEmbeddings, tiedEmbeddings)
        useRope = g(.useRope, useRope); ropeTheta = g(.ropeTheta, ropeTheta)
        norm = g(.norm, norm); normEps = g(.normEps, normEps)
        activation = g(.activation, activation); dropout = g(.dropout, dropout)
        quant = g(.quant, quant); qkNorm = g(.qkNorm, qkNorm)
        finalSoftcap = g(.finalSoftcap, finalSoftcap)
        scaleEmbeddings = g(.scaleEmbeddings, scaleEmbeddings)
        attentionBias = g(.attentionBias, attentionBias)
        headDimOverride = g(.headDimOverride, headDimOverride)
        nopeEvery = g(.nopeEvery, nopeEvery)
        normPlacement = g(.normPlacement, normPlacement)
        ropeScaleFactor = g(.ropeScaleFactor, ropeScaleFactor)
        ropeScaleLow = g(.ropeScaleLow, ropeScaleLow)
        ropeScaleHigh = g(.ropeScaleHigh, ropeScaleHigh)
        ropeScaleOrigCtx = g(.ropeScaleOrigCtx, ropeScaleOrigCtx)
        slidingWindow = g(.slidingWindow, slidingWindow)
        slidingGlobalEvery = g(.slidingGlobalEvery, slidingGlobalEvery)
        ropeThetaGlobal = g(.ropeThetaGlobal, ropeThetaGlobal)
        attnSoftcap = g(.attnSoftcap, attnSoftcap)
        nExperts = g(.nExperts, nExperts); moeTopK = g(.moeTopK, moeTopK)
        moeAuxWeight = g(.moeAuxWeight, moeAuxWeight)
        nSharedExperts = g(.nSharedExperts, nSharedExperts)
        moeScoring = g(.moeScoring, moeScoring)
        moeNormTopk = g(.moeNormTopk, moeNormTopk)
        routedScalingFactor = g(.routedScalingFactor, routedScalingFactor)
        moeDFf = g(.moeDFf, moeDFf); firstKDense = g(.firstKDense, firstKDense)
        moeBiasGamma = g(.moeBiasGamma, moeBiasGamma)
    }

    var headDim: Int { headDimOverride > 0 ? headDimOverride : dModel / max(nHeads, 1) }

    // Same formula as ModelConfig::num_params() — cross-checked by
    // ParamCountTests against `forge info`.
    var paramCount: Int {
        let hd = headDim
        var attn = dModel * nHeads * hd + 2 * dModel * nKvHeads * hd
            + nHeads * hd * dModel
        if attentionBias { attn += (nHeads + 2 * nKvHeads) * hd }
        let actMats = activation == "swiglu" ? 3 : 2
        let mlpDense = actMats * dModel * dFf
        let expertDff = (nExperts > 0 && moeDFf > 0) ? moeDFf : dFf
        let mlpMoe = (nExperts + nSharedExperts) * actMats * dModel * expertDff
            + nExperts * dModel
        let nMoeLayers = nExperts > 0 ? nLayers - min(firstKDense, nLayers) : 0
        let mlpTotal = nMoeLayers * mlpMoe + (nLayers - nMoeLayers) * mlpDense
        let normsPerLayer = normPlacement == "sandwich" ? 4 : 2
        var norms = (norm == "layernorm" ? 2 : 1) * dModel
            * (normsPerLayer * nLayers + 1)
        if qkNorm { norms += 2 * hd * nLayers }
        var total = nLayers * attn + mlpTotal + norms + vocabSize * dModel
        if !tiedEmbeddings { total += vocabSize * dModel }
        if !useRope { total += contextLength * dModel }
        return total
    }

    // Mirrors the C++ parse-time validation; returns human-actionable errors.
    var validationErrors: [String] {
        var e: [String] = []
        if headDimOverride == 0 && nHeads > 0 && dModel % nHeads != 0 {
            e.append("d_model doit être divisible par n_heads (ou fixer head_dim)")
        }
        if nKvHeads > 0 && nHeads % nKvHeads != 0 {
            e.append("n_heads doit être divisible par n_kv_heads")
        }
        if headDim % 2 != 0 { e.append("head_dim doit être pair (paires RoPE)") }
        if !["rmsnorm", "layernorm"].contains(norm) { e.append("norm invalide") }
        if !["swiglu", "gelu", "relu2"].contains(activation) {
            e.append("activation invalide")
        }
        if !["pre", "post", "sandwich"].contains(normPlacement) {
            e.append("norm_placement invalide")
        }
        if !["none", "int8", "ternary"].contains(quant) { e.append("quant invalide") }
        if nExperts > 0 && !(1...nExperts).contains(moeTopK) {
            e.append("moe_top_k doit être dans [1, n_experts]")
        }
        if nSharedExperts > 0 && nExperts == 0 {
            e.append("n_shared_experts requiert n_experts > 0")
        }
        if firstKDense < 0 || firstKDense > nLayers {
            e.append("first_k_dense doit être dans [0, n_layers]")
        }
        for (v, n) in [(nLayers, "n_layers"), (dModel, "d_model"), (nHeads, "n_heads"),
                       (dFf, "d_ff"), (vocabSize, "vocab_size"),
                       (contextLength, "context_length")] where v <= 0 {
            e.append("\(n) doit être > 0")
        }
        return e
    }
}

struct TrainConfig: Codable, Equatable {
    var lr = 6e-4
    var minLrRatio = 0.1
    var warmupSteps = 2000
    var maxSteps = 100_000
    var schedule = "cosine"       // cosine | wsd
    var wsdDecayFrac = 0.15
    var optimizer = "adamw"       // adamw | muon
    var muonLr = 0.02
    var muonMomentum = 0.95
    var beta1 = 0.9
    var beta2 = 0.95
    var eps = 1e-8
    var weightDecay = 0.1
    var gradClip = 1.0
    var batchSize = 32
    var gradAccumSteps = 1
    var precision = "f32"
    var checkpointEvery = 1000
    var forgeSave = true
    var forgeDtype = "f32"
    var evalEvery = 500
    var evalBatches = 20
    var seed = 1337
    var deterministic = false

    enum CodingKeys: String, CodingKey {
        case lr, minLrRatio = "min_lr_ratio", warmupSteps = "warmup_steps"
        case maxSteps = "max_steps", schedule, wsdDecayFrac = "wsd_decay_frac"
        case optimizer, muonLr = "muon_lr", muonMomentum = "muon_momentum"
        case beta1, beta2, eps, weightDecay = "weight_decay"
        case gradClip = "grad_clip", batchSize = "batch_size"
        case gradAccumSteps = "grad_accum_steps", precision
        case checkpointEvery = "checkpoint_every", forgeSave = "forge_save"
        case forgeDtype = "forge_dtype", evalEvery = "eval_every"
        case evalBatches = "eval_batches", seed, deterministic
    }

    init() {}
    init(from decoder: Decoder) throws {
        let c = try decoder.container(keyedBy: CodingKeys.self)
        func g<T: Decodable>(_ k: CodingKeys, _ d: T) -> T {
            (try? c.decodeIfPresent(T.self, forKey: k)) as? T ?? d
        }
        lr = g(.lr, lr); minLrRatio = g(.minLrRatio, minLrRatio)
        warmupSteps = g(.warmupSteps, warmupSteps); maxSteps = g(.maxSteps, maxSteps)
        schedule = g(.schedule, schedule)
        wsdDecayFrac = g(.wsdDecayFrac, wsdDecayFrac)
        optimizer = g(.optimizer, optimizer); muonLr = g(.muonLr, muonLr)
        muonMomentum = g(.muonMomentum, muonMomentum)
        beta1 = g(.beta1, beta1); beta2 = g(.beta2, beta2); eps = g(.eps, eps)
        weightDecay = g(.weightDecay, weightDecay); gradClip = g(.gradClip, gradClip)
        batchSize = g(.batchSize, batchSize)
        gradAccumSteps = g(.gradAccumSteps, gradAccumSteps)
        precision = g(.precision, precision)
        checkpointEvery = g(.checkpointEvery, checkpointEvery)
        forgeSave = g(.forgeSave, forgeSave); forgeDtype = g(.forgeDtype, forgeDtype)
        evalEvery = g(.evalEvery, evalEvery); evalBatches = g(.evalBatches, evalBatches)
        seed = g(.seed, seed); deterministic = g(.deterministic, deterministic)
    }

    var validationErrors: [String] {
        var e: [String] = []
        if lr <= 0 { e.append("lr doit être > 0") }
        if warmupSteps > maxSteps { e.append("warmup_steps ≤ max_steps requis") }
        if !["cosine", "wsd"].contains(schedule) { e.append("schedule invalide") }
        if !["adamw", "muon"].contains(optimizer) { e.append("optimizer invalide") }
        if !(0.0..<1.0).contains(wsdDecayFrac) || wsdDecayFrac <= 0 {
            e.append("wsd_decay_frac doit être dans (0, 1)")
        }
        if !["f32", "f16", "bf16"].contains(forgeDtype) {
            e.append("forge_dtype invalide")
        }
        if batchSize <= 0 || gradAccumSteps <= 0 || maxSteps <= 0 {
            e.append("batch/accum/max_steps doivent être > 0")
        }
        return e
    }

    // LR schedule preview (same math as src/train/scheduler.h).
    func lrAt(step: Int) -> Double {
        let minLr = lr * minLrRatio
        if step < warmupSteps {
            return lr * Double(step + 1) / Double(warmupSteps + 1)
        }
        if schedule == "wsd" {
            let decaySteps = max(1, Int(Double(maxSteps) * wsdDecayFrac))
            let decayStart = maxSteps - decaySteps
            if step < decayStart { return lr }
            if step >= maxSteps { return minLr }
            let ratio = Double(step - decayStart) / Double(decaySteps)
            return minLr + (lr - minLr) * (1.0 - ratio.squareRoot())
        }
        if step >= maxSteps { return minLr }
        let ratio = Double(step - warmupSteps) / Double(maxSteps - warmupSteps)
        return minLr + 0.5 * (1.0 + cos(.pi * ratio)) * (lr - minLr)
    }
}

struct ForgeConfig: Codable, Equatable {
    var model = ModelConfig()
    var train = TrainConfig()

    var tokensPerStep: Int { train.batchSize * model.contextLength * train.gradAccumSteps }
    var totalTokens: Int { tokensPerStep * train.maxSteps }
    var validationErrors: [String] {
        model.validationErrors + train.validationErrors
    }

    static func load(from url: URL) throws -> ForgeConfig {
        try JSONDecoder().decode(ForgeConfig.self, from: Data(contentsOf: url))
    }

    func exportJSON() throws -> Data {
        let enc = JSONEncoder()
        enc.outputFormatting = [.prettyPrinted, .sortedKeys]
        return try enc.encode(self)
    }
}