6. Batch Prediction

// Batch prediction dengan MLArrayBatchProvider
func batchClassify(images: [CVPixelBuffer]) throws -> [String] {
    // Buat array of feature providers
    let providers: [MLFeatureProvider] = try images.map { buffer in
        let input = MyClassifierInput(image: buffer)
        return input
    }

    // Buat batch provider
    let batchProvider = MLArrayBatchProvider(array: providers)

    // Predict semua sekaligus
    let options = MLPredictionOptions()
    options.usesCPUOnly = false  // izinkan GPU/ANE

    let batchOutput = try model.predictions(
        fromBatch: batchProvider,
        options: options
    )

    // Extract hasil
    var labels: [String] = []
    for i in 0..<batchOutput.count {
        let output = batchOutput.features(at: i)
        if let label = output.featureValue(for: "classLabel")?.stringValue {
            labels.append(label)
        }
    }
    return labels
}

// Adaptive batch size berdasarkan available memory
func optimalBatchSize(for model: MLModel) -> Int {
    // Perkiraan memory per sample (dari model description)
    let inputDesc = model.modelDescription.inputDescriptionsByName
    var bytesPerSample = 0
    for (_, desc) in inputDesc {
        if let imageConstraint = desc.imageConstraint {
            let w = imageConstraint.pixelsWide
            let h = imageConstraint.pixelsHigh
            bytesPerSample += w * h * 4  // RGBA Float32
        } else if let multiArrayConstraint = desc.multiArrayConstraint {
            let count = multiArrayConstraint.shape.reduce(1) { $0 * $1.intValue }
            bytesPerSample += count * 4  // Float32
        }
    }

    // Gunakan 10% dari available memory sebagai batch budget
    let availableMemory = ProcessInfo.processInfo.physicalMemory
    let budget = Int(Double(availableMemory) * 0.10)
    let rawBatchSize = budget / max(bytesPerSample, 1)

    // Clamp ke rentang yang wajar
    return min(max(rawBatchSize, 1), 64)
}

// Streaming batch processor
actor BatchProcessor {
    private let model: MLModel
    private let batchSize: Int

    init(model: MLModel) {
        self.model = model
        self.batchSize = optimalBatchSize(for: model)
    }

    func processAll(inputs: [CVPixelBuffer]) async throws -> [ClassificationResult] {
        var results: [ClassificationResult] = []

        for chunk in inputs.chunks(ofSize: batchSize) {
            let providers = chunk.map { MyModelInput(image: $0) as MLFeatureProvider }
            let batch = MLArrayBatchProvider(array: providers)
            let batchOutput = try model.predictions(fromBatch: batch)

            for i in 0..<batchOutput.count {
                let output = batchOutput.features(at: i)
                results.append(ClassificationResult(from: output))
            }
        }

        return results
    }
}

struct ClassificationResult {
    let label: String
    let confidence: Float

    init(from output: MLFeatureProvider) {
        label = output.featureValue(for: "classLabel")?.stringValue ?? "unknown"
        confidence = output.featureValue(for: "classLabelProbs")?
            .dictionaryValue?[label as NSObject] as? Float ?? 0
    }
}

extension Array {
    func chunks(ofSize size: Int) -> [[Element]] {
        stride(from: 0, to: count, by: size).map { Array(self[$0..<Swift.min($0 + size, count)]) }
    }
}