import XCTest
@testable import MarkBase

final class SimpleComparisonTest: XCTestCase {
    
    func testE2BAudioOnly() throws {
        print("\n═══════════════════════════════════════")
        print("  E2B Audio Only - Simple Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-e2b-it-4bit/snapshots/2c3e507453b4f218d05fe3cc97bea5c5a654257e"
        
        print("Step 1: Create engine...")
        let engine = try MarkBaseEngine(autoCompile: true)
        print("  ✓ Engine created")
        
        print("\nStep 2: Load model...")
        let mmModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 256)
        print("  ✓ Model loaded: hidden=\(mmModel.textModel.hiddenSize)")
        print("  Audio tower: \(mmModel.audioTowerE2B != nil ? "E2B (Full)" : "N/A")")
        
        print("\nStep 3: Generate synthetic mel spectrogram...")
        let seqLen = 100
        let nMels = 128
        var melFeatures: [[Float]] = []
        for t in 0..<seqLen {
            var frame: [Float] = []
            for m in 0..<nMels {
                frame.append(Float.random(in: -0.5...0.5))
            }
            melFeatures.append(frame)
        }
        print("  ✓ Mel shape: [\(seqLen), \(nMels)]")
        
        print("\nStep 4: Process audio...")
        let start = Date()
        let audioEmbeds = try mmModel.processAudio(audioFeatures: melFeatures)
        let elapsed = Date().timeIntervalSince(start) * 1000
        
        let outputSeqLen = seqLen / 4
        let outputDim = 1536
        print("  Output shape: [\(outputSeqLen), \(outputDim)]")
        print("  Range: [\(audioEmbeds.min() ?? 0), \(audioEmbeds.max() ?? 0)]")
        print("  NaN count: \(audioEmbeds.filter { $0.isNaN }.count)")
        print("  Time: \(elapsed) ms")
        
        XCTAssertFalse(audioEmbeds.contains { $0.isNaN }, "No NaN in audio output")
        
        print("\n═══════════════════════════════════════")
        print("✓ E2B audio test passed")
        print("═══════════════════════════════════════\n")
    }
    
    func testE4BAudioOnly() throws {
        print("\n═══════════════════════════════════════")
        print("  E4B Audio Only - Simple Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        
        print("Step 1: Create engine...")
        let engine = try MarkBaseEngine(autoCompile: true)
        print("  ✓ Engine created")
        
        print("\nStep 2: Load model...")
        let mmModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 256)
        print("  ✓ Model loaded: hidden=\(mmModel.textModel.hiddenSize)")
        print("  Audio tower: \(mmModel.audioTowerFull != nil ? "E4B (Full)" : "N/A")")
        
        print("\nStep 3: Generate synthetic mel spectrogram...")
        let seqLen = 100
        let nMels = 128
        var melFeatures: [[Float]] = []
        for t in 0..<seqLen {
            var frame: [Float] = []
            for m in 0..<nMels {
                frame.append(Float.random(in: -0.5...0.5))
            }
            melFeatures.append(frame)
        }
        print("  ✓ Mel shape: [\(seqLen), \(nMels)]")
        
        print("\nStep 4: Process audio...")
        let start = Date()
        let audioEmbeds = try mmModel.processAudio(audioFeatures: melFeatures)
        let elapsed = Date().timeIntervalSince(start) * 1000
        
        let outputSeqLen = seqLen / 4
        let outputDim = 1536
        print("  Output shape: [\(outputSeqLen), \(outputDim)]")
        print("  Range: [\(audioEmbeds.min() ?? 0), \(audioEmbeds.max() ?? 0)]")
        print("  NaN count: \(audioEmbeds.filter { $0.isNaN }.count)")
        print("  Time: \(elapsed) ms")
        
        XCTAssertFalse(audioEmbeds.contains { $0.isNaN }, "No NaN in audio output")
        
        print("\n═══════════════════════════════════════")
        print("✓ E4B audio test passed")
        print("═══════════════════════════════════════\n")
    }
    
    func test12BAudioOnly() throws {
        print("\n═══════════════════════════════════════")
        print("  12B Audio Only - Simple Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Step 1: Create engine...")
        let engine = try MarkBaseEngine(autoCompile: true)
        print("  ✓ Engine created")
        
        print("\nStep 2: Load model...")
        let mmModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 256)
        print("  ✓ Model loaded: hidden=\(mmModel.textModel.hiddenSize)")
        print("  Audio tower: \(mmModel.audioTower != nil ? "12B (Projection)" : "N/A")")
        
        print("\nStep 3: Generate 640-dim audio embeddings (12B expects this, not mel)...")
        let seqLen = 100
        let audioDim = 640
        var audioEmbeds: [Float] = []
        for _ in 0..<seqLen * audioDim {
            audioEmbeds.append(Float.random(in: -0.5...0.5))
        }
        print("  ✓ Audio shape: [\(seqLen), \(audioDim)]")
        
        print("\nStep 4: Process audio through projection...")
        let inputBuffer = engine.device.makeBuffer(bytes: audioEmbeds, length: audioEmbeds.count * 4)!
        let outputBuffer = engine.device.makeBuffer(length: seqLen * mmModel.textModel.hiddenSize * 4)!
        
        let start = Date()
        if let tower = mmModel.audioTower {
            try tower.forward(inputBuffer: inputBuffer, seqLen: seqLen, outputBuffer: outputBuffer)
        } else {
            throw NSError(domain: "Audio", code: -1, userInfo: [NSLocalizedDescriptionKey: "No audio tower"])
        }
        let elapsed = Date().timeIntervalSince(start) * 1000
        
        let ptr = outputBuffer.contents().assumingMemoryBound(to: Float.self)
        let output = Array(UnsafeBufferPointer(start: ptr, count: seqLen * mmModel.textModel.hiddenSize))
        
        print("  Output shape: [\(seqLen), \(mmModel.textModel.hiddenSize)]")
        print("  Range: [\(output.min() ?? 0), \(output.max() ?? 0)]")
        print("  NaN count: \(output.filter { $0.isNaN }.count)")
        print("  Mean: \(output.reduce(0, +) / Float(output.count))")
        print("  Time: \(elapsed) ms")
        
        XCTAssertFalse(output.contains { $0.isNaN }, "No NaN in audio output")
        
        print("\n═══════════════════════════════════════")
        print("✓ 12B audio test passed")
        print("═══════════════════════════════════════\n")
    }
    
    func testE4BVisionOnly() throws {
        print("\n═══════════════════════════════════════")
        print("  E4B Vision Only - Simple Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        
        print("Step 1: Create engine...")
        let engine = try MarkBaseEngine(autoCompile: true)
        print("  ✓ Engine created")
        
        print("\nStep 2: Load model...")
        let mmModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 256)
        print("  ✓ Model loaded: hidden=\(mmModel.textModel.hiddenSize)")
        print("  Vision tower: \(mmModel.visionTowerFull != nil ? "E4B (Full 16 layers)" : "N/A")")
        
        print("\nStep 3: Generate synthetic patch embeddings...")
        let numPatches = 256
        let patchDim = 768
        var patches: [Float] = []
        for _ in 0..<numPatches * patchDim {
            patches.append(Float.random(in: -0.5...0.5))
        }
        print("  ✓ Patch shape: [\(numPatches), \(patchDim)]")
        
        print("\nStep 4: Process vision...")
        let start = Date()
        let visionEmbeds = try mmModel.processVision(patchEmbeddings: patches, numPatches: numPatches)
        let elapsed = Date().timeIntervalSince(start) * 1000
        
        print("  Output shape: [\(numPatches), \(visionEmbeds.count / numPatches)]")
        print("  Range: [\(visionEmbeds.min() ?? 0), \(visionEmbeds.max() ?? 0)]")
        print("  NaN count: \(visionEmbeds.filter { $0.isNaN }.count)")
        print("  Time: \(elapsed) ms")
        
        XCTAssertFalse(visionEmbeds.contains { $0.isNaN }, "No NaN in vision output")
        
        print("\n═══════════════════════════════════════")
        print("✓ E4B vision test passed")
        print("═══════════════════════════════════════\n")
    }
    
    func test12BVisionOnly() throws {
        print("\n═══════════════════════════════════════")
        print("  12B Vision Only - Simple Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Step 1: Create engine...")
        let engine = try MarkBaseEngine(autoCompile: true)
        print("  ✓ Engine created")
        
        print("\nStep 2: Load model...")
        let mmModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 256)
        print("  ✓ Model loaded: hidden=\(mmModel.textModel.hiddenSize)")
        print("  Vision tower: \(mmModel.visionTower != nil ? "12B (Simplified)" : "N/A")")
        
        print("\nStep 3: Generate synthetic patch embeddings...")
        let numPatches = 256
        let patchDim = 768
        var patches: [Float] = []
        for _ in 0..<numPatches * patchDim {
            patches.append(Float.random(in: -0.5...0.5))
        }
        print("  ✓ Patch shape: [\(numPatches), \(patchDim)]")
        
        print("\nStep 4: Process vision...")
        let start = Date()
        let visionEmbeds = try mmModel.processVision(patchEmbeddings: patches, numPatches: numPatches)
        let elapsed = Date().timeIntervalSince(start) * 1000
        
        print("  Output shape: [\(numPatches), \(visionEmbeds.count / numPatches)]")
        print("  Range: [\(visionEmbeds.min() ?? 0), \(visionEmbeds.max() ?? 0)]")
        print("  NaN count: \(visionEmbeds.filter { $0.isNaN }.count)")
        print("  Time: \(elapsed) ms")
        
        XCTAssertFalse(visionEmbeds.contains { $0.isNaN }, "No NaN in vision output")
        
        print("\n═══════════════════════════════════════")
        print("✓ 12B vision test passed")
        print("═══════════════════════════════════════\n")
    }
    
    func testE2BVisionOnly() throws {
        print("\n═══════════════════════════════════════")
        print("  E2B Vision Only - Simple Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-e2b-it-4bit/snapshots/2c3e507453b4f218d05fe3cc97bea5c5a654257e"
        
        print("Step 1: Create engine...")
        let engine = try MarkBaseEngine(autoCompile: true)
        print("  ✓ Engine created")
        
        print("\nStep 2: Load model...")
        let mmModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 256)
        print("  ✓ Model loaded: hidden=\(mmModel.textModel.hiddenSize)")
        print("  Vision tower: \(mmModel.visionTower != nil ? "E2B (12B variant)" : "N/A")")
        print("  Vision tower full: \(mmModel.visionTowerFull != nil ? "Available" : "N/A")")
        
        print("\nStep 3: Generate synthetic patch embeddings...")
        let numPatches = 256
        let patchDim = 768
        var patches: [Float] = []
        for _ in 0..<numPatches * patchDim {
            patches.append(Float.random(in: -0.5...0.5))
        }
        print("  ✓ Patch shape: [\(numPatches), \(patchDim)]")
        
        print("\nStep 4: Process vision...")
        let start = Date()
        do {
            let visionEmbeds = try mmModel.processVision(patchEmbeddings: patches, numPatches: numPatches)
            let elapsed = Date().timeIntervalSince(start) * 1000
            
            print("  Output shape: [\(numPatches), \(visionEmbeds.count / numPatches)]")
            print("  Range: [\(visionEmbeds.min() ?? 0), \(visionEmbeds.max() ?? 0)]")
            print("  NaN count: \(visionEmbeds.filter { $0.isNaN }.count)")
            print("  Time: \(elapsed) ms")
            
            XCTAssertFalse(visionEmbeds.contains { $0.isNaN }, "No NaN in vision output")
        } catch {
            print("  ⚠ Vision processing failed: \(error)")
            print("  E2B may not have VisionTower loaded, skipping...")
        }
        
        print("\n═══════════════════════════════════════")
        print("✓ E2B vision test completed")
        print("═══════════════════════════════════════\n")
    }
    
    func testEndToEndE4B() throws {
        print("\n═══════════════════════════════════════")
        print("  E4B End-to-End Multimodal Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        
        print("Step 1: Load model...")
        let loadStart = Date()
        let engine = try MarkBaseEngine(autoCompile: true)
        let mmModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let loadTime = Date().timeIntervalSince(loadStart) * 1000
        print("  ✓ Model loaded in \(loadTime) ms")
        
        print("\nStep 2: Process audio (mel spectrogram)...")
        let seqLen = 100
        let nMels = 128
        var melFeatures: [[Float]] = []
        for _ in 0..<seqLen {
            var frame: [Float] = []
            for _ in 0..<nMels {
                frame.append(Float.random(in: -0.5...0.5))
            }
            melFeatures.append(frame)
        }
        
        let audioStart = Date()
        let audioEmbeds = try mmModel.processAudio(audioFeatures: melFeatures)
        let audioTime = Date().timeIntervalSince(audioStart) * 1000
        print("  ✓ Audio processed in \(audioTime) ms")
        print("  Output: \(audioEmbeds.count) floats")
        
        print("\nStep 3: Process vision (patches)...")
        let numPatches = 256
        let patchDim = 768
        var patches: [Float] = []
        for _ in 0..<numPatches * patchDim {
            patches.append(Float.random(in: -0.5...0.5))
        }
        
        let visionStart = Date()
        let visionEmbeds = try mmModel.processVision(patchEmbeddings: patches, numPatches: numPatches)
        let visionTime = Date().timeIntervalSince(visionStart) * 1000
        print("  ✓ Vision processed in \(visionTime) ms")
        print("  Output: \(visionEmbeds.count) floats")
        
        print("\nStep 4: Generate text tokens...")
        var tokens: [Int] = [2]  // BOS
        let genStart = Date()
        let numTokens = 20
        for _ in 0..<numTokens {
            let logits = try mmModel.textModel.forward(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)
        let tokPerSec = Double(numTokens) / genTime
        print("  ✓ Generated \(numTokens) tokens at \(tokPerSec) tok/s")
        print("  Tokens: \(tokens.suffix(10))")
        
        let totalTime = Date().timeIntervalSince(loadStart) * 1000
        
        print("\n═══════════════════════════════════════")
        print("E4B End-to-End Summary:")
        print("  Load:     \(String(format: "%.1f", loadTime)) ms")
        print("  Audio:    \(String(format: "%.1f", audioTime)) ms")
        print("  Vision:   \(String(format: "%.1f", visionTime)) ms")
        print("  Gen:      \(String(format: "%.1f", tokPerSec)) tok/s")
        print("  Total:    \(String(format: "%.1f", totalTime)) ms")
        print("═══════════════════════════════════════\n")
    }
    
    func testEndToEnd12B() throws {
        print("\n═══════════════════════════════════════")
        print("  12B End-to-End Multimodal Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Step 1: Load model...")
        let loadStart = Date()
        let engine = try MarkBaseEngine(autoCompile: true)
        let mmModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let loadTime = Date().timeIntervalSince(loadStart) * 1000
        print("  ✓ Model loaded in \(loadTime) ms")
        
        print("\nStep 2: Process audio (640-dim embeddings)...")
        let seqLen = 100
        let audioDim = 640
        var audioEmbedsInput: [Float] = []
        for _ in 0..<seqLen * audioDim {
            audioEmbedsInput.append(Float.random(in: -0.5...0.5))
        }
        
        let inputBuffer = engine.device.makeBuffer(bytes: audioEmbedsInput, length: audioEmbedsInput.count * 4)!
        let outputBuffer = engine.device.makeBuffer(length: seqLen * mmModel.textModel.hiddenSize * 4)!
        
        let audioStart = Date()
        if let tower = mmModel.audioTower {
            try tower.forward(inputBuffer: inputBuffer, seqLen: seqLen, outputBuffer: outputBuffer)
        }
        let audioTime = Date().timeIntervalSince(audioStart) * 1000
        print("  ✓ Audio processed in \(audioTime) ms")
        
        print("\nStep 3: Process vision (patches)...")
        let numPatches = 256
        let patchDim = 768
        var patches: [Float] = []
        for _ in 0..<numPatches * patchDim {
            patches.append(Float.random(in: -0.5...0.5))
        }
        
        let visionStart = Date()
        let visionEmbeds = try mmModel.processVision(patchEmbeddings: patches, numPatches: numPatches)
        let visionTime = Date().timeIntervalSince(visionStart) * 1000
        print("  ✓ Vision processed in \(visionTime) ms")
        print("  Output: \(visionEmbeds.count) floats")
        
        print("\nStep 4: Generate text tokens...")
        var tokens: [Int] = [2]  // BOS
        let genStart = Date()
        let numTokens = 20
        for _ in 0..<numTokens {
            let logits = try mmModel.textModel.forward(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)
        let tokPerSec = Double(numTokens) / genTime
        print("  ✓ Generated \(numTokens) tokens at \(tokPerSec) tok/s")
        print("  Tokens: \(tokens.suffix(10))")
        
        let totalTime = Date().timeIntervalSince(loadStart) * 1000
        
        print("\n═══════════════════════════════════════")
        print("12B End-to-End Summary:")
        print("  Load:     \(String(format: "%.1f", loadTime)) ms")
        print("  Audio:    \(String(format: "%.1f", audioTime)) ms")
        print("  Vision:   \(String(format: "%.1f", visionTime)) ms")
        print("  Gen:      \(String(format: "%.1f", tokPerSec)) tok/s")
        print("  Total:    \(String(format: "%.1f", totalTime)) ms")
        print("═══════════════════════════════════════\n")
    }
}