import XCTest
@testable import MarkBase

class StressTest: XCTestCase {
    
    func testConcurrentInference() throws {
        print("\n═════════════════════════════════════════════════════════════════════")
        print("  Stress Test 1: Concurrent Inference (5 sequences)")
        print("═══════════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 256)
        
        print("✓ Model loaded")
        
        let concurrentCount = 5
        let start = Date()
        var totalTokens = 0
        var nanCount = 0
        
        for i in 0..<concurrentCount {
            for pos in 0..<20 {
                let tokenId = 2 + i + pos
                let logits = try model.forwardOptimized(tokenId: tokenId, position: pos)
                nanCount += logits.filter { $0.isNaN }.count
                totalTokens += 1
            }
        }
        
        let elapsed = Date().timeIntervalSince(start) * 1000
        
        print("\nResults:")
        print("  Sequences: \(concurrentCount)")
        print("  Tokens per seq: 20")
        print("  Total tokens: \(totalTokens)")
        print("  Time: \(String(format: "%.1f", elapsed))ms")
        print("  Throughput: \(String(format: "%.1f", Double(totalTokens) / elapsed * 1000)) tok/s")
        print("  NaN count: \(nanCount)")
        
        if nanCount == 0 {
            print("✅ PASS - Concurrent inference stable")
        } else {
            print("⚠ FAIL - NaN detected")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════════")
    }
    
    func testMemoryStress() throws {
        print("\n═══════════════════════════════════════════════════════════════════════")
        print("  Stress Test 2: Memory Pressure (256 context)")
        print("═════════════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
        
        print("✓ Model loaded with maxContext=512")
        
        let start = Date()
        var tokenCount = 0
        var nanCount = 0
        
        for pos in 0..<256 {
            let tokenId = 2 + (pos % 100)
            let logits = try model.forwardOptimized(tokenId: tokenId, position: pos)
            nanCount += logits.filter { $0.isNaN }.count
            tokenCount += 1
        }
        
        let elapsed = Date().timeIntervalSince(start) * 1000
        
        print("\nResults:")
        print("  Context length: 256")
        print("  Tokens processed: \(tokenCount)")
        print("  Time: \(String(format: "%.1f", elapsed))ms")
        print("  Speed: \(String(format: "%.1f", Double(tokenCount) / elapsed * 1000)) tok/s")
        print("  NaN count: \(nanCount)")
        
        if nanCount == 0 {
            print("✅ PASS - Memory stress test passed")
        } else {
            print("⚠ FAIL - NaN in long context")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════════")
    }
    
    func testContinuousGeneration() throws {
        print("\n═══════════════════════════════════════════════════════════════════════")
        print("  Stress Test 3: Continuous Generation (100 tokens)")
        print("═════════════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
        
        print("✓ Model loaded")
        
        let start = Date()
        var currentToken = 2
        var nanCount = 0
        
        for i in 0..<100 {
            let logits = try model.forwardOptimized(tokenId: currentToken, position: i)
            nanCount += logits.filter { $0.isNaN }.count
            
            var maxIdx = 0
            var maxVal = logits[0]
            for j in 1..<logits.count {
                if logits[j] > maxVal {
                    maxVal = logits[j]
                    maxIdx = j
                }
            }
            currentToken = maxIdx
        }
        
        let elapsed = Date().timeIntervalSince(start) * 1000
        
        print("\nResults:")
        print("  Tokens generated: 100")
        print("  Time: \(String(format: "%.1f", elapsed))ms")
        print("  Speed: \(String(format: "%.1f", 100.0 / elapsed * 1000)) tok/s")
        print("  NaN count: \(nanCount)")
        
        if nanCount == 0 {
            print("✅ PASS - Continuous generation stable")
        } else {
            print("⚠ FAIL - NaN during generation")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════════")
    }
    
    func testBatchProcessing() throws {
        print("\n═══════════════════════════════════════════════════════════════════════")
        print("  Stress Test 4: Batch Processing (10 batches)")
        print("═════════════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
        
        print("✓ Model loaded")
        
        let start = Date()
        var totalTokens = 0
        var nanCount = 0
        
        for batch in 0..<10 {
            for pos in 0..<10 {
                let tokenId = 2 + batch + pos
                let logits = try model.forwardOptimized(tokenId: tokenId, position: pos)
                nanCount += logits.filter { $0.isNaN }.count
                totalTokens += 1
            }
        }
        
        let elapsed = Date().timeIntervalSince(start) * 1000
        
        print("\nResults:")
        print("  Batches: 10")
        print("  Tokens per batch: 10")
        print("  Total tokens: \(totalTokens)")
        print("  Time: \(String(format: "%.1f", elapsed))ms")
        print("  Throughput: \(String(format: "%.1f", Double(totalTokens) / elapsed * 1000)) tok/s")
        print("  NaN count: \(nanCount)")
        
        if nanCount == 0 {
            print("✅ PASS - Batch processing stable")
        } else {
            print("⚠ FAIL - NaN in batches")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════════")
    }
    
    func testLongRunningStability() throws {
        print("\n═══════════════════════════════════════════════════════════════════════")
        print("  Stress Test 5: Long Running Stability (30s)")
        print("═════════════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 256)
        
        print("✓ Model loaded")
        
        let start = Date()
        var totalTokens = 0
        var nanCount = 0
        var errorCount = 0
        let duration = 30.0
        
        while Date().timeIntervalSince(start) < duration {
            for pos in 0..<20 {
                let tokenId = 2 + pos
                do {
                    let logits = try model.forwardOptimized(tokenId: tokenId, position: pos)
                    nanCount += logits.filter { $0.isNaN }.count
                    totalTokens += 1
                } catch {
                    errorCount += 1
                }
            }
        }
        
        let elapsed = Date().timeIntervalSince(start) * 1000
        
        print("\nResults:")
        print("  Duration: \(duration)s")
        print("  Total tokens: \(totalTokens)")
        print("  NaN count: \(nanCount)")
        print("  Errors: \(errorCount)")
        print("  Avg speed: \(String(format: "%.1f", Double(totalTokens) / elapsed)) tok/s")
        
        if errorCount == 0 && nanCount == 0 {
            print("✅ PASS - Long running stability OK")
        } else {
            print("⚠ FAIL - Stability issues")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════════")
    }
}