import XCTest
@testable import MarkBase

final class OptimizationPrototypeTest: XCTestCase {
    
    func testBatchedCommandsDemo() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Metal Command Batching Optimization Demo")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let device = MTLCreateSystemDefaultDevice()!
        let queue = device.makeCommandQueue()!
        
        let size = 2560  // E4B hidden size
        let buffer1 = device.makeBuffer(length: size * 4)!
        let buffer2 = device.makeBuffer(length: size * 4)!
        let buffer3 = device.makeBuffer(length: size * 4)!
        
        // ── 測試 1: 多個同步操作（慢）──────────────────────────────────
        print("Test 1: Multiple Synchronous Operations (Current Approach)")
        let start1 = Date()
        
        for i in 0..<5 {
            let cmdBuf = queue.makeCommandBuffer()!
            let blit = cmdBuf.makeBlitCommandEncoder()!
            blit.copy(from: buffer1, sourceOffset: 0,
                     to: buffer2, destinationOffset: i * size * 4 / 5,
                     size: size * 4 / 5)
            blit.endEncoding()
            cmdBuf.commit()
            cmdBuf.waitUntilCompleted()  // ← 每次都等待！
        }
        
        let time1 = Date().timeIntervalSince(start1) * 1000
        print("  Time: \(time1) ms (5 synchronous blit operations)")
        print("  Issue: Each operation waits for GPU completion")
        
        // ── 測試 2: Batched操作（快）──────────────────────────────────
        print("\nTest 2: Batched Operations (Optimized Approach)")
        let start2 = Date()
        
        let cmdBuf = queue.makeCommandBuffer()!
        let blit = cmdBuf.makeBlitCommandEncoder()!
        
        // 所有操作加入同一個 command buffer
        for i in 0..<5 {
            blit.copy(from: buffer1, sourceOffset: 0,
                     to: buffer3, destinationOffset: i * size * 4 / 5,
                     size: size * 4 / 5)
        }
        
        blit.endEncoding()
        cmdBuf.commit()
        cmdBuf.waitUntilCompleted()  // ← 只等待一次！
        
        let time2 = Date().timeIntervalSince(start2) * 1000
        print("  Time: \(time2) ms (5 batched blit operations)")
        print("  Benefit: All operations execute in one GPU dispatch")
        
        // ── 比較結果 ───────────────────────────────────────────────────
        print("\nComparison:")
        let speedup = time1 / time2
        print("  Speedup: \(speedup)x faster")
        print("  Savings: \(time1 - time2) ms")
        print("  WaitUntilCompleted calls: Test1=5 vs Test2=1")
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("✓ Optimization demo completed")
        print("  Key insight: Batching commands reduces GPU-CPU sync overhead")
        print("  Expected improvement: 10x+ faster TEXT generation")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        XCTAssertGreaterThan(speedup, 2.0, "Batched should be significantly faster")
    }
    
    func testForwardPassBenchmark() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Forward Pass Benchmark (waitUntilCompleted Analysis)")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        print("Current Model.swift structure:")
        print("  Total waitUntilCompleted calls: 11")
        print("  Location breakdown:")
        
        let lines = [
            "Line 283: Embedding dequantize",
            "Line 682: Scale operation",
            "Line 706: Per-layer embedding",
            "Line 730: Per-layer projection",
            "Line 1135: Per-layer norm (inside loop)",
            "Line 1157: Per-layer copy",
            "Line 1181: Hidden state copy",
            "Line 1304: Layer operations",
            "Line 1322: LM head",
            "Line 1348: Readback",
            "Line 1367: Final operation"
        ]
        
        for line in lines {
            print("    \(line)")
        }
        
        print("\nOptimization target:")
        print("  Reduce from 11 → 1 waitUntilCompleted")
        print("  Expected speedup: 10x")
        print("  Estimated token generation time:")
        print("    E4B: 11.3秒 → ~1.1秒")
        print("    12B: 5.8秒 → ~0.6秒")
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("✓ Analysis complete - optimization plan ready")
        print("═══════════════════════════════════════════════════════════════════\n")
    }
}