Core AI Framework

import CoreAI
import SwiftUI

// MARK: - Core AI inference wrapper
actor DirectionPredictor {
    private let inferenceFunction: InferenceFunction

    // hiddenDim must match the model's fixed second dimension
    let hiddenDim: Int = 32

    init(modelURL: URL) throws {
        let model = try AIModel(contentsOf: modelURL)
        self.inferenceFunction = try model.loadMainFunction()
    }

    /// Predict the best direction index (0-3) from a sequence of feature vectors.
    func predict(featureSequence: [[Float]]) throws -> Int {
        let seqLen = featureSequence.count
        // Create a 2-D NDArray [seqLen, hiddenDim] backed by float32
        var inputArray = NDArray(
            shape: [seqLen, hiddenDim],
            dataType: .float32
        )
        // Fill NDArray row-by-row via a MutableView
        inputArray.withMutableView { view in
            for (row, features) in featureSequence.enumerated() {
                for (col, value) in features.prefix(hiddenDim).enumerated() {
                    view[row, col] = value
                }
            }
        }
        // Run inference — dispatches to CPU/GPU/Neural Engine automatically
        let outputs = try inferenceFunction.run(["features": inputArray])
        guard let logitsArray = outputs["logits"] else {
            throw PredictorError.missingOutput
        }
        // Read logits and return the index of the highest value
        let logits: [Float] = logitsArray.withView { view in
            (0..<4).map { view[0, $0] as Float }
        }
        return logits.indices.max(by: { logits[$0] < logits[$1] }) ?? 0
    }

    enum PredictorError: Error {
        case missingOutput
    }
}

// MARK: - SwiftUI demo view
struct CoreAIDemoView: View {
    @State private var predictedDirection: String = "—"
    @State private var isRunning = false
    private let directions = ["Up", "Down", "Left", "Right"]

    // In a real app, bundle the .aimodel and resolve with Bundle.main.url
    private let modelURL = Bundle.main.url(
        forResource: "SnakeTransformer", withExtension: "aimodel"
    )!

    var body: some View {
        VStack(spacing: 24) {
            Text("Core AI Demo")
                .font(.title2.bold())
            Text("Predicted direction")
                .foregroundStyle(.secondary)
            Text(predictedDirection)
                .font(.largeTitle.bold())
                .contentTransition(.numericText())
            Button(isRunning ? "Running…" : "Run Inference") {
                Task { await runInference() }
            }
            .buttonStyle(.borderedProminent)
            .disabled(isRunning)
        }
        .padding()
    }

    @MainActor
    func runInference() async {
        isRunning = true
        defer { isRunning = false }
        do {
            let predictor = try DirectionPredictor(modelURL: modelURL)
            // Simulate 3 game-state feature vectors of length 32
            let fakeSequence = (0..<3).map { _ in
                (0..<32).map { _ in Float.random(in: -1...1) }
            }
            let index = try await predictor.predict(featureSequence: fakeSequence)
            predictedDirection = directions[index]
        } catch {
            predictedDirection = "Error: \(error.localizedDescription)"
        }
    }
}

#Preview {
    CoreAIDemoView()
}