import XCTest
@testable import MarkBase

final class GenerationNaNTest: XCTestCase {
    func test26BA4BGenerationNoNaN() throws {
        print("\n" + String(repeating: "=", count: 60))
        print("26B-A4B Generation NaN Test - Final Verification")
        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("\nModel loaded successfully")
        
        // Run generation like the original test
        let inputIds: [Int32] = [1, 3234, 357, 659, 198]
        print("\nRunning generation with input: \(inputIds)")
        
        do {
            // Process input tokens
            for (index, tokenId) in inputIds.enumerated() {
                let _ = try model.forward(tokenId: Int(tokenId), position: index)
            }
            
            print("\nInput tokens processed, generating new tokens...")
            
            // Generate new tokens
            var generatedTokens: [Int32] = []
            var lastToken = inputIds.last!
            
            for i in 0..<5 {
                let logits = try model.forward(tokenId: Int(lastToken), position: inputIds.count + i)
                let logitsNaNCount = logits.filter { $0.isNaN }.count
                
                print("  Generated token \(i): logits NaN count = \(logitsNaNCount)")
                
                if logitsNaNCount > 0 {
                    print("ERROR: NaN detected during generation!")
                    break
                }
                
                // Sample next token (simple argmax)
                let maxIdx = logits.enumerated().max(by: { $0.element < $1.element })!.offset
                generatedTokens.append(Int32(maxIdx))
                lastToken = Int32(maxIdx)
            }
            
            print("\n✓ Generation completed successfully!")
            print("Generated tokens: \(generatedTokens)")
            print("No NaN detected in any forward pass!")
            
        } catch {
            print("Generation failed: \(error)")
        }
        
        print("\n" + String(repeating: "=", count: 60))
    }
    
    func testGenerationPerformance() throws {
        print("\n" + String(repeating: "=", count: 60))
        print("26B-A4B Performance Benchmark")
        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)
        
        let inputIds: [Int32] = [1, 3234, 357, 659, 198]
        
        // Benchmark generation
        let start = Date()
        for (index, tokenId) in inputIds.enumerated() {
            _ = try model.forward(tokenId: Int(tokenId), position: index)
        }
        
        // 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 maxIdx = logits.enumerated().max(by: { $0.element < $1.element })!.offset
            lastToken = Int32(maxIdx)
        }
        
        let elapsed = Date().timeIntervalSince(start)
        let totalTokens = inputIds.count + 5
        let speed = Double(totalTokens) / elapsed
        
        print("\n✓ Generated \(totalTokens) tokens in \(String(format: "%.3f", elapsed))s")
        print("✓ Speed: \(String(format: "%.2f", speed)) tok/s")
        
        print("\n" + String(repeating: "=", count: 60))
    }
}