import XCTest
@testable import MarkBase

final class BatchGenerationTest: XCTestCase {
    
    func testBatchGenerationPerformance() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Batch Generation Performance 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: \(textModel.numHiddenLayers) layers, hidden=\(textModel.hiddenSize)\n")
        
        // Warm up
        print("Warm up...")
        _ = try textModel.forwardOptimized(tokenId: 2, position: 0)
        print("  ✓ Warm up complete\n")
        
        // Test 1: Sequential generation (baseline)
        print("Test 1: Sequential generation (10 tokens)")
        let seqStart = Date()
        var seqTokens: [Int] = [2]
        for _ in 0..<10 {
            let logits = try textModel.forwardOptimized(
                tokenId: seqTokens.last!,
                position: seqTokens.count - 1
            )
            var maxIdx = 0
            var maxLogit = logits[0]
            for i in 1..<logits.count {
                if logits[i] > maxLogit {
                    maxLogit = logits[i]
                    maxIdx = i
                }
            }
            seqTokens.append(maxIdx)
        }
        let seqTime = Date().timeIntervalSince(seqStart) * 1000
        print("  ✓ Generated 10 tokens in \(seqTime) ms")
        print("  Average: \(seqTime / 10) ms/token")
        print("  Tokens: \(seqTokens)")
        
        // Test 2: Batch generation
        print("\nTest 2: Batch generation (10 tokens in batches of 8)")
        let batchStart = Date()
        let batchTokens = try textModel.generateBatch(startToken: 2, numTokens: 10)
        let batchTime = Date().timeIntervalSince(batchStart) * 1000
        print("  ✓ Generated 10 tokens in \(batchTime) ms")
        print("  Average: \(batchTime / 10) ms/token")
        print("  Tokens: \(batchTokens)")
        
        // Comparison
        print("\n═══════════════════════════════════════════════════════════════════")
        let speedup = seqTime / batchTime
        let improvement = (seqTime - batchTime) / seqTime * 100
        print("Comparison:")
        print("  Sequential: \(seqTime) ms (\(seqTime/10) ms/token)")
        print("  Batch: \(batchTime) ms (\(batchTime/10) ms/token)")
        print("  Speedup: \(speedup)x")
        print("  Improvement: \(improvement)%")
        
        if speedup > 1.0 {
            print("\n✓ Batch generation is faster!")
        } else {
            print("\n⚠ Batch generation needs optimization")
        }
        print("═══════════════════════════════════════════════════════════════════\n")
    }
    
    func testSingleVsBatchComparison() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Single vs Batch Token Generation")
        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)
        
        // Warm up
        print("Warm up...")
        _ = try textModel.forwardOptimized(tokenId: 2, position: 0)
        _ = try textModel.forwardBatch(tokenIds: [2], positions: [0])
        print("  ✓ Warm up complete\n")
        
        // Test single token
        print("Test 1: Single token generation")
        var singleTimes: [Double] = []
        for _ in 0..<5 {
            let start = Date()
            let logits = try textModel.forwardOptimized(tokenId: 2, position: 0)
            let elapsed = Date().timeIntervalSince(start) * 1000
            singleTimes.append(elapsed)
            XCTAssertFalse(logits.contains { $0.isNaN }, "Single logits should not have NaN")
        }
        let singleAvg = singleTimes.reduce(0, +) / Double(singleTimes.count)
        print("  Average: \(singleAvg) ms")
        
        // Test batch of 1
        print("\nTest 2: Batch of 1 token")
        var batch1Times: [Double] = []
        for _ in 0..<5 {
            let start = Date()
            let logits = try textModel.forwardBatch(tokenIds: [2], positions: [0])
            let elapsed = Date().timeIntervalSince(start) * 1000
            batch1Times.append(elapsed)
            XCTAssertFalse(logits[0].contains { $0.isNaN }, "Batch logits should not have NaN")
        }
        let batch1Avg = batch1Times.reduce(0, +) / Double(batch1Times.count)
        print("  Average: \(batch1Avg) ms")
        
        // Test batch of 4
        print("\nTest 3: Batch of 4 tokens")
        var batch4Times: [Double] = []
        for _ in 0..<5 {
            let start = Date()
            let logits = try textModel.forwardBatch(tokenIds: [2, 2, 2, 2], positions: [0, 1, 2, 3])
            let elapsed = Date().timeIntervalSince(start) * 1000
            batch4Times.append(elapsed)
            for l in logits {
                XCTAssertFalse(l.contains { $0.isNaN }, "Batch logits should not have NaN")
            }
        }
        let batch4Avg = batch4Times.reduce(0, +) / Double(batch4Times.count)
        print("  Average: \(batch4Avg) ms")
        print("  Per-token: \(batch4Avg / 4) ms")
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("Results:")
        print("  Single: \(singleAvg) ms")
        print("  Batch(1): \(batch1Avg) ms")
        print("  Batch(4): \(batch4Avg) ms (\(batch4Avg/4) ms/token)")
        
        let batch4Speedup = (singleAvg * 4) / batch4Avg
        print("  Batch(4) speedup vs 4x single: \(batch4Speedup)x")
        print("═══════════════════════════════════════════════════════════════════\n")
    }
}