import XCTest
@testable import MarkBase

final class Layer4NaNTest: XCTestCase {
    func testLayer4Expert127NaN() throws {
        print("\n" + String(repeating: "=", count: 60))
        print("Layer 4 Expert 127 NaN Debug 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("\nModel loaded successfully")
        
        // Run one forward pass to trigger Layer 4 Expert 127
        let inputIds: [Int32] = [1, 3234, 357, 659, 198]
        print("\nRunning forward pass with input: \(inputIds)")
        
        // Run multiple forward passes like generation
        do {
            for (index, tokenId) in inputIds.enumerated() {
                print("\n--- Processing token \(index): \(tokenId) ---")
                let logits = try model.forward(tokenId: Int(tokenId), position: index)
                let logitsNaNCount = logits.filter { $0.isNaN }.count
                print("Logits NaN count: \(logitsNaNCount)/\(logits.count)")
                
                if logitsNaNCount > 0 {
                    print("ERROR: NaN detected in logits!")
                    break
                }
            }
            print("\nAll forward passes completed successfully")
        } catch {
            print("Forward pass failed: \(error)")
        }
        
        print("\n" + String(repeating: "=", count: 60))
    }
}