import XCTest
@testable import MarkBase

final class OptimizationVerificationTest: XCTestCase {
    
    func testOptimizationEffect() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Optimization Verification Test")
        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 loaded: \(textModel.numHiddenLayers) layers\n")
        
        // Warm up shaders (run once to compile)
        print("Warm up Metal shaders...")
        _ = try textModel.forward(tokenId: 2, position: 0)
        _ = try textModel.forwardOptimized(tokenId: 2, position: 0)
        print("  ✓ Shaders compiled\n")
        
        // Test original forward (10 tokens)
        print("Test 1: Original forward (10 tokens)...")
        var tokens1: [Int] = [2]
        let start1 = Date()
        for _ in 0..<10 {
            let logits = try textModel.forward(tokenId: tokens1.last!, position: tokens1.count - 1)
            var maxIdx = 0
            var maxLogit = logits[0]
            for i in 1..<logits.count {
                if logits[i] > maxLogit { maxLogit = logits[i]; maxIdx = i }
            }
            tokens1.append(maxIdx)
        }
        let time1 = Date().timeIntervalSince(start1) * 1000
        print("  ✓ Time: \(time1) ms (avg \(time1/10) ms/token)")
        
        // Test optimized forward (10 tokens)
        print("\nTest 2: Optimized forward (10 tokens)...")
        var tokens2: [Int] = [2]
        let start2 = Date()
        for _ in 0..<10 {
            let logits = try textModel.forwardOptimized(tokenId: tokens2.last!, position: tokens2.count - 1)
            var maxIdx = 0
            var maxLogit = logits[0]
            for i in 1..<logits.count {
                if logits[i] > maxLogit { maxLogit = logits[i]; maxIdx = i }
            }
            tokens2.append(maxIdx)
        }
        let time2 = Date().timeIntervalSince(start2) * 1000
        print("  ✓ Time: \(time2) ms (avg \(time2/10) ms/token)")
        
        // Comparison
        print("\nResults:")
        let speedup = time1 / time2
        let improvement = (time1 - time2) / time1 * 100
        print("  Speedup: \(speedup)x")
        print("  Improvement: \(improvement)%")
        print("  Time saved: \(time1 - time2) ms")
        
        // Debug: check which path is being used
        print("\nDebug info:")
        print("  Dense layers should have: 0 waits per layer")
        print("  Original: \(textModel.numHiddenLayers) layers × 1 wait = \(textModel.numHiddenLayers) waits")
        print("  Optimized: 1 wait total (at end)")
        print("  Expected speedup: \(textModel.numHiddenLayers)x")
        print("  Actual speedup: \(speedup)x")
        
        if speedup < 2.0 {
            print("\n⚠ Warning: Optimization not working as expected!")
            print("  Possible issues:")
            print("    1. Layer.forwardOptimized not being called")
            print("    2. Internal helper functions creating their own command buffers")
            print("    3. MoE layers being used (require router read)")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("✓ Verification test complete")
        print("═══════════════════════════════════════════════════════════════════\n")
    }
}