// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Multi-run overlay: val-loss (and optional train-EMA) curves for 2–8 runs,
// X axis in steps, TOKENS (the honest axis when batch sizes differ) or
// wall-clock; legend annotated with the auto-computed hyperparameter diff
// (only the fields that differ across the selection); summary table + CSV
// export.
import Charts
import SwiftUI

struct CompareView: View {
    @Environment(AppModel.self) private var app
    @State private var selected: Set<UUID> = []
    @State private var xMode: XMode = .steps
    @State private var showTrainEMA = false
    @State private var series: [UUID: CompareSeries] = [:]

    enum XMode: String, CaseIterable {
        case steps = "steps"
        case tokens = "tokens"
        case time = "temps"
    }

    struct CompareSeries {
        var val: [(x: Double, y: Double)] = []
        var ema: [(x: Double, y: Double)] = []
        var bestVal: Double?
        var finalPpl: Double?
        var meanTps: Double?
        var totalTime: Double?
    }

    private let palette: [Color] = [.blue, .orange, .green, .purple, .red, .teal,
                                    .pink, .brown]

    private var candidates: [Run] {
        app.store.runs.filter { FileManager.default.fileExists(atPath: $0.logCSVPath) }
            .sorted { $0.createdAt > $1.createdAt }
    }
    private var selectedRuns: [Run] { candidates.filter { selected.contains($0.id) } }

    var body: some View {
        HSplitView {
            List(candidates, selection: $selected) { run in
                RunRow(run: run).tag(run.id)
            }
            .frame(minWidth: 230, maxWidth: 300)

            VStack(alignment: .leading, spacing: 12) {
                HStack {
                    Text("Comparaison").font(.title2.weight(.semibold))
                    Spacer()
                    Toggle("train EMA", isOn: $showTrainEMA).toggleStyle(.checkbox)
                    Picker("", selection: $xMode) {
                        ForEach(XMode.allCases, id: \.self) { Text($0.rawValue) }
                    }
                    .pickerStyle(.segmented)
                    .frame(width: 220)
                    Button("Exporter CSV") { exportCSV() }
                        .disabled(selectedRuns.isEmpty)
                }
                if selectedRuns.count < 2 {
                    ContentUnavailableView(
                        "Sélectionnez 2 à 8 runs",
                        systemImage: "chart.xyaxis.line",
                        description: Text("⌘-clic dans la liste pour comparer leurs courbes de validation."))
                } else {
                    diffLegend
                    chart
                    summaryTable
                }
                Spacer(minLength: 0)
            }
            .padding()
        }
        .task(id: "\(selected.hashValue)|\(xMode.rawValue)") { await loadSeries() }
        .navigationTitle("Comparer")
    }

    private var chart: some View {
        Chart {
            ForEach(Array(selectedRuns.enumerated()), id: \.element.id) { idx, run in
                let s = series[run.id]
                ForEach(s?.val ?? [], id: \.x) { p in
                    LineMark(x: .value("x", p.x), y: .value("val", p.y),
                             series: .value("run", run.name + " val"))
                        .foregroundStyle(palette[idx % palette.count])
                }
                if showTrainEMA {
                    ForEach(s?.ema ?? [], id: \.x) { p in
                        LineMark(x: .value("x", p.x), y: .value("ema", p.y),
                                 series: .value("run", run.name + " ema"))
                            .foregroundStyle(palette[idx % palette.count].opacity(0.4))
                    }
                }
            }
        }
        .chartXAxisLabel(xMode == .tokens ? "tokens vus"
                         : xMode == .time ? "secondes" : "steps")
        .frame(minHeight: 300)
    }

    // Only the config fields that differ across the selection.
    private var diffLegend: some View {
        let diffs = configDiffs()
        return VStack(alignment: .leading, spacing: 4) {
            ForEach(Array(selectedRuns.enumerated()), id: \.element.id) { idx, run in
                HStack(spacing: 8) {
                    Circle().fill(palette[idx % palette.count]).frame(width: 9, height: 9)
                    Text(run.name).font(.callout.weight(.medium))
                    Text(diffs[run.id] ?? "").font(.caption)
                        .foregroundStyle(.secondary)
                }
            }
        }
    }

    private func configDiffs() -> [UUID: String] {
        guard selectedRuns.count >= 2 else { return [:] }
        // Flatten each config to key: value strings via JSON.
        var flat: [UUID: [String: String]] = [:]
        for run in selectedRuns {
            guard let data = try? run.config.exportJSON(),
                  let obj = try? JSONSerialization.jsonObject(with: data)
                      as? [String: [String: Any]] else { continue }
            var kv: [String: String] = [:]
            for (section, fields) in obj {
                for (k, v) in fields { kv["\(section).\(k)"] = "\(v)" }
            }
            flat[run.id] = kv
        }
        let allKeys = Set(flat.values.flatMap(\.keys))
        let differing = allKeys.filter { key in
            Set(flat.values.map { $0[key] ?? "" }).count > 1
        }.sorted()
        var out: [UUID: String] = [:]
        for run in selectedRuns {
            out[run.id] = differing
                .filter { $0 != "model.name" }
                .map { "\($0.split(separator: ".").last!)=\(flat[run.id]?[$0] ?? "?")" }
                .joined(separator: "  ")
        }
        return out
    }

    private var summaryTable: some View {
        Table(selectedRuns) {
            TableColumn("Run") { Text($0.name) }
            TableColumn("Params") {
                Text($0.config.model.paramCount.formatted(.number.notation(.compactName)))
            }
            TableColumn("Best val") { run in
                Text(series[run.id]?.bestVal.map { String(format: "%.4f", $0) } ?? "—")
            }
            TableColumn("PPL finale") { run in
                Text(series[run.id]?.finalPpl.map { String(format: "%.1f", $0) } ?? "—")
            }
            TableColumn("tok/s moyen") { run in
                Text(series[run.id]?.meanTps.map { String(format: "%.0f", $0) } ?? "—")
            }
            TableColumn("Durée") { run in
                Text(series[run.id]?.totalTime.map {
                    Duration.seconds($0).formatted(.time(pattern: .hourMinuteSecond))
                } ?? "—")
            }
        }
        .frame(height: CGFloat(40 + selectedRuns.count * 28))
    }

    private func loadSeries() async {
        var out: [UUID: CompareSeries] = [:]
        for run in selectedRuns {
            let store = MetricsStore()
            await store.ingestCSV(at: URL(fileURLWithPath: run.logCSVPath))
            let snap = await store.snapshot(maxPoints: 800, smoothing: 0.6)
            let tps = Double(run.config.tokensPerStep)
            func xform(_ p: Downsampler.XY, elapsedLookup: Bool = false) -> (Double, Double) {
                switch xMode {
                case .steps: return (p.x, p.y)
                case .tokens: return ((p.x + 1) * tps, p.y)
                case .time: return (p.x, p.y) // refined below with elapsed
                }
            }
            var s = CompareSeries()
            if xMode == .time {
                // elapsed_s per step from the raw points.
                let pts = await store.points
                s.val = pts.compactMap { p in
                    guard let v = p.valLoss, let e = p.elapsedS else { return nil }
                    return (e, v)
                }
                s.ema = zip(pts, Smoothing.ema(pts.map(\.trainLoss), smoothing: 0.6))
                    .compactMap { p, e in p.elapsedS.map { ($0, e) } }
            } else {
                s.val = snap.val.map { xform($0) }
                s.ema = snap.trainEMA.map { xform($0) }
            }
            s.bestVal = snap.bestVal?.loss
            s.finalPpl = snap.lastPoint.map { exp($0.trainLoss) }
            let allPts = await store.points
            if !allPts.isEmpty {
                s.meanTps = allPts.map(\.tokensPerSec).reduce(0, +) / Double(allPts.count)
                s.totalTime = allPts.last?.elapsedS
            }
            out[run.id] = s
        }
        series = out
    }

    private func exportCSV() {
        var csv = "run,params,best_val,final_ppl,mean_tok_s,total_s\n"
        for run in selectedRuns {
            let s = series[run.id]
            csv += "\(run.name),\(run.config.model.paramCount),"
            csv += "\(s?.bestVal.map { String($0) } ?? ""),"
            csv += "\(s?.finalPpl.map { String($0) } ?? ""),"
            csv += "\(s?.meanTps.map { String($0) } ?? ""),"
            csv += "\(s?.totalTime.map { String($0) } ?? "")\n"
        }
        let panel = NSSavePanel()
        panel.nameFieldStringValue = "compare.csv"
        if panel.runModal() == .OK, let url = panel.url {
            try? csv.write(to: url, atomically: true, encoding: .utf8)
        }
    }
}