import XCTest
@testable import MarkBase

final class OptimizedForwardTest: XCTestCase {
    
    func testE4BOptimizedForward() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  E4B Optimized Forward Pass Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        
        print("Step 1: Load TEXT model...")
        let engine = try MarkBaseEngine(autoCompile: true)
        let loadStart = Date()
        let textModel = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 256)
        let loadTime = Date().timeIntervalSince(loadStart) * 1000
        print("  ✓ Model loaded in \(loadTime) ms")
        print("  Hidden: \(textModel.hiddenSize), Layers: \(textModel.numHiddenLayers)")
        
        print("\nStep 2: Test ORIGINAL forward (baseline)...")
        var tokens1: [Int] = [2]
        let start1 = Date()
        let logits1 = try textModel.forward(tokenId: 2, position: 0)
        let time1 = Date().timeIntervalSince(start1) * 1000
        print("  ✓ Original forward: \(time1) ms")
        print("  Logits: max=\(logits1.max() ?? 0), min=\(logits1.min() ?? 0)")
        print("  Top 5: \(logits1.enumerated().sorted(by: { $0.element > $1.element }).prefix(5).map { "\($0.offset): \($0.element)" })")
        
        print("\nStep 3: Test OPTIMIZED forward...")
        let start2 = Date()
        let logits2 = try textModel.forwardOptimized(tokenId: 2, position: 0)
        let time2 = Date().timeIntervalSince(start2) * 1000
        print("  ✓ Optimized forward: \(time2) ms")
        print("  Logits: max=\(logits2.max() ?? 0), min=\(logits2.min() ?? 0)")
        print("  Top 5: \(logits2.enumerated().sorted(by: { $0.element > $1.element }).prefix(5).map { "\($0.offset): \($0.element)" })")
        
        print("\nStep 4: Compare results...")
        let speedup = time1 / time2
        let maxDiff = abs(logits1.max()! - logits2.max()!)
        let minDiff = abs(logits1.min()! - logits2.min()!)
        print("  Speedup: \(speedup)x")
        print("  Max difference: \(maxDiff)")
        print("  Min difference: \(minDiff)")
        
        // Compare top tokens
        let top1 = logits1.enumerated().sorted(by: { $0.element > $1.element }).prefix(5).map { $0.offset }
        let top2 = logits2.enumerated().sorted(by: { $0.element > $1.element }).prefix(5).map { $0.offset }
        print("  Top 5 tokens match: \(top1 == top2)")
        
        print("\nStep 5: Generate 10 tokens (optimized)...")
        var tokens: [Int] = [2]
        let genStart = Date()
        for _ in 0..<10 {
            let logits = try textModel.forwardOptimized(tokenId: tokens.last!, position: tokens.count - 1)
            var maxIdx = 0
            var maxLogit = logits[0]
            for i in 1..<logits.count {
                if logits[i] > maxLogit {
                    maxLogit = logits[i]
                    maxIdx = i
                }
            }
            tokens.append(maxIdx)
        }
        let genTime = Date().timeIntervalSince(genStart) * 1000
        print("  ✓ Generated 10 tokens in \(genTime) ms")
        print("  Tokens: \(tokens.suffix(10))")
        print("  Average token time: \(genTime / 10) ms")
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("✓ Optimized forward test passed")
        print("  Load: \(loadTime) ms")
        print("  Single forward: Original=\(time1)ms, Optimized=\(time2)ms")
        print("  10 tokens: \(genTime) ms (avg \(genTime/10) ms/token)")
        print("  Speedup: \(speedup)x")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        XCTAssertFalse(logits2.contains { $0.isNaN }, "Optimized logits should not have NaN")
        XCTAssertLessThan(maxDiff, 1.0, "Logits should be similar")
    }
    
    func test12BOptimizedForward() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  12B Optimized Forward Pass Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Step 1: Load TEXT model...")
        let engine = try MarkBaseEngine(autoCompile: true)
        let loadStart = Date()
        let textModel = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 256)
        let loadTime = Date().timeIntervalSince(loadStart) * 1000
        print("  ✓ Model loaded in \(loadTime) ms")
        print("  Hidden: \(textModel.hiddenSize), Layers: \(textModel.numHiddenLayers)")
        
        print("\nStep 2: Generate 10 tokens (optimized)...")
        var tokens: [Int] = [2]
        let genStart = Date()
        for _ in 0..<10 {
            let logits = try textModel.forwardOptimized(tokenId: tokens.last!, position: tokens.count - 1)
            var maxIdx = 0
            var maxLogit = logits[0]
            for i in 1..<logits.count {
                if logits[i] > maxLogit {
                    maxLogit = logits[i]
                    maxIdx = i
                }
            }
            tokens.append(maxIdx)
        }
        let genTime = Date().timeIntervalSince(genStart) * 1000
        print("  ✓ Generated 10 tokens in \(genTime) ms")
        print("  Tokens: \(tokens.suffix(10))")
        print("  Average token time: \(genTime / 10) ms")
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("✓ 12B Optimized forward test passed")
        print("  Load: \(loadTime) ms")
        print("  10 tokens: \(genTime) ms (avg \(genTime/10) ms/token)")
        print("═══════════════════════════════════════════════════════════════════\n")
    }
}