import Foundation
import MarkBase

print("\n=== 测试 NaN Bug ===\n")

let modelPath = "./models/E4B-MarkBase"
print("加载模型...")
let engine = try MarkBaseEngine(autoCompile: true)
let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
let tokenizer = try TokenizerFactory.load(modelDir: modelPath)
print("✓ 模型加载完成\n")

// 测试 prompt encoding
let prompt = "Hello"
print("Prompt: '\(prompt)'")
let promptTokens = tokenizer.encode(text: prompt)
print("Tokens: \(promptTokens)")
print("Token details:")
for (i, tokenId) in promptTokens.enumerated() {
    let raw = tokenizer.rawToken(for: tokenId) ?? "nil"
    print("  [\(i)] ID=\(tokenId) → '\(raw)'")
}

print("\n=== 测试 Forward Pass ===\n")

// 测试第一个 token 的 forward pass
var lastLogits: [Float] = []
for (position, tokenId) in promptTokens.enumerated() {
    print("Forward pass: tokenId=\(tokenId), position=\(position)")
    lastLogits = try model.forward(tokenId: tokenId, position: position)
    
    // 检查 NaN
    let hasNaN = lastLogits.contains { $0.isNaN }
    let nanCount = lastLogits.filter { $0.isNaN }.count
    let maxVal = lastLogits.max() ?? 0
    let minVal = lastLogits.min() ?? 0
    
    print("  Logits: count=\(lastLogits.count), hasNaN=\(hasNaN), nanCount=\(nanCount)")
    print("  Stats: max=\(maxVal), min=\(minVal)")
    if nanCount < 10 {
        print("  Sample (first 20): \(lastLogits.prefix(20))")
    }
    print()
}

print("=== 测试完成 ===")