Core AI Framework

import CoreAI
import Foundation

// MARK: - Model Player using Core AI Framework

struct GameFeatures {
    let wallDistances: [Float]   // 4 values
    let foodDelta: (x: Float, y: Float)
    let currentDirection: [Float] // 4 one-hot values
    let opponentDistance: Float
    let opponentDirection: [Float] // 4 one-hot values
    
    var featureVector: [Float] {
        wallDistances + [foodDelta.x, foodDelta.y] + currentDirection + [opponentDistance] + opponentDirection
    }
    static let featureDim = 15
}

enum SnakeDirection: Int { case up = 0, down, left, right }

final class SnakeAIPlayer {
    private let inferenceFunction: InferenceFunction

    init(modelURL: URL) throws {
        // Load the .aimodel file
        let model = try AIModel(contentsOf: modelURL)
        // Grab the single main inference function
        self.inferenceFunction = try model.loadInferenceFunction()
    }

    func chooseDirection(history: [GameFeatures]) throws -> SnakeDirection {
        let seqLen = history.count
        let featureDim = GameFeatures.featureDim

        // Create a 2D NDArray: [sequenceLength, featureDim] of Float32
        var input = NDArray(
            shape: [seqLen, featureDim],
            dataType: .float32
        )

        // Populate features row by row using a MutableView (non-escapable, memory-safe)
        input.withMutableView { view in
            for (row, features) in history.enumerated() {
                let vec = features.featureVector
                for (col, value) in vec.enumerated() {
                    view[row, col] = value
                }
            }
        }

        // Run inference — returns a dictionary of named output NDArrays
        let outputs = try inferenceFunction.run(inputs: ["features": input])

        // Read output logits NDArray [4] — one logit per direction
        guard let logitsArray = outputs["logits"] else {
            throw InferenceError.missingOutput
        }

        // Find argmax over the 4 direction logits
        let bestIndex = logitsArray.withView { view -> Int in
            var maxIdx = 0
            var maxVal: Float = view[0]
            for i in 1..<4 {
                let v: Float = view[i]
                if v > maxVal { maxVal = v; maxIdx = i }
            }
            return maxIdx
        }

        return SnakeDirection(rawValue: bestIndex) ?? .up
    }
}

enum InferenceError: Error {
    case missingOutput
}