import XCTest
@testable import MarkBase

final class FinalMoETest: XCTestCase {
    func test26BA4BFullGeneration() throws {
        print("\n" + String(repeating: "=", count: 60))
        print("26B-A4B MoE Full Generation Test")
        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 successfully")
        
        // Test input
        let inputIds: [Int32] = [1, 3234, 357, 659, 198]
        
        // Forward pass for all 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 at position \(index)")
                return
            }
        }
        
        // Generate 10 new tokens
        print("\nGenerating 10 tokens...")
        var lastToken = inputIds.last!
        var generatedTokens: [Int32] = []
        
        for i in 0..<10 {
            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 at generation position \(i)")
                return
            }
            
            // Get next token (simple greedy)
            let maxIdx = logits.enumerated().max(by: { $0.element < $1.element })!.offset
            lastToken = Int32(maxIdx)
            generatedTokens.append(lastToken)
        }
        
        print("✓ Generated tokens: \(generatedTokens)")
        print("✓ No NaN in entire forward pass!")
        
        print("\n" + String(repeating: "=", count: 60))
        print("SUCCESS: 26B-A4B MoE fully operational!")
        print(String(repeating: "=", count: 60))
    }
}