import XCTest
@testable import MarkBase

final class KernelFusionPerformanceTest: XCTestCase {
    
    func testFusedKernelPerformance() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Kernel Fusion Performance Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine(autoCompile: true)
        let device = engine.device
        
        // Test parameters
        let hiddenSize = 2560
        let groupSize = 64
        
        // Create test buffers
        let inputBuffer = device.makeBuffer(length: hiddenSize * 4)!
        let outputBuffer = device.makeBuffer(length: hiddenSize * 4)!
        let scalesBuffer = device.makeBuffer(length: hiddenSize / groupSize * 4)!
        let biasesBuffer = device.makeBuffer(length: hiddenSize / groupSize * 4)!
        
        // Fill with test data
        let inputPtr = inputBuffer.contents().assumingMemoryBound(to: Float.self)
        for i in 0..<hiddenSize {
            inputPtr[i] = Float.random(in: -1...1)
        }
        
        let scalesPtr = scalesBuffer.contents().assumingMemoryBound(to: Float.self)
        let biasesPtr = biasesBuffer.contents().assumingMemoryBound(to: Float.self)
        for i in 0..<(hiddenSize / groupSize) {
            scalesPtr[i] = 0.1
            biasesPtr[i] = 0.0
        }
        
        // Create quantized weights (simplified)
        let weightBuffer = device.makeBuffer(length: hiddenSize * 4)!
        let weightPtr = weightBuffer.contents().assumingMemoryBound(to: UInt32.self)
        for i in 0..<hiddenSize {
            weightPtr[i] = UInt32.random(in: 0...UInt32.max)
        }
        
        // Warm up
        print("Warm up kernels...")
        _ = try engine.pipeline(named: "dequantize_row")
        _ = try engine.pipeline(named: "eltwise_scale")
        print("  ✓ Kernels loaded\n")
        
        // Test 1: Separate kernels (baseline)
        print("Test 1: Separate kernels (dequantize + scale)")
        var separateTimes: [Double] = []
        for _ in 0..<10 {
            let cmdBuf = engine.commandQueue.makeCommandBuffer()!
            
            // Dequantize
            let pso1 = try engine.pipeline(named: "dequantize_row")
            let enc1 = cmdBuf.makeComputeCommandEncoder()!
            enc1.setComputePipelineState(pso1)
            enc1.setBuffer(weightBuffer, offset: 0, index: 0)
            enc1.setBuffer(scalesBuffer, offset: 0, index: 1)
            enc1.setBuffer(biasesBuffer, offset: 0, index: 2)
            enc1.setBuffer(outputBuffer, offset: 0, index: 3)
            var n = UInt32(hiddenSize)
            enc1.setBytes(&n, length: 4, index: 4)
            var row = Int32(0)
            enc1.setBytes(&row, length: 4, index: 5)
            var gs = UInt32(groupSize)
            enc1.setBytes(&gs, length: 4, index: 6)
            enc1.dispatchThreads(MTLSize(width: hiddenSize, height: 1, depth: 1),
                                threadsPerThreadgroup: MTLSize(width: 256, height: 1, depth: 1))
            enc1.endEncoding()
            
            // Scale
            let pso2 = try engine.pipeline(named: "eltwise_scale")
            let enc2 = cmdBuf.makeComputeCommandEncoder()!
            enc2.setComputePipelineState(pso2)
            enc2.setBuffer(outputBuffer, offset: 0, index: 0)
            var scale = Float(1.5)
            enc2.setBytes(&scale, length: 4, index: 1)
            enc2.setBytes(&n, length: 4, index: 2)
            enc2.dispatchThreads(MTLSize(width: hiddenSize, height: 1, depth: 1),
                                threadsPerThreadgroup: MTLSize(width: 256, height: 1, depth: 1))
            enc2.endEncoding()
            
            let start = Date()
            cmdBuf.commit()
            cmdBuf.waitUntilCompleted()
            let elapsed = Date().timeIntervalSince(start) * 1000
            separateTimes.append(elapsed)
        }
        let separateAvg = separateTimes.reduce(0, +) / Double(separateTimes.count)
        print("  Average time: \(separateAvg) ms")
        
        // Test 2: Fused kernel
        print("\nTest 2: Fused kernel (dequantize + scale)")
        var fusedTimes: [Double] = []
        for _ in 0..<10 {
            let cmdBuf = engine.commandQueue.makeCommandBuffer()!
            
            // Try to use fused kernel
            do {
                let pso = try engine.pipeline(named: "fused_dequantize_scale")
                let enc = cmdBuf.makeComputeCommandEncoder()!
                enc.setComputePipelineState(pso)
                enc.setBuffer(weightBuffer, offset: 0, index: 0)
                enc.setBuffer(scalesBuffer, offset: 0, index: 1)
                enc.setBuffer(biasesBuffer, offset: 0, index: 2)
                enc.setBuffer(outputBuffer, offset: 0, index: 3)
                var n = UInt32(hiddenSize)
                enc.setBytes(&n, length: 4, index: 4)
                var row = Int32(0)
                enc.setBytes(&row, length: 4, index: 5)
                var gs = UInt32(groupSize)
                enc.setBytes(&gs, length: 4, index: 6)
                var scale = Float(1.5)
                enc.setBytes(&scale, length: 4, index: 7)
                enc.dispatchThreads(MTLSize(width: hiddenSize, height: 1, depth: 1),
                                   threadsPerThreadgroup: MTLSize(width: 256, height: 1, depth: 1))
                enc.endEncoding()
                
                let start = Date()
                cmdBuf.commit()
                cmdBuf.waitUntilCompleted()
                let elapsed = Date().timeIntervalSince(start) * 1000
                fusedTimes.append(elapsed)
            } catch {
                print("  ⚠ Fused kernel not available: \(error)")
                return
            }
        }
        let fusedAvg = fusedTimes.reduce(0, +) / Double(fusedTimes.count)
        print("  Average time: \(fusedAvg) ms")
        
        // Comparison
        print("\n═══════════════════════════════════════════════════════════════════")
        let speedup = separateAvg / fusedAvg
        let improvement = (separateAvg - fusedAvg) / separateAvg * 100
        print("Comparison:")
        print("  Separate kernels: \(separateAvg) ms")
        print("  Fused kernel: \(fusedAvg) ms")
        print("  Speedup: \(speedup)x")
        print("  Improvement: \(improvement)%")
        
        if speedup < 1.0 {
            print("\n⚠ Fused kernel is SLOWER than separate kernels!")
            print("  Issue: Kernel fusion needs optimization")
        } else {
            print("\n✓ Fused kernel is faster than separate kernels")
        }
        print("═══════════════════════════════════════════════════════════════════\n")
    }
}