// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// New Run editor: preset picker, the high-traffic model/train fields, the
// live derived panel (params, tokens/step, epochs, memory badge) and the LR
// schedule preview. Inline validation gates the Start button. (The long
// tail of variant knobs lives in the "Avancé" JSON editor — full form in M2.)
import Charts
import SwiftUI

struct NewRunSheet: View {
    @Environment(AppModel.self) private var app
    @Environment(\.dismiss) private var dismiss
    @State private var config = ForgeConfig()
    @State private var runName = "run"
    @State private var datasetPath = ""
    @State private var startError: String?

    private var datasets: [Dataset] { app.datasets() }
    private var dataset: Dataset? { datasets.first { $0.path == datasetPath } }
    private var errors: [String] {
        var e = config.validationErrors
        if datasetPath.isEmpty { e.append("choisir un dataset") }
        if let v = dataset?.vocabSize, v != config.model.vocabSize {
            e.append("vocab_size (\(config.model.vocabSize)) ≠ tokenizer du dataset (\(v))")
        }
        return e
    }

    var body: some View {
        VStack(spacing: 0) {
            HStack {
                Text("Nouveau run").font(.title2.weight(.semibold))
                Spacer()
                presetMenu
            }
            .padding()
            Divider()
            HSplitView {
                Form {
                    Section("Run") {
                        TextField("Nom", text: $runName)
                        Picker("Dataset", selection: $datasetPath) {
                            Text("—").tag("")
                            ForEach(datasets) { ds in
                                Text("\(ds.name) · vocab \(ds.vocabSize ?? 0) · \((ds.trainTokens ?? 0).formatted(.number.notation(.compactName))) tokens")
                                    .tag(ds.path)
                            }
                        }
                        .onChange(of: datasetPath) {
                            if let v = dataset?.vocabSize { config.model.vocabSize = v }
                        }
                    }
                    Section("Modèle") {
                        intField("n_layers", $config.model.nLayers)
                        intField("d_model", $config.model.dModel)
                        intField("n_heads", $config.model.nHeads)
                        intField("n_kv_heads (GQA)", $config.model.nKvHeads)
                        intField("d_ff", $config.model.dFf)
                        intField("context_length", $config.model.contextLength)
                        Picker("activation", selection: $config.model.activation) {
                            ForEach(["swiglu", "gelu", "relu2"], id: \.self) { Text($0) }
                        }
                        Picker("norm", selection: $config.model.norm) {
                            ForEach(["rmsnorm", "layernorm"], id: \.self) { Text($0) }
                        }
                        Toggle("tied_embeddings", isOn: $config.model.tiedEmbeddings)
                        Toggle("qk_norm", isOn: $config.model.qkNorm)
                        intField("sliding_window (0 = full)", $config.model.slidingWindow)
                    }
                    Section("Attention & positions") {
                        Toggle("use_rope", isOn: $config.model.useRope)
                        doubleField("rope_theta", $config.model.ropeTheta)
                        doubleField("rope_theta_global (0 = idem)", $config.model.ropeThetaGlobal)
                        intField("sliding_global_every (Gemma3: 6)", $config.model.slidingGlobalEvery)
                        intField("nope_every (SmolLM3: 4)", $config.model.nopeEvery)
                        intField("head_dim (0 = auto)", $config.model.headDimOverride)
                        Toggle("attention_bias (Qwen2.5)", isOn: $config.model.attentionBias)
                        doubleField("attn_softcap (Gemma2: 50)", $config.model.attnSoftcap)
                        doubleField("rope_scale_factor (llama3: 32, 0 = off)",
                                    $config.model.ropeScaleFactor)
                        if config.model.ropeScaleFactor > 0 {
                            doubleField("rope_scale_low", $config.model.ropeScaleLow)
                            doubleField("rope_scale_high", $config.model.ropeScaleHigh)
                            intField("rope_scale_orig_ctx", $config.model.ropeScaleOrigCtx)
                        }
                    }
                    Section("Normalisation & sorties") {
                        Picker("norm_placement", selection: $config.model.normPlacement) {
                            ForEach(["pre", "post", "sandwich"], id: \.self) { Text($0) }
                        }
                        doubleField("norm_eps", $config.model.normEps)
                        doubleField("final_softcap (Gemma2: 30)", $config.model.finalSoftcap)
                        Toggle("scale_embeddings (Gemma)", isOn: $config.model.scaleEmbeddings)
                        Picker("quant (QAT)", selection: $config.model.quant) {
                            ForEach(["none", "int8", "ternary"], id: \.self) { Text($0) }
                        }
                    }
                    Section("Mixture of Experts") {
                        intField("n_experts (0 = dense)", $config.model.nExperts)
                        if config.model.nExperts > 0 {
                            intField("moe_top_k", $config.model.moeTopK)
                            intField("n_shared_experts (V3: 1)", $config.model.nSharedExperts)
                            Picker("moe_scoring", selection: $config.model.moeScoring) {
                                ForEach(["softmax", "sigmoid"], id: \.self) { Text($0) }
                            }
                            Toggle("moe_norm_topk", isOn: $config.model.moeNormTopk)
                            doubleField("routed_scaling_factor (V3: 2.5)",
                                        $config.model.routedScalingFactor)
                            intField("moe_d_ff (0 = d_ff)", $config.model.moeDFf)
                            intField("first_k_dense (V3: 3)", $config.model.firstKDense)
                            doubleField("moe_bias_gamma (noaux, V3: 0.001)",
                                        $config.model.moeBiasGamma)
                            doubleField("moe_aux_weight", $config.model.moeAuxWeight)
                        }
                    }
                    Section("Entraînement") {
                        doubleField("lr", $config.train.lr)
                        intField("max_steps", $config.train.maxSteps)
                        intField("warmup_steps", $config.train.warmupSteps)
                        Picker("schedule", selection: $config.train.schedule) {
                            ForEach(["cosine", "wsd"], id: \.self) { Text($0) }
                        }
                        Picker("optimizer", selection: $config.train.optimizer) {
                            ForEach(["adamw", "muon"], id: \.self) { Text($0) }
                        }
                        intField("batch_size (micro)", $config.train.batchSize)
                        intField("grad_accum_steps", $config.train.gradAccumSteps)
                        intField("checkpoint_every", $config.train.checkpointEvery)
                        intField("eval_every", $config.train.evalEvery)
                    }
                    Section("Optimiseur avancé") {
                        if config.train.optimizer == "muon" {
                            doubleField("muon_lr", $config.train.muonLr)
                            doubleField("muon_momentum", $config.train.muonMomentum)
                        }
                        if config.train.schedule == "wsd" {
                            doubleField("wsd_decay_frac", $config.train.wsdDecayFrac)
                        }
                        doubleField("min_lr_ratio", $config.train.minLrRatio)
                        doubleField("beta1", $config.train.beta1)
                        doubleField("beta2", $config.train.beta2)
                        doubleField("weight_decay", $config.train.weightDecay)
                        doubleField("grad_clip (0 = off)", $config.train.gradClip)
                        intField("eval_batches", $config.train.evalBatches)
                        intField("seed", $config.train.seed)
                        Toggle("forge_save (.forge à chaque checkpoint)",
                               isOn: $config.train.forgeSave)
                        Picker("forge_dtype", selection: $config.train.forgeDtype) {
                            ForEach(["f32", "f16", "bf16"], id: \.self) { Text($0) }
                        }
                    }
                }
                .formStyle(.grouped)
                .frame(minWidth: 380)

                DerivedPanel(config: config, dataset: dataset)
                    .frame(minWidth: 300)
                    .padding()
            }
            Divider()
            HStack {
                if let first = errors.first {
                    Label(first, systemImage: "exclamationmark.triangle")
                        .foregroundStyle(.red)
                        .font(.callout)
                }
                if let startError {
                    Text(startError).foregroundStyle(.red).font(.callout)
                }
                Spacer()
                Button("Annuler") { dismiss() }
                Button("Démarrer l'entraînement") {
                    Task {
                        do {
                            config.model.name = runName
                            try await app.supervisor.start(
                                config: config, datasetPath: datasetPath,
                                name: runName)
                            dismiss()
                        } catch {
                            startError = error.localizedDescription
                        }
                    }
                }
                .keyboardShortcut(.defaultAction)
                .disabled(!errors.isEmpty)
            }
            .padding()
        }
        .frame(minWidth: 860, minHeight: 620)
        .onAppear {
            if let first = datasets.first { datasetPath = first.path }
        }
    }

    private var presetMenu: some View {
        HStack {
            Menu("Presets") {
                ForEach(Presets.all, id: \.name) { preset in
                    Button(preset.name) {
                        config = preset.config
                        runName = preset.config.model.name
                        if let v = dataset?.vocabSize { config.model.vocabSize = v }
                    }
                }
            }
            .frame(width: 110)
            Button("Importer…") { importConfig() }
            Button("Exporter…") { exportConfig() }
        }
    }

    private func importConfig() {
        let panel = NSOpenPanel()
        panel.allowedContentTypes = [.json]
        if panel.runModal() == .OK, let url = panel.url,
           let cfg = try? ForgeConfig.load(from: url) {
            config = cfg
            runName = cfg.model.name
        }
    }

    private func exportConfig() {
        let panel = NSSavePanel()
        panel.allowedContentTypes = [.json]
        panel.nameFieldStringValue = "\(runName).json"
        if panel.runModal() == .OK, let url = panel.url {
            var cfg = config
            cfg.model.name = runName
            try? cfg.exportJSON().write(to: url)
        }
    }

    private func intField(_ label: String, _ value: Binding<Int>) -> some View {
        TextField(label, value: value, format: .number)
    }

    private func doubleField(_ label: String, _ value: Binding<Double>) -> some View {
        TextField(label, value: value, format: .number.precision(.significantDigits(1...6)))
    }
}

struct DerivedPanel: View {
    let config: ForgeConfig
    let dataset: Dataset?

    var body: some View {
        VStack(alignment: .leading, spacing: 14) {
            Text("Dérivés").font(.headline)
            derived("Paramètres",
                    config.model.paramCount.formatted(.number.notation(.compactName)))
            derived("Tokens/step", config.tokensPerStep.formatted())
            derived("Tokens totaux",
                    config.totalTokens.formatted(.number.notation(.compactName)))
            if let t = dataset?.trainTokens, t > 0 {
                derived("Epochs sur le dataset",
                        String(format: "%.2f", Double(config.totalTokens) / Double(t)))
            }
            memoryBadge
            Divider()
            Text("Schedule LR").font(.headline)
            Chart {
                let steps = stride(from: 0, to: config.train.maxSteps,
                                   by: max(1, config.train.maxSteps / 200))
                ForEach(Array(steps), id: \.self) { s in
                    LineMark(x: .value("step", s),
                             y: .value("lr", config.train.lrAt(step: s)))
                        .foregroundStyle(.purple)
                }
            }
            .frame(height: 140)
            Spacer()
        }
    }

    private var memoryBadge: some View {
        let bytes = SystemInfo.estimatedTrainingBytes(config: config)
        let machine = SystemInfo.memoryBytes
        let gb = Double(bytes) / 1e9
        let ok = bytes < machine * 8 / 10
        return HStack {
            Image(systemName: ok ? "memorychip" : "exclamationmark.triangle.fill")
            Text(String(format: "~%.1f GB estimés / %.0f GB unifiés", gb,
                        Double(machine) / 1e9))
        }
        .font(.callout)
        .foregroundStyle(ok ? Color.secondary : Color.orange)
    }

    private func derived(_ label: String, _ value: String) -> some View {
        HStack {
            Text(label).foregroundStyle(.secondary)
            Spacer()
            Text(value).monospacedDigit()
        }
        .font(.callout)
    }
}

enum Presets {
    struct Preset {
        let name: String
        let config: ForgeConfig
    }

    // Mirrors configs/ in the forge repo (gpt-10m-1epoch, gpt-50m…).
    static let all: [Preset] = {
        var p10 = ForgeConfig()
        p10.model.name = "gpt-10m"
        p10.model.nLayers = 6; p10.model.dModel = 384; p10.model.nHeads = 6
        p10.model.nKvHeads = 6; p10.model.dFf = 1024; p10.model.vocabSize = 4096
        p10.model.contextLength = 512
        p10.train.lr = 6e-4; p10.train.warmupSteps = 58; p10.train.maxSteps = 584
        p10.train.batchSize = 64; p10.train.gradAccumSteps = 1
        p10.train.checkpointEvery = 200; p10.train.evalEvery = 100

        var p50 = ForgeConfig()
        p50.model.name = "gpt-50m"
        p50.model.nLayers = 10; p50.model.dModel = 640; p50.model.nHeads = 10
        p50.model.nKvHeads = 10; p50.model.dFf = 1728; p50.model.vocabSize = 4096
        p50.model.contextLength = 1024
        p50.train.lr = 5e-4; p50.train.warmupSteps = 117; p50.train.maxSteps = 1170
        p50.train.batchSize = 8; p50.train.gradAccumSteps = 8
        p50.train.checkpointEvery = 200; p50.train.evalEvery = 100

        var muon = p50
        muon.model.name = "gpt-50m-muon"
        muon.train.optimizer = "muon"; muon.train.schedule = "wsd"

        return [Preset(name: "gpt-10m (1 epoch)", config: p10),
                Preset(name: "gpt-50m", config: p50),
                Preset(name: "gpt-50m Muon+WSD", config: muon)]
    }()
}