import XCTest
@testable import MarkBase

final class G12BPerformanceTests: XCTestCase {
    
    func testInferenceSpeed() throws {
        print("\n═══════════════════════════════════════")
        print("      12B Performance Benchmark")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Step 1: Initialize Metal engine...")
        let engine = try MarkBaseEngine(autoCompile: true)
        print("  ✓ Engine ready\n")
        
        print("Step 2: Load 12B model...")
        let startLoad = Date()
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let loadTime = Date().timeIntervalSince(startLoad)
        print("  ✓ Model loaded in \(loadTime) seconds\n")
        
        print("Step 3: Warm-up inference (5 tokens)...")
        for i in 0..<5 {
            _ = try model.forward(tokenId: i, position: i)
        }
        print("  ✓ Warm-up completed\n")
        
        print("Step 4: Benchmark inference (20 tokens)...")
        let startInference = Date()
        for i in 0..<20 {
            let tokenStart = Date()
            _ = try model.forward(tokenId: i, position: i)
            let tokenTime = Date().timeIntervalSince(tokenStart)
            if i >= 5 { // Skip warm-up in stats
                print("  Token \(i): \(tokenTime) seconds")
            }
        }
        let totalInferenceTime = Date().timeIntervalSince(startInference)
        let avgTokenTime = totalInferenceTime / 20.0
        let tokensPerSecond = 20.0 / totalInferenceTime
        
        print("\nPerformance Summary:")
        print("  Total time (20 tokens): \(totalInferenceTime) seconds")
        print("  Average per token: \(avgTokenTime) seconds")
        print("  Tokens per second: \(tokensPerSecond) tok/s")
        print("  Comparison: E4B achieves ~0.051s/token (~19.7 tok/s)")
        print("  12B ratio: \(avgTokenTime / 0.051)x slower than E4B\n")
        
        print("═══════════════════════════════════════")
        print("✓ Performance benchmark completed")
        print("═══════════════════════════════════════\n")
        
        // Log results for tracking
        print("RESULTS:")
        print("  avg_token_time=\(avgTokenTime)")
        print("  tokens_per_second=\(tokensPerSecond)")
        print("  model_load_time=\(loadTime)")
    }

    func test31BPerformance() throws {
        print("\n═══════════════════════════════════════")
        print("      31B Dense Performance Benchmark")
        print("═══════════════════════════════════════\n")

        let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-31b-it-4bit"

        guard FileManager.default.fileExists(atPath: modelDir + "/config.json") else {
            print("✗ 31B model not found")
            return
        }

        print("Step 1: Initialize Metal engine...")
        let engine = try MarkBaseEngine(autoCompile: true)
        print("  ✓ Engine ready\n")

        print("Step 2: Load 31B model (~18 GB)...")
        let startLoad = Date()
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)
        let loadTime = Date().timeIntervalSince(startLoad)
        print("  ✓ Model loaded in \(String(format: "%.1f", loadTime))s")
        print("  Layers: \(model.numHiddenLayers)")
        print("  Hidden: \(model.hiddenSize)\n")

        print("Step 3: Warm-up inference (3 tokens)...")
        for i in 0..<3 {
            _ = try model.forward(tokenId: i, position: i)
        }
        print("  ✓ Warm-up completed\n")

        print("Step 4: Benchmark inference (10 tokens)...")
        let startInference = Date()
        var tokenTimes: [TimeInterval] = []
        for i in 0..<10 {
            let tokenStart = Date()
            _ = try model.forward(tokenId: 2, position: i)
            let tokenTime = Date().timeIntervalSince(tokenStart)
            tokenTimes.append(tokenTime)
            print("  Token \(i): \(String(format: "%.3f", tokenTime))s")
        }
        let totalInferenceTime = Date().timeIntervalSince(startInference)
        let avgTokenTime = totalInferenceTime / 10.0
        let tokensPerSecond = 10.0 / totalInferenceTime

        print("\nPerformance Summary:")
        print("  Model load time: \(String(format: "%.1f", loadTime))s")
        print("  Total time (10 tokens): \(String(format: "%.3f", totalInferenceTime))s")
        print("  Average per token: \(String(format: "%.3f", avgTokenTime))s")
        print("  Tokens per second: \(String(format: "%.1f", tokensPerSecond)) tok/s")

        print("\n═══════════════════════════════════════")
        print("✓ 31B Performance benchmark completed")
        print("═══════════════════════════════════════\n")

        print("RESULTS:")
        print("  avg_token_time=\(avgTokenTime)")
        print("  tokens_per_second=\(tokensPerSecond)")
        print("  model_load_time=\(loadTime)")
    }
}