// Simple test to verify multimodal inference works
import MarkBase

let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"

print("Loading model...")
let engine = try E4BEngine(autoCompile: true)
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
let sampler = Sampler()

print("✓ Model loaded")
print("  Layers: \(model.numHiddenLayers)")
print("  Vocab: \(model.vocabSize)")

// Test text generation
print("\n=== Test 1: Text Generation ===")
let prompt = "Hello"
let tokens = tokenizer.encode(text: prompt)
print("Prompt: '\(prompt)'")
print("Tokens: \(tokens)")

// Reset KV cache
for cache in model.kvCaches {
    cache.reset()
}

// Generate
print("Generating...")
var generated = tokens
for i in 0..<10 {
    let pos = generated.count - 1
    let logits = try model.forward(tokenId: generated.last!, position: pos)
    let nextToken = sampler.sample(logits: logits, temperature: 0.7, topK: 50, topP: 0.95, filterUnusedTokens: true)
    generated.append(nextToken)
    let tokenStr = tokenizer.decode(tokens: [nextToken])
    print("  \(i+1): \(nextToken) '\(tokenStr)'")
}

print("\nOutput: '\(tokenizer.decode(tokens: generated))'")

// Test multimodal
print("\n=== Test 2: Multimodal ===")
let mm = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
print("Multimodal model loaded")

if let tower = mm.visionTowerFull {
    print("Vision tower: \(tower.config.numHiddenLayers) layers")
}