import XCTest
@testable import MarkBase

final class CleanMoETest: XCTestCase {
    func testCleanMoEPerformance() throws {
        print("\n" + String(repeating: "=", count: 60))
        print("Clean MoE Performance Test (No Debug Output)")
        print(String(repeating: "=", count: 60))
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)
        
        print("\n✓ Model loaded")
        
        // Benchmark without debug output
        let inputIds: [Int32] = [1, 3234, 357, 659, 198]
        let iterations = 3
        
        var totalSpeed: Double = 0
        
        for iter in 0..<iterations {
            print("\nIteration \(iter + 1):")
            
            let start = Date()
            
            // Process input tokens
            for (index, tokenId) in inputIds.enumerated() {
                let logits = try model.forward(tokenId: Int(tokenId), position: index)
                let nanCount = logits.filter { $0.isNaN }.count
                if nanCount > 0 {
                    print("  ERROR: NaN detected in logits!")
                    return
                }
            }
            
            // Generate 5 new tokens
            var lastToken = inputIds.last!
            for i in 0..<5 {
                let logits = try model.forward(tokenId: Int(lastToken), position: inputIds.count + i)
                let nanCount = logits.filter { $0.isNaN }.count
                if nanCount > 0 {
                    print("  ERROR: NaN detected in generation!")
                    return
                }
                let maxIdx = logits.enumerated().max(by: { $0.element < $1.element })!.offset
                lastToken = Int32(maxIdx)
            }
            
            let elapsed = Date().timeIntervalSince(start)
            let tokens = inputIds.count + 5
            let speed = Double(tokens) / elapsed
            
            print("  Tokens: \(tokens), Time: \(String(format: "%.3f", elapsed))s")
            print("  Speed: \(String(format: "%.2f", speed)) tok/s")
            
            totalSpeed += speed
        }
        
        let avgSpeed = totalSpeed / Double(iterations)
        print("\n✓ Average speed: \(String(format: "%.2f", avgSpeed)) tok/s")
        print("✓ No NaN detected in all iterations!")
        
        print("\n" + String(repeating: "=", count: 60))
    }
}