import XCTest
@testable import MarkBase

final class MoERouterOnlyTest: XCTestCase {
    
    func testRouterProjectionOnly() throws {
        print("\n═══════════════════════════════════════")
        print("      Router Projection Only Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        
        print("Step 1: Load model...")
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 32)
        print("  ✓ Model loaded")
        
        print("\nStep 2: Get router from layer 0...")
        let layer0 = model.layers[0]
        
        guard let router = layer0.routerProj else {
            XCTFail("Router not present")
            return
        }
        print("  ✓ Router retrieved")
        print("    Router bits: \(router.bits)")
        print("    Router inDim: \(router.inDim)")
        print("    Router outDim: \(router.outDim)")
        
        print("\nStep 3: Create minimal buffers...")
        let hs = model.hiddenSize
        
        // Create input buffer (normalized layer output)
        let inputBuffer = engine.device.makeBuffer(length: hs * MemoryLayout<Float>.size)!
        var inputData = [Float](repeating: 0.1, count: hs)  // Small values to avoid overflow
        memcpy(inputBuffer.contents(), &inputData, inputData.count * MemoryLayout<Float>.size)
        
        // Create output buffer for router logits
        let routerOutDim = router.outDim  // 128 experts
        let routerOutputBuffer = engine.device.makeBuffer(length: routerOutDim * MemoryLayout<Float>.size)!
        
        print("  ✓ Buffers created")
        print("    Input: \(hs) floats")
        print("    Output: \(routerOutDim) floats (expert scores)")
        
        print("\nStep 4: Test router projection alone...")
        let start = Date()
        
        do {
            print("  Calling quantizedMatmul...")
            fflush(stdout)
            
            let cmdBuf = engine.commandQueue.makeCommandBuffer()!
            try layer0.quantizedMatmul(
                engine: engine,
                cmdBuf: cmdBuf,
                input: inputBuffer,
                weights: router,
                output: routerOutputBuffer
            )
            
            print("  Command buffer created, committing...")
            fflush(stdout)
            
            cmdBuf.commit()
            print("  Waiting for completion...")
            fflush(stdout)
            
            cmdBuf.waitUntilCompleted()
            
            let elapsed = Date().timeIntervalSince(start)
            print("  ✓ Router projection completed in \(String(format: "%.3f", elapsed))s")
            print("    Command buffer status: \(cmdBuf.status.rawValue)")
            
            if elapsed > 5.0 {
                print("  ⚠️ Router took >5s - potential issue")
            }
            
            if cmdBuf.status != .completed {
                let errorMsg = cmdBuf.error?.localizedDescription ?? "unknown"
                print("  ❌ Router command failed: \(errorMsg)")
                XCTFail("Router projection failed")
                return
            }
            
            // Read router output
            let routerOutput = engine.readFloats(from: routerOutputBuffer, count: routerOutDim)
            print("\n  Router logits (first 10): \(routerOutput[0..<10])")
            print("  Router logits max: \(routerOutput.max() ?? 0)")
            print("  Router logits min: \(routerOutput.min() ?? 0)")
            
            let hasNaN = routerOutput.contains { $0.isNaN }
            if hasNaN {
                print("  ❌ NaN detected in router output!")
                XCTFail("Router output has NaN")
            } else {
                print("  ✓ Router output valid (no NaN)")
            }
            
        } catch {
            print("  ❌ Router projection failed: \(error)")
            throw error
        }
        
        print("\n═══════════════════════════════════════")
        print("✓ Router projection works!")
        print("═══════════════════════════════════════\n")
    }
}