// Step 0 diagnostic: where does CoreML place each op of the GLiNER2.5-Decide package?
// Usage: plan <path.mlpackage|.mlmodelc> <out.json>
import CoreML
import Foundation

func deviceName(_ d: MLComputeDevice) -> String {
    switch d {
    case .cpu: return "CPU"
    case .gpu: return "GPU"
    case .neuralEngine: return "ANE"
    @unknown default: return "?"
    }
}

struct OpRecord: Codable {
    let index: Int
    let op: String
    let outputs: [String]
    let preferred: String
    let supported: [String]
    let cost: Double
}

func walk(_ block: MLModelStructure.Program.Block, plan: MLComputePlan, into records: inout [OpRecord]) {
    for op in block.operations {
        let usage = plan.deviceUsage(for: op)
        let cost = plan.estimatedCost(of: op)?.weight ?? 0
        records.append(OpRecord(
            index: records.count,
            op: op.operatorName,
            outputs: op.outputs.map { $0.name },
            preferred: usage.map { deviceName($0.preferred) } ?? "none",
            supported: usage.map { $0.supported.map(deviceName) } ?? [],
            cost: cost))
        for inner in op.blocks { walk(inner, plan: plan, into: &records) }
    }
}

let args = CommandLine.arguments
let modelURL = URL(fileURLWithPath: args[1])
let outURL = URL(fileURLWithPath: args[2])

let compiledURL: URL
if modelURL.pathExtension == "mlmodelc" {
    compiledURL = modelURL
} else {
    let t0 = Date()
    let tmp = try await MLModel.compileModel(at: modelURL)
    compiledURL = modelURL.deletingPathExtension().appendingPathExtension("mlmodelc")
    try? FileManager.default.removeItem(at: compiledURL)
    try FileManager.default.moveItem(at: tmp, to: compiledURL)
    print("compiled in \(String(format: "%.1f", Date().timeIntervalSince(t0)))s -> \(compiledURL.lastPathComponent)")
}

var result: [String: [OpRecord]] = [:]
for (label, units) in [("cpuAndNeuralEngine", MLComputeUnits.cpuAndNeuralEngine), ("all", MLComputeUnits.all)] {
    let config = MLModelConfiguration()
    config.computeUnits = units
    let plan = try await MLComputePlan.load(contentsOf: compiledURL, configuration: config)
    guard case let .program(program) = plan.modelStructure, let main = program.functions["main"] else {
        fatalError("not an ML program")
    }
    var records: [OpRecord] = []
    walk(main.block, plan: plan, into: &records)
    result[label] = records
    print("\(label): \(records.count) ops")
}
let enc = JSONEncoder()
enc.outputFormatting = [.prettyPrinted, .sortedKeys]
try enc.encode(result).write(to: outURL)
print("wrote \(outURL.path)")