// Native Core ML benchmark for the hand-written CiteNet builds.
// usage: bench <build dir>
import CoreML
import Foundation

let dir = URL(fileURLWithPath: CommandLine.arguments[1])
let seq = 64
let raw = try Data(contentsOf: dir.appendingPathComponent("bench_codes.f16"))
let nCites = raw.count / (seq * 2)

func load(_ b: Int, _ units: MLComputeUnits) throws -> MLModel {
    let cfg = MLModelConfiguration()
    cfg.computeUnits = units
    return try MLModel(contentsOf: dir.appendingPathComponent("CiteNet_b\(b).mlmodelc"), configuration: cfg)
}

func batch(_ b: Int, offset: Int) throws -> MLDictionaryFeatureProvider {
    let arr = try MLMultiArray(shape: [NSNumber(value: b), 1, 1, NSNumber(value: seq)], dataType: .float16)
    raw.withUnsafeBytes { src in
        let dst = arr.dataPointer
        for i in 0..<b {
            let c = (offset + i) % nCites
            memcpy(dst + i * seq * 2, src.baseAddress! + c * seq * 2, seq * 2)
        }
    }
    return try MLDictionaryFeatureProvider(dictionary: ["codes": arr])
}

func now() -> UInt64 { DispatchTime.now().uptimeNanoseconds }

func pct(_ v: [Double], _ p: Double) -> Double { v.sorted()[min(v.count - 1, Int(Double(v.count) * p))] }

let unitsName: [(String, MLComputeUnits)] = [("cpu", .cpuOnly), ("ane", .cpuAndNeuralEngine), ("gpu", .cpuAndGPU)]
let only = CommandLine.arguments.count > 2 ? CommandLine.arguments[2] : "all"
for b in [1, 32, 512, 2048, 8192] where only == "all" {
    for (uname, units) in unitsName {
        if b == 1 && uname == "gpu" { continue }          // b1/b32 plans stay on CPU for GPU units
        if b >= 2048 && uname == "cpu" { continue }
        let m = try load(b, units)
        let inputs = try (0..<8).map { try batch(b, offset: $0 * b) }
        for i in 0..<50 { _ = try m.prediction(from: inputs[i % 8]) }
        let iters = b == 1 ? 5000 : (b == 32 ? 3000 : max(40, 300_000 / b))
        var lat = [Double]()
        let t0 = now()
        for i in 0..<iters {
            let s = now()
            _ = try m.prediction(from: inputs[i % 8])
            lat.append(Double(now() - s) / 1e3)
        }
        let total = Double(now() - t0) / 1e9
        print(String(format: "b%-4d %@  p50 %7.1f us  p99 %7.1f us  %8.0f cites/s", b, uname,
                     pct(lat, 0.5), pct(lat, 0.99), Double(iters * b) / total))
    }
}

// ---- concurrent: one thread per compute unit, each with its own model, for a fixed time
func concurrent(_ lanes: [(String, Int, MLComputeUnits)], seconds: Double) throws {
    let models = try lanes.map { try load($0.1, $0.2) }
    let inputs = try lanes.map { l in try (0..<4).map { try batch(l.1, offset: $0 * l.1) } }
    for (i, m) in models.enumerated() { for j in 0..<20 { _ = try m.prediction(from: inputs[i][j % 4]) } }
    var counts = [Int](repeating: 0, count: lanes.count)
    let lock = NSLock()
    let start = now(), deadline = start + UInt64(seconds * 1e9)
    DispatchQueue.concurrentPerform(iterations: lanes.count) { i in
        var n = 0, j = 0
        while now() < deadline {
            _ = try! models[i].prediction(from: inputs[i][j % 4]); j += 1; n += lanes[i].1
        }
        lock.lock(); counts[i] = n; lock.unlock()
    }
    let total = Double(now() - start) / 1e9
    let parts = zip(lanes, counts).map { String(format: "%@ b%d %.0fk/s", $0.0.0, $0.0.1, Double($0.1) / total / 1e3) }
    print(String(format: "concurrent %@ -> %.0f cites/s", parts.joined(separator: " + "), Double(counts.reduce(0, +)) / total))
}
try concurrent([("ane", 512, .cpuAndNeuralEngine)], seconds: 3)
try concurrent([("gpu", 8192, .cpuAndGPU)], seconds: 3)
try concurrent([("ane", 512, .cpuAndNeuralEngine), ("gpu", 8192, .cpuAndGPU)], seconds: 3)
try concurrent([("ane", 2048, .cpuAndNeuralEngine), ("gpu", 8192, .cpuAndGPU)], seconds: 3)
try concurrent([("ane", 512, .cpuAndNeuralEngine), ("ane", 512, .cpuAndNeuralEngine), ("gpu", 8192, .cpuAndGPU), ("gpu", 8192, .cpuAndGPU)], seconds: 3)