import XCTest
@testable import MarkBase

final class MoEHiddenStateNaNDebugTest: XCTestCase {
    
    func testHiddenStateNaNDebug() throws {
        print("\n═══════════════════════════════════════")
        print("      Hidden State NaN Debug (Layer 5)")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        
        print("Step 1: Load model...")
        fflush(stdout)
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 32)
        print("  ✓ Model loaded")
        fflush(stdout)
        
        // Get Layer 5 info
        print("\nStep 2: Layer 5 info...")
        fflush(stdout)
        let layer5 = model.layers[5]
        print("  Layer 5: isFull=\(layer5.config.isSliding ? false : true), headDim=\(layer5.config.headDim), nKvHeads=\(layer5.config.nKvHeads)")
        print("  Layer 5: useMoE=\(layer5.useMoE)")
        print("  Layer 5: routerScale=\(layer5.routerScale)")
        fflush(stdout)
        
        print("\nStep 3: Run forward with monitoring...")
        fflush(stdout)
        
        let h = model.temps.io
        let hs = model.hiddenSize
        
        // Run first token
        let _ = try model.forward(tokenId: 10979, position: 0, debug: true)
        
        // Check hidden state after each layer
        print("\n  Hidden state after each layer:")
        fflush(stdout)
        
        let hiddenVals = engine.readFloats(from: h, count: min(20, hs))
        let maxVal = hiddenVals.max() ?? 0
        let minVal = hiddenVals.min() ?? 0
        let hasNaN = hiddenVals.contains { $0.isNaN }
        
        print("  After all layers: max=\(maxVal), min=\(minVal), hasNaN=\(hasNaN)")
        print("  Hidden sample: \(hiddenVals[0..<min(10, hs)])")
        fflush(stdout)
        
        print("\n═══════════════════════════════════════")
        print("✓ NaN debug completed")
        print("═══════════════════════════════════════\n")
        fflush(stdout)
    }
}