// Run coreml/laya-fp16-512.mlpackage on the Neural Engine and print the option probabilities.
// The token ids come from one case of tests/vectors.json (the first one, or the id given as the
// second argument), so this file needs no tokenizer. In an app, build the ids with a port of
// examples/laya_sequence.py.
//
//   swiftc -O examples/RouteCoreML.swift -o /tmp/route && /tmp/route . pay_wrong_address
import CoreML
import Foundation

let length = 512
let pad: Int32 = 50283
let temperature = 1.7601518630981445  // rl_agent_config.json, "choice:3-5"

func multiArray(_ values: [Int32], count: Int, fill: Int32) throws -> MLMultiArray {
    let array = try MLMultiArray(shape: [1, NSNumber(value: count)], dataType: .int32)
    let pointer = array.dataPointer.bindMemory(to: Int32.self, capacity: count)
    for i in 0..<count { pointer[i] = i < values.count ? values[i] : fill }
    return array
}

let repo = URL(fileURLWithPath: CommandLine.arguments.count > 1 ? CommandLine.arguments[1] : ".")
let vectors = try JSONSerialization.jsonObject(
    with: Data(contentsOf: repo.appendingPathComponent("tests/vectors.json"))) as! [String: Any]
let cases = vectors["cases"] as! [[String: Any]]
let wanted = CommandLine.arguments.count > 2 ? CommandLine.arguments[2] : cases[0]["id"] as! String
guard let chosen = cases.first(where: { $0["id"] as? String == wanted }) else {
    fatalError("no case \(wanted) in tests/vectors.json")
}
let ids = (chosen["input_ids"] as! [Int]).map(Int32.init)
let markers = (chosen["markers"] as! [Int]).map(Int32.init)

// Compile once and keep the .mlmodelc (an app does this after download, then reuses it).
let compiled = try MLModel.compileModel(at: repo.appendingPathComponent("coreml/laya-fp16-512.mlpackage"))
let configuration = MLModelConfiguration()
configuration.computeUnits = .cpuAndNeuralEngine
let model = try MLModel(contentsOf: compiled, configuration: configuration)

let features = try MLDictionaryFeatureProvider(dictionary: [
    "input_ids": try multiArray(ids, count: length, fill: pad),
    "attention_mask": try multiArray([Int32](repeating: 1, count: ids.count), count: length, fill: 0),
    "marker_pos": try multiArray(markers, count: 5, fill: 0),
])
let logits = try model.prediction(from: features).featureValue(for: "logits")!.multiArrayValue!

let z = (0..<5).map { logits[$0].doubleValue / temperature }
let top = z.max()!
let e = z.map { exp($0 - top) }
let probabilities = e.map { $0 / e.reduce(0, +) }
let labels = ["billing", "technical", "account", "shipping", "other"]  // option order in the sequence
let expected = chosen["probabilities"] as! [String: Double]
print("\(wanted): \(ids.count) tokens")
for (label, p) in zip(labels, probabilities) {
    let name = label.padding(toLength: 10, withPad: " ", startingAt: 0)
    print(name + String(format: " %.4f  (laya %.4f)", p, expected[label]!))
}

var times: [Double] = []
for _ in 0..<30 {
    let start = DispatchTime.now().uptimeNanoseconds
    _ = try model.prediction(from: features)
    times.append(Double(DispatchTime.now().uptimeNanoseconds - start) / 1e6)
}
times.sort()
print(String(format: "Swift .cpuAndNeuralEngine: median %.1f ms, p90 %.1f ms", times[15], times[26]))

if #available(macOS 14.4, iOS 17.4, *) {
    let plan = try await MLComputePlan.load(contentsOf: compiled, configuration: configuration)
    if case let .program(program) = plan.modelStructure, let main = program.functions["main"] {
        var devices: [String: Int] = [:]
        for operation in main.block.operations {
            guard let usage = plan.deviceUsage(for: operation) else { continue }
            let name: String
            switch usage.preferred {
            case .cpu: name = "CPU"
            case .gpu: name = "GPU"
            case .neuralEngine: name = "NeuralEngine"
            @unknown default: name = "other"
            }
            devices[name, default: 0] += 1
        }
        print("MLComputePlan preferred devices:", devices)
    }
}
try? FileManager.default.removeItem(at: compiled)