import XCTest
@testable import MarkBase

class Model31BForwardTest: XCTestCase {
    
    func testModel31BForwardNaN() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  31B Forward Pass NaN Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-31b-it-4bit"
        
        guard FileManager.default.fileExists(atPath: modelPath) else {
            print("⚠ Model not found")
            return
        }
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        print("Loading 31B...")
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
        print("✓ Model loaded (Layers: \(model.numHiddenLayers), Hidden: \(model.hiddenSize))")
        
        // Test multiple tokenIds (like 26B-A4B)
        print("\nTesting tokenIds 0-10...")
        var nanPattern: [Int: Int] = [:]
        for testTokenId in 0..<10 {
            let result = try model.forwardOptimized(tokenId: testTokenId, position: 0)
            let nanCount = result.filter { $0.isNaN }.count
            if nanCount > 0 {
                nanPattern[testTokenId] = nanCount
            }
        }
        
        if nanPattern.isEmpty {
            print("  ✓ No NaN for tokenIds 0-10")
            print("  31B may handle scales differently (no NaN despite wrong scales)")
        } else {
            print("  ⚠ NaN detected:")
            for (tokenId, count) in nanPattern.sorted(by: { $0.key < $1.key }) {
                print("    tokenId=\(tokenId): NaN=\(count)")
            }
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
    }
}