import XCTest
@testable import MarkBase

class CompleteModelComparisonTest: XCTestCase {
    
    let models = [
        ("E4B", "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"),
        ("12B", "/Users/accusys/MarkBaseEngine/models/gemma-4-12b-it-4bit"),
        ("31B", "/Users/accusys/MarkBaseEngine/models/gemma-4-31b-it-4bit"),
        ("E2B", "/Users/accusys/MarkBaseEngine/models/gemma-4-e2b-it-4bit"),
        ("26B-Std", "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-standard"),
        ("26B-A4B", "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit")
    ]
    
    func testAllModelsBasicLoading() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  PHASE 1: Basic Loading Test (All 6 Models)")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        var results: [(String, Int, Int, Int, Int, Int, Double)] = []
        
        for (name, path) in models {
            print("───────────────────────────────────────────────────────────────────")
            print("Testing: \(name)")
            print("───────────────────────────────────────────────────────────────────")
            
            guard FileManager.default.fileExists(atPath: path) else {
                print("  ⚠ Model not found at \(path)")
                continue
            }
            
            let startTime = Date()
            
            do {
                let model = try E4BModel(modelDir: path, engine: engine, maxContextLength: 128)
                let loadTime = Date().timeIntervalSince(startTime)
                
                print("  ✓ Model loaded in \(String(format: "%.2f", loadTime))s")
                print("    Layers: \(model.numHiddenLayers)")
                print("    Hidden: \(model.hiddenSize)")
                print("    Vocab: \(model.vocabSize)")
                
                let safetensorsPath = path.contains(".cache") ? path + "/model.safetensors" :
                    (FileManager.default.fileExists(atPath: path + "/model.safetensors") ?
                    path + "/model.safetensors" : path + "/model-00001-of-00002.safetensors")
                
                if let reader = try? SafeTensorsReader(path: safetensorsPath) {
                    let allTensors = reader.allDescriptors()
                    let audioTensors = allTensors.filter { $0.name.contains("audio") || $0.name.contains("audio_tower") || $0.name.contains("embed_audio") }
                    let visionTensors = allTensors.filter { $0.name.contains("vision") || $0.name.contains("vision_tower") || $0.name.contains("embed_vision") }
                    let textTensors = allTensors.filter { !$0.name.contains("audio") && !$0.name.contains("vision") }
                    
                    print("    Audio tensors: \(audioTensors.count)")
                    print("    Vision tensors: \(visionTensors.count)")
                    print("    Text tensors: \(textTensors.count)")
                    
                    results.append((name, model.numHiddenLayers, model.hiddenSize, audioTensors.count, visionTensors.count, textTensors.count, loadTime))
                } else {
                    results.append((name, model.numHiddenLayers, model.hiddenSize, 0, 0, 0, loadTime))
                }
                
            } catch {
                print("  ✗ Failed: \(error)")
            }
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  BASIC LOADING SUMMARY")
        print("═══════════════════════════════════════════════════════════════════\n")
        print("| Model | Layers | Hidden | Audio | Vision | Text | Load Time |")
        print("|-------|--------|--------|-------|--------|------|-----------|")
        for r in results {
            print("| \(r.0) | \(r.1) | \(r.2) | \(r.3) | \(r.4) | \(r.5) | \(String(format: "%.2fs", r.6)) |")
        }
        print()
    }
    
    func testAllModelsMultimodal() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  PHASE 2: Multimodal Loading Test (Audio/Vision)")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        var audioResults: [(String, String, Int, Int)] = []
        var visionResults: [(String, String, Int, Int)] = []
        
        for (name, path) in models {
            print("───────────────────────────────────────────────────────────────────")
            print("Testing: \(name)")
            print("───────────────────────────────────────────────────────────────────")
            
            guard FileManager.default.fileExists(atPath: path) else {
                print("  ⚠ Model not found")
                continue
            }
            
            do {
                let mmModel = try MultimodalModel(modelDir: path, engine: engine, maxContextLength: 128)
                
                var audioType = "None"
                var audioLayers = 0
                var audioHidden = 0
                
                if let audioFull = mmModel.audioTowerFull {
                    audioType = "E4B-Tower"
                    audioLayers = audioFull.config.numHiddenLayers
                    audioHidden = audioFull.config.hiddenSize
                    print("  ✓ AudioTower (E4B): \(audioLayers) layers, \(audioHidden) hidden")
                } else if let audioE2B = mmModel.audioTowerE2B {
                    audioType = "E2B-Tower"
                    audioLayers = audioE2B.config.numHiddenLayers
                    audioHidden = audioE2B.config.hiddenSize
                    print("  ✓ AudioTowerE2B: \(audioLayers) layers, \(audioHidden) hidden")
                } else if let audio12B = mmModel.audioTower {
                    audioType = "12B-Embed"
                    audioLayers = 0
                    audioHidden = audio12B.config.outputDim
                    print("  ✓ AudioTower12B (Embed): outputDim=\(audioHidden)")
                } else {
                    print("  ✗ No Audio Tower")
                }
                
                audioResults.append((name, audioType, audioLayers, audioHidden))
                
                var visionType = "None"
                var visionLayers = 0
                var visionHidden = 0
                
                if let visionFull = mmModel.visionTowerFull {
                    visionType = "E4B-Tower"
                    visionLayers = visionFull.config.numHiddenLayers
                    visionHidden = visionFull.config.hiddenSize
                    print("  ✓ VisionTower (E4B): \(visionLayers) layers, \(visionHidden) hidden")
                } else if let visionE2B = mmModel.visionTowerE2B {
                    visionType = "E2B-Tower"
                    visionLayers = visionE2B.config.numHiddenLayers
                    visionHidden = visionE2B.config.hiddenSize
                    print("  ✓ VisionTowerE2B: \(visionLayers) layers, \(visionHidden) hidden")
                } else if let vision12B = mmModel.visionTower {
                    visionType = "12B-Embed"
                    visionLayers = 0
                    visionHidden = vision12B.config.hiddenDim
                    print("  ✓ VisionTower12B (Embed): hiddenDim=\(visionHidden)")
                } else {
                    print("  ✗ No Vision Tower")
                }
                
                visionResults.append((name, visionType, visionLayers, visionHidden))
                
                print("  Token IDs: Audio=\(mmModel.audioTokenId), Vision=\(mmModel.imageTokenId)")
                
            } catch {
                print("  ✗ Failed: \(error)")
                audioResults.append((name, "Error", 0, 0))
                visionResults.append((name, "Error", 0, 0))
            }
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  MULTIMODAL SUMMARY")
        print("═══════════════════════════════════════════════════════════════════\n")
        print("| Model | Audio Type | Audio Layers | Audio Hidden |")
        print("|-------|------------|-------------|-------------|")
        for r in audioResults {
            print("| \(r.0) | \(r.1) | \(r.2) | \(r.3) |")
        }
        print()
        print("| Model | Vision Type | Vision Layers | Vision Hidden |")
        print("|-------|-------------|--------------|--------------|")
        for r in visionResults {
            print("| \(r.0) | \(r.1) | \(r.2) | \(r.3) |")
        }
        print()
    }
    
    func testAllModelsForward() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  PHASE 3: Forward Pass Test (Text + Audio + Vision)")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        var results: [(String, Int, Int, Int, String)] = []
        
        for (name, path) in models {
            print("───────────────────────────────────────────────────────────────────")
            print("Testing: \(name)")
            print("───────────────────────────────────────────────────────────────────")
            
            guard FileManager.default.fileExists(atPath: path) else {
                print("  ⚠ Model not found")
                continue
            }
            
            var textNaN = 0
            var audioNaN = 0
            var visionNaN = 0
            var stability = "Perfect"
            
            do {
                let model = try E4BModel(modelDir: path, engine: engine, maxContextLength: 128)
                
                let textResult = try model.forwardOptimized(tokenId: 2, position: 0)
                textNaN = textResult.filter { $0.isNaN }.count
                print("  Text forward: NaN=\(textNaN)/\(textResult.count)")
                
                if textNaN > 0 {
                    stability = "Issue"
                }
                
            } catch {
                print("  ✗ Text forward failed: \(error)")
                stability = "Failed"
            }
            
            do {
                let mmModel = try MultimodalModel(modelDir: path, engine: engine, maxContextLength: 128)
                
                var audioForward = false
                if let audioFull = mmModel.audioTowerFull {
                    let seqLen = 50
                    let nMels = 128
                    let melFeatures = Array(repeating: Float(0.1), count: seqLen * nMels)
                    let inputBuffer = engine.device.makeBuffer(bytes: melFeatures, length: melFeatures.count * 4)!
                    let outputLen = seqLen / 4 * audioFull.config.outputProjDims
                    let outputBuffer = engine.device.makeBuffer(length: outputLen * 4)!
                    
                    try audioFull.forward(inputBuffer: inputBuffer, seqLen: seqLen, outputBuffer: outputBuffer)
                    
                    let ptr = outputBuffer.contents().assumingMemoryBound(to: Float.self)
                    audioNaN = 0
                    for i in 0..<outputLen {
                        if ptr[i].isNaN { audioNaN += 1 }
                    }
                    print("  Audio forward (E4B): NaN=\(audioNaN)/\(outputLen)")
                    audioForward = true
                    
                    if audioNaN > 0 && stability == "Perfect" {
                        stability = "Audio Issue"
                    }
                }
                
                if let audioE2B = mmModel.audioTowerE2B {
                    let seqLen = 50
                    let nMels = 128
                    let melFeatures = Array(repeating: Float(0.1), count: seqLen * nMels)
                    let inputBuffer = engine.device.makeBuffer(bytes: melFeatures, length: melFeatures.count * 4)!
                    let outputLen = seqLen / 4 * audioE2B.config.outputProjDims
                    let outputBuffer = engine.device.makeBuffer(length: outputLen * 4)!
                    
                    try audioE2B.forward(inputBuffer: inputBuffer, seqLen: seqLen, outputBuffer: outputBuffer)
                    
                    let ptr = outputBuffer.contents().assumingMemoryBound(to: Float.self)
                    audioNaN = 0
                    for i in 0..<outputLen {
                        if ptr[i].isNaN { audioNaN += 1 }
                    }
                    print("  Audio forward (E2B): NaN=\(audioNaN)/\(outputLen)")
                    audioForward = true
                    
                    if audioNaN > 0 && stability == "Perfect" {
                        stability = "Audio Issue"
                    }
                }
                
                if !audioForward {
                    if mmModel.audioTower != nil {
                        print("  Audio forward: 12B-Embed (no tower forward test)")
                    } else {
                        print("  Audio forward: N/A (no audio)")
                    }
                }
                
            } catch {
                print("  Audio forward skipped: \(error)")
            }
            
            results.append((name, textNaN, audioNaN, visionNaN, stability))
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  FORWARD PASS SUMMARY")
        print("═══════════════════════════════════════════════════════════════════\n")
        print("| Model | Text NaN | Audio NaN | Vision NaN | Stability |")
        print("|-------|---------|----------|-----------|----------|")
        for r in results {
            print("| \(r.0) | \(r.1) | \(r.2) | \(r.3) | \(r.4) |")
        }
        print()
    }
    
    func testAllModelsLongContext() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  PHASE 4: Long Context Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        let contextLengths = [128, 512]
        
        var results: [(String, Int, Int, Double, String)] = []
        
        for (name, path) in models {
            print("───────────────────────────────────────────────────────────────────")
            print("Testing: \(name)")
            print("───────────────────────────────────────────────────────────────────")
            
            guard FileManager.default.fileExists(atPath: path) else {
                print("  ⚠ Model not found")
                continue
            }
            
            var maxContextTested = 0
            var nanAt512 = 0
            var kvMemoryMB = 0.0
            var rating = "Standard"
            
            for maxLen in contextLengths {
                do {
                    let model = try E4BModel(modelDir: path, engine: engine, maxContextLength: maxLen)
                    
                    maxContextTested = maxLen
                    
                    let result = try model.forwardOptimized(tokenId: 2, position: 0)
                    let nanCount = result.filter { $0.isNaN }.count
                    
                    if maxLen == 512 {
                        nanAt512 = nanCount
                    }
                    
                    let kvMemory = 2.0 * Double(model.numHiddenLayers) * Double(maxLen) * 4.0 * 256.0 * 4.0
                    kvMemoryMB = kvMemory / 1024.0 / 1024.0
                    
                    print("  Context \(maxLen): NaN=\(nanCount), KV=\(String(format: "%.2f", kvMemoryMB))MB")
                    
                } catch {
                    print("  Context \(maxLen): Failed - \(error)")
                    break
                }
            }
            
            if maxContextTested == 512 && nanAt512 == 0 {
                rating = "Good"
            }
            
            results.append((name, maxContextTested, nanAt512, kvMemoryMB, rating))
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  LONG CONTEXT SUMMARY")
        print("═══════════════════════════════════════════════════════════════════\n")
        print("| Model | Max Tested | 512 NaN | KV Memory | Rating |")
        print("|-------|-----------|---------|----------|-------|")
        for r in results {
            print("| \(r.0) | \(r.1) | \(r.2) | \(String(format: "%.2fMB", r.3)) | \(r.4) |")
        }
        print()
    }
    
    func testAllModelsStability() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  PHASE 5: Stability Test (Repeated Forward)")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        var results: [(String, Int, String)] = []
        
        for (name, path) in models {
            print("───────────────────────────────────────────────────────────────────")
            print("Testing: \(name)")
            print("───────────────────────────────────────────────────────────────────")
            
            guard FileManager.default.fileExists(atPath: path) else {
                print("  ⚠ Model not found")
                continue
            }
            
            var totalNaN = 0
            
            do {
                let model = try E4BModel(modelDir: path, engine: engine, maxContextLength: 128)
                
                for _ in 0..<5 {
                    let result = try model.forwardOptimized(tokenId: 2, position: 0)
                    totalNaN += result.filter { $0.isNaN }.count
                }
                
                print("  Total NaN (5 passes): \(totalNaN)")
                
            } catch {
                print("  ✗ Failed: \(error)")
            }
            
            let rating = totalNaN == 0 ? "⭐⭐⭐⭐⭐ Perfect" :
                        totalNaN < 10 ? "⭐⭐⭐⭐ Good" :
                        totalNaN < 100 ? "⭐⭐⭐ OK" : "⭐⭐ Issue"
            
            results.append((name, totalNaN, rating))
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  STABILITY SUMMARY")
        print("═══════════════════════════════════════════════════════════════════\n")
        print("| Model | Total NaN | Stability |")
        print("|-------|----------|----------|")
        for r in results {
            print("| \(r.0) | \(r.1) | \(r.2) |")
        }
        print()
    }
}