import XCTest
@testable import MarkBase

final class MultimodalComparisonTest: XCTestCase {
    
    var reporter: ComparisonReporter!
    var audioGenerator: AudioSampleGenerator!
    var visionGenerator: VisionSampleGenerator!
    
    var e2bModelDir: String!
    var e4bModelDir: String!
    var g12bModelDir: String!
    
    override func setUp() {
        reporter = ComparisonReporter()
        audioGenerator = AudioSampleGenerator()
        visionGenerator = VisionSampleGenerator()
        
        e2bModelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-e2b-it-4bit/snapshots/2c3e507453b4f218d05fe3cc97bea5c5a654257e"
        e4bModelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        g12bModelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("              MULTIMODAL COMPARISON: E2B vs E4B vs 12B")
        print("═══════════════════════════════════════════════════════════════════")
    }
    
    func testAudioComparison() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("                     AUDIO COMPARISON TEST")
        print("═══════════════════════════════════════════════════════════════════")
        
        let samples = audioGenerator.generateAll()
        var allResults: [AudioTestResult] = []
        
        for sample in samples {
            print("\n  Sample: \(sample.name)")
            print("  Description: \(sample.description)")
            
            let results = try runAudioTest(sample: sample)
            allResults.append(contentsOf: results)
            reporter.printAudioTable(results: results, sampleName: sample.name)
        }
        
        let passed = allResults.allSatisfy { $0.passed && $0.nanCount == 0 }
        XCTAssertTrue(passed, "All audio tests should pass with no NaN")
    }
    
    func testVisionComparison() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("                    VISION COMPARISON TEST")
        print("═══════════════════════════════════════════════════════════════════")
        
        let samples = visionGenerator.generateAll()
        var allResults: [VisionTestResult] = []
        
        for sample in samples {
            print("\n  Sample: \(sample.name)")
            print("  Description: \(sample.description)")
            
            let (results, _, _) = try runVisionTest(sample: sample)
            allResults.append(contentsOf: results)
            reporter.printVisionTable(results: results, sampleName: sample.name)
        }
        
        let passed = allResults.allSatisfy { $0.passed && $0.nanCount == 0 }
        XCTAssertTrue(passed, "All vision tests should pass with no NaN")
    }
    
    func testEndToEndComparison() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("                   END-TO-END COMPARISON TEST")
        print("═══════════════════════════════════════════════════════════════════")
        
        let results = try runEndToEndTest()
        reporter.printEndToEndTable(results: results)
        
        let passed = results.allSatisfy { $0.passed }
        XCTAssertTrue(passed, "All end-to-end tests should pass")
        
        reporter.printSummary()
    }
    
    func testFullReportGeneration() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("                    GENERATING FULL REPORT")
        print("═══════════════════════════════════════════════════════════════════")
        
        let audioSamples = audioGenerator.generateAll()
        var audioResults: [AudioTestResult] = []
        for sample in audioSamples {
            audioResults.append(contentsOf: try runAudioTest(sample: sample))
        }
        
        let visionSamples = visionGenerator.generateAll()
        var visionResults: [VisionTestResult] = []
        for sample in visionSamples {
            let (results, _, _) = try runVisionTest(sample: sample)
            visionResults.append(contentsOf: results)
        }
        
        let e2eResults = try runEndToEndTest()
        
        let result = ComparisonResult(
            audioResults: audioResults,
            visionResults: visionResults,
            endToEndResults: e2eResults,
            timestamp: Date()
        )
        
        try reporter.writeJSON(result: result)
        try reporter.writeCSV(result: result)
        
        print("\n  ✓ Full report generated successfully")
    }
    
    private func runAudioTest(sample: AudioSample) throws -> [AudioTestResult] {
        var results: [AudioTestResult] = []
        
        let e2bResult = try testAudioModel(
            modelDir: e2bModelDir,
            modelName: "E2B",
            sample: sample
        )
        results.append(e2bResult)
        
        let e4bResult = try testAudioModel(
            modelDir: e4bModelDir,
            modelName: "E4B",
            sample: sample
        )
        results.append(e4bResult)
        
        let g12bResult = try testAudioModel(
            modelDir: g12bModelDir,
            modelName: "12B",
            sample: sample
        )
        results.append(g12bResult)
        
        return results
    }
    
    private func testAudioModel(modelDir: String, modelName: String, sample: AudioSample) throws -> AudioTestResult {
        let engine = try MarkBaseEngine(autoCompile: true)
        let mmModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 256)
        
        let flatMel = sample.melFeatures.flatMap { $0 }
        let inputBuffer = engine.device.makeBuffer(bytes: flatMel, length: flatMel.count * 4)!
        
        let outputDim: Int
        let outputSeqLen: Int
        if modelName == "12B" {
            outputDim = mmModel.textModel.hiddenSize
            outputSeqLen = sample.seqLen
        } else {
            outputDim = 1536
            outputSeqLen = sample.seqLen / 4
        }
        
        let outputBuffer = engine.device.makeBuffer(length: outputSeqLen * outputDim * 4)!
        
        let start = Date()
        let memoryBefore = getMemoryUsage()
        
        if let tower = mmModel.audioTowerFull {
            try tower.forward(inputBuffer: inputBuffer, seqLen: sample.seqLen, outputBuffer: outputBuffer)
        } else if let tower = mmModel.audioTowerE2B {
            try tower.forward(inputBuffer: inputBuffer, seqLen: sample.seqLen, outputBuffer: outputBuffer)
        } else if let tower = mmModel.audioTower {
            try tower.forward(inputBuffer: inputBuffer, seqLen: sample.seqLen, outputBuffer: outputBuffer)
        } else {
            throw NSError(domain: "AudioTower", code: -1, userInfo: [NSLocalizedDescriptionKey: "No audio tower loaded"])
        }
        
        let elapsed = Date().timeIntervalSince(start) * 1000
        let memoryAfter = getMemoryUsage()
        
        let ptr = outputBuffer.contents().assumingMemoryBound(to: Float.self)
        let output = Array(UnsafeBufferPointer(start: ptr, count: outputSeqLen * outputDim))
        
        let stats = computeStats(data: output)
        
        return AudioTestResult(
            modelName: modelName,
            sampleName: sample.name,
            outputShape: (outputSeqLen, outputDim),
            min: stats.min,
            max: stats.max,
            mean: stats.mean,
            std: stats.std,
            nanCount: stats.nanCount,
            infCount: stats.infCount,
            forwardTimeMs: elapsed,
            memoryPeakMB: memoryAfter - memoryBefore,
            passed: stats.nanCount == 0 && stats.infCount == 0
        )
    }
    
    private func runVisionTest(sample: VisionSample) throws -> (results: [VisionTestResult], e2bOutput: [Float], e4bOutput: [Float]) {
        var results: [VisionTestResult] = []
        var e2bOut: [Float] = []
        var e4bOut: [Float] = []
        
        let (e2bResult, e2bOutput) = try testVisionModel(
            modelDir: e2bModelDir,
            modelName: "E2B",
            sample: sample
        )
        results.append(e2bResult)
        e2bOut = e2bOutput
        
        let (e4bResult, e4bOutput) = try testVisionModel(
            modelDir: e4bModelDir,
            modelName: "E4B",
            sample: sample
        )
        results.append(e4bResult)
        e4bOut = e4bOutput
        
        let (g12bResult, _) = try testVisionModel(
            modelDir: g12bModelDir,
            modelName: "12B",
            sample: sample
        )
        results.append(g12bResult)
        
        let cosSim = computeCosineSimilarity(a: e2bOut, b: e4bOut)
        if let e4bIdx = results.firstIndex(where: { $0.modelName == "E4B" }) {
            let existing = results[e4bIdx]
            results[e4bIdx] = VisionTestResult(
                modelName: existing.modelName,
                sampleName: existing.sampleName,
                outputShape: existing.outputShape,
                min: existing.min,
                max: existing.max,
                mean: existing.mean,
                std: existing.std,
                nanCount: existing.nanCount,
                infCount: existing.infCount,
                cosineSimilarity: cosSim,
                forwardTimeMs: existing.forwardTimeMs,
                passed: existing.passed
            )
        }
        
        return (results, e2bOut, e4bOut)
    }
    
    private func testVisionModel(modelDir: String, modelName: String, sample: VisionSample) throws -> (VisionTestResult, [Float]) {
        let engine = try MarkBaseEngine(autoCompile: true)
        let mmModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 256)
        
        let inputBuffer = engine.device.makeBuffer(bytes: sample.patchEmbeddings, length: sample.patchEmbeddings.count * 4)!
        
        let outputDim: Int
        if modelName == "12B" {
            outputDim = mmModel.textModel.hiddenSize
        } else {
            outputDim = mmModel.visionTowerFull?.config.hiddenSize ?? 2560
        }
        
        let outputBuffer = engine.device.makeBuffer(length: sample.numPatches * outputDim * 4)!
        
        let start = Date()
        
        if let tower = mmModel.visionTowerFull {
            try tower.forward(patchEmbeddings: inputBuffer, numPatches: sample.numPatches, outputBuffer: outputBuffer)
        } else if let tower = mmModel.visionTower {
            try tower.forward(patchEmbeddings: inputBuffer, numPatches: sample.numPatches, outputBuffer: outputBuffer)
        } else {
            throw NSError(domain: "VisionTower", code: -1, userInfo: [NSLocalizedDescriptionKey: "No vision tower loaded"])
        }
        
        let elapsed = Date().timeIntervalSince(start) * 1000
        
        let ptr = outputBuffer.contents().assumingMemoryBound(to: Float.self)
        let output = Array(UnsafeBufferPointer(start: ptr, count: sample.numPatches * outputDim))
        
        let stats = computeStats(data: output)
        
        let result = VisionTestResult(
            modelName: modelName,
            sampleName: sample.name,
            outputShape: (sample.numPatches, outputDim),
            min: stats.min,
            max: stats.max,
            mean: stats.mean,
            std: stats.std,
            nanCount: stats.nanCount,
            infCount: stats.infCount,
            cosineSimilarity: nil,
            forwardTimeMs: elapsed,
            passed: stats.nanCount == 0 && stats.infCount == 0
        )
        
        return (result, output)
    }
    
    private func runEndToEndTest() throws -> [EndToEndResult] {
        var results: [EndToEndResult] = []
        
        let audioSample = audioGenerator.generate(type: .syntheticSpeechLike)
        let visionSample = visionGenerator.generate(type: .syntheticNatural)
        
        let e2bResult = try testEndToEnd(
            modelDir: e2bModelDir,
            modelName: "E2B",
            audioSample: audioSample,
            visionSample: visionSample
        )
        results.append(e2bResult)
        
        let e4bResult = try testEndToEnd(
            modelDir: e4bModelDir,
            modelName: "E4B",
            audioSample: audioSample,
            visionSample: visionSample
        )
        results.append(e4bResult)
        
        let g12bResult = try testEndToEnd(
            modelDir: g12bModelDir,
            modelName: "12B",
            audioSample: audioSample,
            visionSample: visionSample
        )
        results.append(g12bResult)
        
        return results
    }
    
    private func testEndToEnd(modelDir: String, modelName: String, audioSample: AudioSample, visionSample: VisionSample) throws -> EndToEndResult {
        let totalStart = Date()
        
        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
        
        let audioStart = Date()
        let _ = try mmModel.processAudio(audioFeatures: audioSample.melFeatures)
        let audioTime = Date().timeIntervalSince(audioStart) * 1000
        
        let visionStart = Date()
        let _ = try mmModel.processVision(patchEmbeddings: visionSample.patchEmbeddings, numPatches: visionSample.numPatches)
        let visionTime = Date().timeIntervalSince(visionStart) * 1000
        
        let genStart = Date()
        var tokens: [Int] = [2]
        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
        
        let totalTime = Date().timeIntervalSince(totalStart) * 1000
        
        return EndToEndResult(
            modelName: modelName,
            loadTimeMs: loadTime,
            audioProcessMs: audioTime,
            visionProcessMs: visionTime,
            genSpeedTokPerSec: tokPerSec,
            totalTimeMs: totalTime,
            generatedTokens: numTokens,
            passed: true
        )
    }
    
    private func computeStats(data: [Float]) -> (min: Float, max: Float, mean: Float, std: Float, nanCount: Int, infCount: Int) {
        var minVal: Float = Float.greatestFiniteMagnitude
        var maxVal: Float = -Float.greatestFiniteMagnitude
        var sum: Float = 0
        var nanCount = 0
        var infCount = 0
        
        for v in data {
            if v.isNaN { nanCount += 1; continue }
            if v.isInfinite { infCount += 1; continue }
            if v < minVal { minVal = v }
            if v > maxVal { maxVal = v }
            sum += v
        }
        
        let validCount = data.count - nanCount - infCount
        let mean = validCount > 0 ? sum / Float(validCount) : 0
        
        var variance: Float = 0
        for v in data {
            if !v.isNaN && !v.isInfinite {
                let diff = v - mean
                variance += diff * diff
            }
        }
        let std = validCount > 0 ? sqrt(variance / Float(validCount)) : 0
        
        return (minVal, maxVal, mean, std, nanCount, infCount)
    }
    
    private func computeCosineSimilarity(a: [Float], b: [Float]) -> Float {
        guard a.count == b.count && a.count > 0 else { return 0 }
        
        var dot: Float = 0
        var normA: Float = 0
        var normB: Float = 0
        
        for i in 0..<a.count {
            if !a[i].isNaN && !b[i].isNaN {
                dot += a[i] * b[i]
                normA += a[i] * a[i]
                normB += b[i] * b[i]
            }
        }
        
        let denom = sqrt(normA) * sqrt(normB)
        return denom > 0 ? dot / denom : 0
    }
    
    private func getMemoryUsage() -> Double {
        var info = mach_task_basic_info()
        var count = mach_msg_type_number_t(MemoryLayout<mach_task_basic_info>.size) / 4
        
        let result = withUnsafeMutablePointer(to: &info) {
            $0.withMemoryRebound(to: integer_t.self, capacity: Int(count)) {
                task_info(mach_task_self_, task_flavor_t(MACH_TASK_BASIC_INFO), $0, &count)
            }
        }
        
        if result == KERN_SUCCESS {
            return Double(info.resident_size) / 1024.0 / 1024.0
        }
        return 0
    }
}