import XCTest
@testable import MarkBase

class InferenceSpeedTest: XCTestCase {
    
    func test26BStandardSpeed() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  26B-Standard Inference Speed Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-standard"
        
        guard FileManager.default.fileExists(atPath: modelPath) else {
            print("⚠ Model not found at \(modelPath)")
            return
        }
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        print("Loading 26B-Standard...")
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
        print("✓ Model loaded (Layers: \(model.numHiddenLayers), Hidden: \(model.hiddenSize))")
        
        // Warm up
        print("\nWarm up (1 token)...")
        let warmupStart = Date()
        _ = try model.forwardOptimized(tokenId: 2, position: 0)
        let warmupTime = Date().timeIntervalSince(warmupStart) * 1000
        print("  Warmup: \(String(format: "%.1f", warmupTime))ms")
        
        // Test single token generation speed
        print("\nTesting 10 tokens...")
        let testStart = Date()
        var currentToken = 2
        for i in 0..<10 {
            let result = try model.forwardOptimized(tokenId: currentToken, position: i)
            
            // Greedy selection (max logits)
            var maxIdx = 0
            var maxVal = result[0]
            for j in 1..<result.count {
                if result[j] > maxVal {
                    maxVal = result[j]
                    maxIdx = j
                }
            }
            currentToken = maxIdx
        }
        let testTime = Date().timeIntervalSince(testStart) * 1000
        let avgTime = testTime / 10.0
        
        print("  Total: \(String(format: "%.1f", testTime))ms for 10 tokens")
        print("  Average: \(String(format: "%.1f", avgTime))ms per token")
        print("  Speed: \(String(format: "%.1f", 1000.0 / avgTime)) tok/s")
        
        // Check if production-ready (<100ms/token)
        if avgTime < 100 {
            print("✓✓✓ PRODUCTION READY (<100ms/token)")
        } else {
            print("⚠ Need optimization (current: \(String(format: "%.1f", avgTime))ms/token, target: <100ms)")
        }
        
        XCTAssertLessThan(avgTime, 200, "Should be <200ms per token")
        
        print("\n═══════════════════════════════════════════════════════════════════")
    }
    
    func testE2BSpeed() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  E2B Inference Speed Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-12b-it-4bit"
        
        guard FileManager.default.fileExists(atPath: modelPath) else {
            print("⚠ Model not found")
            return
        }
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        print("Loading E2B...")
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
        print("✓ Model loaded")
        
        // Warm up
        _ = try model.forwardOptimized(tokenId: 2, position: 0)
        
        // Test speed
        print("\nTesting 10 tokens...")
        let testStart = Date()
        var currentToken = 2
        for i in 0..<10 {
            let result = try model.forwardOptimized(tokenId: currentToken, position: i)
            var maxIdx = 0
            var maxVal = result[0]
            for j in 1..<result.count {
                if result[j] > maxVal {
                    maxVal = result[j]
                    maxIdx = j
                }
            }
            currentToken = maxIdx
        }
        let testTime = Date().timeIntervalSince(testStart) * 1000
        let avgTime = testTime / 10.0
        
        print("  Average: \(String(format: "%.1f", avgTime))ms per token")
        print("  Speed: \(String(format: "%.1f", 1000.0 / avgTime)) tok/s")
        
        if avgTime < 100 {
            print("✓ PRODUCTION READY")
        } else {
            print("⚠ Optimization needed")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
    }
    
    func test31BSpeed() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  31B Inference Speed Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-31b-it-4bit"
        
        guard FileManager.default.fileExists(atPath: modelPath) else {
            print("⚠ Model not found")
            return
        }
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        print("Loading 31B (60 layers, 5376 hidden)...")
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
        print("✓ Model loaded")
        
        // Warm up
        _ = try model.forwardOptimized(tokenId: 2, position: 0)
        
        // Test speed
        print("\nTesting 10 tokens...")
        let testStart = Date()
        var currentToken = 2
        for i in 0..<10 {
            let result = try model.forwardOptimized(tokenId: currentToken, position: i)
            var maxIdx = 0
            var maxVal = result[0]
            for j in 1..<result.count {
                if result[j] > maxVal {
                    maxVal = result[j]
                    maxIdx = j
                }
            }
            currentToken = maxIdx
        }
        let testTime = Date().timeIntervalSince(testStart) * 1000
        let avgTime = testTime / 10.0
        
        print("  Average: \(String(format: "%.1f", avgTime))ms per token")
        print("  Speed: \(String(format: "%.1f", 1000.0 / avgTime)) tok/s")
        
        if avgTime < 100 {
            print("✓ PRODUCTION READY")
        } else {
            print("⚠ Optimization needed (larger model)")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
    }
}