import XCTest
@testable import MarkBase

class MoE26BStandardTest: XCTestCase {
    
    func testMoE26BStandardForward() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  MoE 26B-Standard Forward Pass Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-standard"
        
        guard FileManager.default.fileExists(atPath: modelPath) else {
            print("⚠ Model not found at \(modelPath)")
            return
        }
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        print("Loading 26B-Standard...")
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
        
        print("✓ Model loaded")
        print("  Layers: \(model.numHiddenLayers)")
        print("  Hidden: \(model.hiddenSize)")
        print("  Vocab: \(model.vocabSize)")
        
        print("\nTesting forward pass...")
        let result = try model.forwardOptimized(tokenId: 2, position: 0)
        
        let nanCount = result.filter { $0.isNaN }.count
        print("Forward result: NaN=\(nanCount)/\(result.count)")
        
        if nanCount == 0 {
            print("✓✓✓ Zero NaN - MoE model success!")
        } else {
            print("✗ NaN detected in MoE forward")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Test completed")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        XCTAssertEqual(nanCount, 0, "MoE forward should not produce NaN")
    }
}