import XCTest
@testable import MarkBase

final class BatchEmbeddingOptimizationTest: XCTestCase {
    
    func testBatchEmbeddingPerformance() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Batch Embedding Optimization Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        // Test with E4B-MarkBase (smaller, faster to load)
        let modelPath = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        
print("Loading E4B-MarkBase...")
let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
print("✓ Model loaded: \(model.layers.count) layers\n")

// Create batch context (using model's helper method)
let maxBatchSize = 8
let context = model.createBatchContext(maxBatchSize: maxBatchSize)
        
        print("Testing batch embedding performance:")
        print("  Batch sizes: 1, 2, 4, 8\n")
        
        for batchSize in [1, 2, 4, 8] {
            print("Testing batch(\(batchSize))...")
            
            // Generate random token IDs
            let tokenIds = (0..<batchSize).map { _ in Int.random(in: 0..<model.vocabSize) }
            let positions = Array(0..<batchSize)
            
            // Test 5 forward passes
            var times: [Double] = []
            for i in 0..<5 {
                let start = Date()
                let logits = try model.forwardBatchTrue(
                    tokenIds: tokenIds,
                    positions: positions,
                    context: context
                )
                let elapsed = Date().timeIntervalSince(start) * 1000
                times.append(elapsed)
                
                // Verify output
                XCTAssertEqual(logits.count, batchSize, "Batch size mismatch")
                XCTAssertEqual(logits[0].count, model.vocabSize, "Vocab size mismatch")
                
                let nanCount = logits.flatMap { $0 }.filter { $0.isNaN }.count
                XCTAssertEqual(nanCount, 0, "NaN detected in batch output")
            }
            
            let avgTime = times.reduce(0, +) / Double(times.count)
            let perToken = avgTime / Double(batchSize)
            
            print("  ✓ Batch(\(batchSize)): avg \(String(format: "%.1f", avgTime))ms total, \(String(format: "%.1f", perToken))ms/token")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Batch Embedding Optimization COMPLETE ✓✓✓")
        print("  Single GPU dispatch for entire batch (eliminated sequential waits)")
        print("═══════════════════════════════════════════════════════════════════\n")
    }
    
    func test31BBatchPerformance() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  31B Batch Performance Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine(autoCompile: true)
        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-31b-it-4bit"
        
        print("Loading 31B model...")
        let startLoad = Date()
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
        let loadTime = Date().timeIntervalSince(startLoad) * 1000
print("✓ Model loaded in \(String(format: "%.1f", loadTime))ms")
print("  Layers: \(model.layers.count)")
print("  Hidden: \(model.hiddenSize)\n")

// Create batch context (using model's helper method)
let maxBatchSize = 4
let context = model.createBatchContext(maxBatchSize: maxBatchSize)
        
        print("Testing 31B batch performance:")
        
        // Single token baseline
        print("\nSingle token baseline...")
        var singleTimes: [Double] = []
        for i in 0..<3 {
            let start = Date()
            let logits = try model.forwardOptimized(tokenId: 2, position: i)
            let elapsed = Date().timeIntervalSince(start) * 1000
            singleTimes.append(elapsed)
            
            let nanCount = logits.filter { $0.isNaN }.count
            XCTAssertEqual(nanCount, 0, "NaN in single forward")
        }
        let singleAvg = singleTimes.reduce(0, +) / Double(singleTimes.count)
        print("  ✓ Single: \(String(format: "%.1f", singleAvg))ms/token")
        
        // Batch(4) test
        print("\nBatch(4) test...")
        let tokenIds = [2, 2, 2, 2]
        let positions = [0, 1, 2, 3]
        
        var batchTimes: [Double] = []
        for i in 0..<3 {
            let start = Date()
            let logits = try model.forwardBatchTrue(
                tokenIds: tokenIds,
                positions: positions,
                context: context
            )
            let elapsed = Date().timeIntervalSince(start) * 1000
            batchTimes.append(elapsed)
            
            XCTAssertEqual(logits.count, 4, "Batch size mismatch")
            let nanCount = logits.flatMap { $0 }.filter { $0.isNaN }.count
            XCTAssertEqual(nanCount, 0, "NaN in batch forward")
        }
        
        let batchAvg = batchTimes.reduce(0, +) / Double(batchTimes.count)
        let batchPerToken = batchAvg / 4.0
        let speedup = singleAvg / batchPerToken
        
        print("  ✓ Batch(4): \(String(format: "%.1f", batchAvg))ms total, \(String(format: "%.1f", batchPerToken))ms/token")
        print("  ✓ Speedup: \(String(format: "%.1f", speedup))x")
        
        // Target check
        if batchPerToken < 50.0 {
            print("\n✓✓✓ TARGET MET! 31B batch <50ms/token")
        } else {
            print("\n⚠ Target not met: \(String(format: "%.1f", batchPerToken))ms/token (target <50ms)")
            print("  Further optimization needed")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  31B Batch Test Complete")
        print("═══════════════════════════════════════════════════════════════════\n")
    }
}