import Foundation
import MarkBase

// Test harness for 26B-standard model
let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-standard"

guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else {
    print("Model not found")
    exit(1)
}

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

print("\nLoading 26B Standard model (~15 GB, may take 2-3 minutes)...")
let start = Date()
let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)
let loadTime = Date().timeIntervalSince(start)

print("✓ Model loaded in \(String(format: "%.1f", loadTime))s")
print("  Layers: \(model.numHiddenLayers)")
print("  Hidden: \(model.hiddenSize)")
print("  Vocab: \(model.vocabSize)")

// Forward pass
print("\n=== Forward pass test ===")
let logits = try model.forward(tokenId: 2, position: 0)
print("✓ Forward pass complete: \(logits.count) logits")
print("  Max logit: \(logits.max() ?? -999)")
let sorted = logits.enumerated().sorted { $0.element > $1.element }
let top5 = sorted.prefix(5)
print("  Top 5: \(top5.map { "\($0.offset):\(String(format: "%.2f", $0.element))" }.joined(separator: ", "))")

print("\n✅ 26B Standard test complete!")