import XCTest
@testable import MarkBase

final class MoEExpertComputationTest: XCTestCase {
    
    func testSingleExpertFusedGateUp() throws {
        print("\n═══════════════════════════════════════")
        print("      Single Expert Fused Gate+Up 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 expert components from layer 0...")
        let layer0 = model.layers[0]
        
        guard let eGate = layer0.expertGate,
              let eUp = layer0.expertUp,
              let eDown = layer0.expertDown else {
            XCTFail("Expert components not present")
            return
        }
        
        print("  ✓ Expert components retrieved")
        print("    Gate: \(eGate.numExperts) experts, outDim=\(eGate.expertOutDim)")
        print("    Up: \(eUp.numExperts) experts, outDim=\(eUp.expertOutDim)")
        print("    Down: \(eDown.numExperts) experts, outDim=\(eDown.expertOutDim)")
        
        print("\nStep 3: Create minimal buffers...")
        let hs = model.hiddenSize
        let moeIntermediate = eGate.expertOutDim  // 704
        
        // Input buffer (normalized)
        let inputBuffer = engine.device.makeBuffer(length: hs * MemoryLayout<Float>.size)!
        var inputData = [Float](repeating: 0.1, count: hs)
        memcpy(inputBuffer.contents(), &inputData, inputData.count * MemoryLayout<Float>.size)
        
        // Output buffer for gate+up (2 * moeIntermediate)
        let gateUpOutputBuffer = engine.device.makeBuffer(length: 2 * moeIntermediate * MemoryLayout<Float>.size)!
        
        print("  ✓ Buffers created")
        print("    Input: \(hs) floats")
        print("    Gate+Up output: \(2 * moeIntermediate) floats")
        
        print("\nStep 4: Test single expert (expert 0) gate+up...")
        let start = Date()
        
        do {
            print("  Calling expertFusedGateUp for expert 0...")
            fflush(stdout)
            
            let cmdBuf = engine.commandQueue.makeCommandBuffer()!
            try layer0.expertFusedGateUp(
                engine: engine,
                cmdBuf: cmdBuf,
                input: inputBuffer,
                gate: eGate,
                up: eUp,
                expertIdx: 0,  // First expert only
                output: gateUpOutputBuffer
            )
            
            print("  Command buffer created, committing...")
            fflush(stdout)
            
            cmdBuf.commit()
            print("  Waiting for completion...")
            fflush(stdout)
            
            cmdBuf.waitUntilCompleted()
            
            let elapsed = Date().timeIntervalSince(start)
            print("  ✓ Expert gate+up completed in \(String(format: "%.3f", elapsed))s")
            print("    Command buffer status: \(cmdBuf.status.rawValue)")
            
            if elapsed > 10.0 {
                print("  ⚠️ Expert took >10s - potential issue")
            }
            
            if cmdBuf.status != .completed {
                let errorMsg = cmdBuf.error?.localizedDescription ?? "unknown"
                print("  ❌ Expert command failed: \(errorMsg)")
                XCTFail("Expert computation failed")
                return
            }
            
            // Read output
            let gateUpOutput = engine.readFloats(from: gateUpOutputBuffer, count: 10)
            print("\n  Gate+Up output (first 10): \(gateUpOutput)")
            
            let hasNaN = gateUpOutput.contains { $0.isNaN }
            if hasNaN {
                print("  ❌ NaN detected in expert output!")
                XCTFail("Expert output has NaN")
            } else {
                print("  ✓ Expert output valid (no NaN)")
            }
            
        } catch {
            print("  ❌ Expert computation failed: \(error)")
            throw error
        }
        
        print("\n═══════════════════════════════════════")
        print("✓ Single expert gate+up works!")
        print("═══════════════════════════════════════\n")
    }
    
    func testExpertLoopIteration() throws {
        print("\n═══════════════════════════════════════")
        print("      Expert Loop Iteration Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        
        print("Loading model...")
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 32)
        
        let layer0 = model.layers[0]
        guard let eGate = layer0.expertGate, let eUp = layer0.expertUp else {
            XCTFail("Experts not present")
            return
        }
        
        let hs = model.hiddenSize
        let moeIntermediate = eGate.expertOutDim
        
        let inputBuffer = engine.device.makeBuffer(length: hs * MemoryLayout<Float>.size)!
        var inputData = [Float](repeating: 0.1, count: hs)
        memcpy(inputBuffer.contents(), &inputData, inputData.count * MemoryLayout<Float>.size)
        
        let gateUpOutputBuffer = engine.device.makeBuffer(length: 2 * moeIntermediate * MemoryLayout<Float>.size)!
        
        print("\nTesting multiple expert iterations...")
        
        // Test iterating through 8 experts (top-k for 26B-A4B)
        let topKExperts = [0, 1, 2, 3, 4, 5, 6, 7]
        
        for (i, expertIdx) in topKExperts.enumerated() {
            print("  Expert \(i): Calling expertFusedGateUp for expert \(expertIdx)...")
            fflush(stdout)
            
            let start = Date()
            let cmdBuf = engine.commandQueue.makeCommandBuffer()!
            
            do {
                try layer0.expertFusedGateUp(
                    engine: engine,
                    cmdBuf: cmdBuf,
                    input: inputBuffer,
                    gate: eGate,
                    up: eUp,
                    expertIdx: expertIdx,
                    output: gateUpOutputBuffer
                )
                
                cmdBuf.commit()
                cmdBuf.waitUntilCompleted()
                
                let elapsed = Date().timeIntervalSince(start)
                print("    ✓ Expert \(expertIdx) completed in \(String(format: "%.3f", elapsed))s")
                
                if elapsed > 5.0 {
                    print("    ⚠️ Expert \(expertIdx) slow (>5s)")
                }
                
                if cmdBuf.status != .completed {
                    print("    ❌ Expert \(expertIdx) failed!")
                    XCTFail("Expert \(expertIdx) iteration failed")
                    return
                }
                
            } catch {
                print("    ❌ Expert \(expertIdx) error: \(error)")
                throw error
            }
        }
        
        print("\n═══════════════════════════════════════")
        print("✓ Expert loop iteration works (8 experts tested)")
        print("═══════════════════════════════════════\n")
    }
}