import XCTest
@testable import MarkBase

final class PerformanceAnalysisTest: XCTestCase {
    
    func testMetalOperationCount() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Metal Operation Count Analysis")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let textModel = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 256)
        
        print("Model: \(textModel.numHiddenLayers) layers")
        print("\nPer-layer operations (estimated):")
        
        // Count operations per layer (dense path)
        let opsPerLayer = [
            "1. input_layernorm (rmsNorm)",
            "2. q_proj (quantizedMatmul)",
            "3. q_norm (groupedRmsNorm)",
            "4. RoPE Q (applyRoPEQ)",
            "5. k_proj (quantizedMatmul)",
            "6. k_norm (groupedRmsNorm)",
            "7. RoPE K (applyRoPEK)",
            "8. v_proj (quantizedMatmul or blit)",
            "9. v_norm (groupedRmsNorm if present)",
            "10. KV cache store",
            "11. Attention (sliding or full)",
            "12. o_proj (quantizedMatmul)",
            "13. Residual add (eltwiseAdd)",
            "14. post_attention_layernorm (rmsNorm)",
            "15. pre_feedforward_layernorm (rmsNorm)",
            "16. gate+up fused (fusedGateUp)",
            "17. down_proj (quantizedMatmul)",
            "18. Residual add (eltwiseAdd)",
            "19. post_feedforward_layernorm (rmsNorm)",
            "20. Per-layer gating (optional, 4-5 ops)"
        ]
        
        for op in opsPerLayer {
            print("  \(op)")
        }
        
        let numOps = opsPerLayer.count
        print("\nTotal ops per layer: ~\(numOps)")
        print("Total ops per forward: ~\(numOps * textModel.numHiddenLayers)")
        
        // Additional embedding/lm head ops
        print("\nEmbedding phase:")
        print("  1. dequantize embedding")
        print("  2. embedding scale")
        print("  3. dequantize per-layer embedding")
        print("  4. per-layer scale")
        print("  5-10. per-layer projection (matmul, scale, norm, add, scale)")
        
        print("\nLM head phase:")
        print("  11. final norm")
        print("  12. lm head (quantizedMatmul)")
        print("  13. logits scaling (if needed)")
        print("  14. logit softcapping")
        
        let embedOps = 10
        let lmOps = 4
        let totalOps = embedOps + numOps * textModel.numHiddenLayers + lmOps
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("Estimated total Metal operations per forward pass:")
        print("  Embedding: \(embedOps)")
        print("  Layers: \(numOps) × \(textModel.numHiddenLayers) = \(numOps * textModel.numHiddenLayers)")
        print("  LM head: \(lmOps)")
        print("  Total: ~\(totalOps)")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        print("Optimization analysis:")
        print("  Original: \(totalOps) operations in \(textModel.numHiddenLayers) command buffers")
        print("  Optimized: \(totalOps) operations in 1 command buffer")
        print("  Expected: reduce \(textModel.numHiddenLayers) → 1 waits")
        print("  But: Each Metal operation has kernel launch overhead (~0.1-0.5ms)")
        print("  Total overhead: \(totalOps) × 0.2ms = \(Double(totalOps) * 0.2)ms")
        print("  This explains why we only see 4x instead of 42x!")
        print("  The bottleneck is kernel dispatch overhead, not waitUntilCompleted")
    }
}