11. Stateful Models (iOS 18)

// Stateful model — state dikelola oleh Core ML runtime
// Berguna untuk: chatbot, speech recognition, real-time sensor fusion

if #available(iOS 18, *) {
    // Load model stateful
    let config = MLModelConfiguration()
    let model = try MLModel(contentsOf: compiledModelURL, configuration: config)

    // Buat MLState — mewakili satu "sesi" dengan state terpisah
    // Berbeda state = berbeda conversation thread
    let state1 = try model.makeState()
    let state2 = try model.makeState()  // state terpisah, independen

    // Prediksi dengan state — state diperbarui otomatis
    let input1 = NextTokenInput(token: 1234)
    let output1 = try model.prediction(from: input1, using: state1)
    // state1 sekarang berisi KV-cache dari token 1234

    let input2 = NextTokenInput(token: 5678)
    let output2 = try model.prediction(from: input2, using: state1)
    // state1 berisi KV-cache dari token 1234 DAN 5678

    // state2 masih kosong (independen dari state1)
    let output3 = try model.prediction(from: NextTokenInput(token: 9999), using: state2)

    // Reset state (mulai conversation baru)
    // Cukup buat state baru: let freshState = try model.makeState()
}

// Placeholder
class NextTokenInput: MLFeatureProvider {
    init(token: Int) {}
    var featureNames: Set<String> { [] }
    func featureValue(for featureName: String) -> MLFeatureValue? { nil }
}

// LLM autoregressive text generation dengan stateful model (iOS 18)
@available(iOS 18, *)
actor StatefulLLMInference {
    private let model: MLModel
    private var activeStates: [String: MLState] = [:]  // sessionID → state

    init(modelURL: URL) throws {
        let config = MLModelConfiguration()
        config.computeUnits = .cpuAndNeuralEngine
        self.model = try MLModel(contentsOf: modelURL, configuration: config)
    }

    func startSession() throws -> String {
        let sessionID = UUID().uuidString
        activeStates[sessionID] = try model.makeState()
        return sessionID
    }

    func generateNextToken(sessionID: String, currentToken: Int) throws -> Int {
        guard let state = activeStates[sessionID] else {
            throw GenerationError.sessionNotFound
        }

        let input = LLMTokenInput(tokenID: currentToken)
        let output = try model.prediction(from: input, using: state)

        // State sudah di-update dengan KV-cache dari currentToken
        guard let logits = output.featureValue(for: "logits")?.multiArrayValue else {
            throw GenerationError.invalidOutput
        }

        return argmax(logits)
    }

    func endSession(sessionID: String) {
        activeStates.removeValue(forKey: sessionID)
        // State di-release secara otomatis
    }

    private func argmax(_ array: MLMultiArray) -> Int {
        let ptr = array.dataPointer.assumingMemoryBound(to: Float.self)
        var maxIdx = 0
        var maxVal = ptr[0]
        for i in 1..<array.count {
            if ptr[i] > maxVal {
                maxVal = ptr[i]
                maxIdx = i
            }
        }
        return maxIdx
    }

    enum GenerationError: Error { case sessionNotFound, invalidOutput }
}

class LLMTokenInput: MLFeatureProvider {
    let tokenID: Int
    var featureNames: Set<String> { ["input_ids"] }
    init(tokenID: Int) { self.tokenID = tokenID }
    func featureValue(for featureName: String) -> MLFeatureValue? {
        guard featureName == "input_ids" else { return nil }
        let array = try? MLMultiArray(shape: [1, 1], dataType: .int32)
        array?.dataPointer.assumingMemoryBound(to: Int32.self)[0] = Int32(tokenID)
        return array.map { MLFeatureValue(multiArray: $0) }
    }
}