// Latency of each GLiNER2.5-Decide bucket package with valid, full-length inputs.
// Usage: benchbuckets <units,...> <iterations> <pkg.mlmodelc>...
import CoreML
import Foundation

let args = CommandLine.arguments
let unitsByName: [String: MLComputeUnits] = [
    "all": .all, "cpuAndGPU": .cpuAndGPU, "cpuOnly": .cpuOnly, "cpuAndNeuralEngine": .cpuAndNeuralEngine,
]
let unitNames = args[1].split(separator: ",").map(String.init)
let iters = Int(args[2])!
let packages = args.dropFirst(3).map { URL(fileURLWithPath: $0) }

func filled(_ shape: [NSNumber], _ type: MLMultiArrayDataType, _ value: (Int) -> Double) throws -> MLMultiArray {
    let a = try MLMultiArray(shape: shape, dataType: type)
    for i in 0..<a.count { a[i] = NSNumber(value: value(i)) }
    return a
}

// Every position is a real (unmasked) token: the worst case for a bucket. Padding does not
// change the compute, so this is also the cost of any request routed to this bucket.
func inputs(for desc: MLModelDescription) throws -> MLFeatureProvider {
    let ids = desc.inputDescriptionsByName["input_ids"]!.multiArrayConstraint!.shape
    let grid = desc.inputDescriptionsByName["marker_indices"]!.multiArrayConstraint!.shape
    let L = ids[1].intValue, K = grid[2].intValue
    let used = [5, 3, 0, 0]  // two questions with 5 and 3 labels, like the example request
    return try MLDictionaryFeatureProvider(dictionary: [
        "input_ids": try filled(ids, .int32) { _ in Double(Int.random(in: 1000..<100_000)) },
        "attention_mask": try filled(ids, .int32) { _ in 1 },
        "marker_indices": try filled(grid, .int32) { i in
            let h = i / K, k = i % K
            return k < used[h] ? Double(min(L - 1, 2 + h * 12 + k * 2)) : 0
        },
        "marker_mask": try filled(grid, .float32) { i in i % K < used[i / K] ? 1 : 0 },
    ])
}

print("bucket   units        load     p50       p90")
for url in packages {
    for label in unitNames {
        let config = MLModelConfiguration()
        config.computeUnits = unitsByName[label]!
        let t0 = Date()
        let model = try MLModel(contentsOf: url, configuration: config)
        let load = Date().timeIntervalSince(t0)
        let input = try inputs(for: model.modelDescription)
        let L = model.modelDescription.inputDescriptionsByName["input_ids"]!.multiArrayConstraint!.shape[1]
        for _ in 0..<5 { _ = try model.prediction(from: input) }
        var t: [Double] = []
        for _ in 0..<iters {
            let s = Date()
            _ = try model.prediction(from: input)
            t.append(Date().timeIntervalSince(s) * 1000)
        }
        t.sort()
        print(String(format: "L%-6@  %-10@  %5.1fs  %6.1f ms  %6.1f ms",
                     L.stringValue as NSString, label as NSString, load, t[t.count / 2], t[Int(Double(t.count) * 0.9)]))
    }
}