import Foundation
import MarkBase

let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-31b-it-4bit"

guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else {
    print("✗ 31B model not found at \(modelDir)")
    exit(1)
}

print("Loading engine...")
let engine = try E4BEngine(autoCompile: true)
print("✓ Engine created")

print("\nLoading 31B model (~18 GB, may take 2-3 minutes)...")
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)

print("✓ Model loaded!")
print("  Layers: \(model.numHiddenLayers)")
print("  Hidden size: \(model.hiddenSize)")
print("  Vocab size: \(model.vocabSize)")
print("  Layer types: \(model.layerTypesIsFull.count) total, \(model.layerTypesIsFull.filter { $0 }.count) full, \(model.layerTypesIsFull.filter { !$0 }.count) sliding")

print("\nLoading tokenizer...")
let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
print("✓ Tokenizer loaded")

// Test forward pass
print("\n=== Forward pass test ===")
let tokens = tokenizer.encode(text: "Hello")
print("Tokenized 'Hello': \(tokens)")

let logits = try model.forward(tokenId: tokens[0], position: 0)
print("Forward pass complete: \(logits.count) logits")
let maxLogit = logits.max() ?? -999
print("Max logit: \(maxLogit)")

let sorted = logits.enumerated().sorted { $0.element > $1.element }
let top5 = sorted.prefix(5)
print("Top 5 predictions:")
for (idx, val) in top5 {
    let tokenStr = tokenizer.decode(tokens: [idx])
    print("  Token \(idx) ('\(tokenStr)'): \(val)")
}

// Generation test
print("\n=== Generation test ===")
let genConfig = GenerationConfig(maxTokens: 5, temperature: 0.7)
let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
let response = try generator.generateComplete(prompt: "Hello", config: genConfig)
print("Generated: '\(response)'")

print("\n✅ 31B test complete!")