import XCTest
@testable import MarkBase

final class AudioSeparateTest: XCTestCase {
    
    func testE2BAudioLoad() throws {
        print("\n═══════════════════════════════════════")
        print("  E2B Audio Load Test (AudioTowerE2B)")
        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 AudioTowerE2B...")
        let acfg = loadAudioConfig(modelDir: modelDir)
        guard let reader = try? SafeTensorsReader(path: modelDir + "/model.safetensors") else {
            print("  ✗ No safetensors file")
            throw NSError(domain: "Audio", code: -1, userInfo: [NSLocalizedDescriptionKey: "No safetensors"])
        }
        
        let loadStart = Date()
        let audioTower = try loadAudioTowerE2B(reader: reader, config: acfg, engine: engine)
        let loadTime = Date().timeIntervalSince(loadStart) * 1000
        
        print("  ✓ Loaded in \(loadTime) ms")
        print("  Type: AudioTowerE2B (\(audioTower.config.numHiddenLayers) layers)")
        print("  Config: hiddenSize=\(audioTower.config.hiddenSize)")
        print("  Output: \(audioTower.config.outputProjDims)")
        
        print("\nStep 3: Generate synthetic 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)
        }
        print("  ✓ Mel shape: [\(seqLen), \(nMels)]")
        
        print("\nStep 4: Process audio forward pass...")
        let flatFeatures = melFeatures.flatMap { $0 }
        let inputBuffer = engine.device.makeBuffer(bytes: flatFeatures, length: flatFeatures.count * 4)!
        let outputBuffer = engine.device.makeBuffer(length: seqLen / 4 * audioTower.config.outputProjDims * 4)!
        
        let forwardStart = Date()
        try audioTower.forward(inputBuffer: inputBuffer, seqLen: seqLen, outputBuffer: outputBuffer)
        let forwardTime = Date().timeIntervalSince(forwardStart) * 1000
        
        let ptr = outputBuffer.contents().assumingMemoryBound(to: Float.self)
        let output = Array(UnsafeBufferPointer(start: ptr, count: seqLen / 4 * audioTower.config.outputProjDims))
        
        print("  ✓ Forward pass completed in \(forwardTime) ms")
        print("  Output shape: [\(seqLen / 4), \(audioTower.config.outputProjDims)]")
        print("  Range: [\(output.min() ?? 0), \(output.max() ?? 0)]")
        print("  NaN count: \(output.filter { $0.isNaN }.count)")
        
        XCTAssertFalse(output.contains { $0.isNaN }, "Audio output should not have NaN")
        
        print("\n═══════════════════════════════════════")
        print("✓ E2B audio test passed")
        print("  Load: \(loadTime) ms, Forward: \(forwardTime) ms")
        print("═══════════════════════════════════════\n")
    }
    
    func testE4BAudioLoad() throws {
        print("\n═══════════════════════════════════════")
        print("  E4B Audio Load Test (AudioTowerFull)")
        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 AudioTowerFull...")
        let acfg = loadAudioConfig(modelDir: modelDir)
        guard let reader = try? SafeTensorsReader(path: modelDir + "/model.safetensors") else {
            print("  ✗ No safetensors file")
            throw NSError(domain: "Audio", code: -1, userInfo: [NSLocalizedDescriptionKey: "No safetensors"])
        }
        
        let loadStart = Date()
        let audioTower = try loadAudioTower(reader: reader, config: acfg, engine: engine)
        let loadTime = Date().timeIntervalSince(loadStart) * 1000
        
        print("  ✓ Loaded in \(loadTime) ms")
        print("  Type: AudioTowerFull (\(audioTower.config.numHiddenLayers) layers)")
        print("  Config: hiddenSize=\(audioTower.config.hiddenSize)")
        print("  Output: \(audioTower.config.outputProjDims)")
        
        print("\nStep 3: Generate synthetic 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)
        }
        print("  ✓ Mel shape: [\(seqLen), \(nMels)]")
        
        print("\nStep 4: Process audio forward pass...")
        let flatFeatures = melFeatures.flatMap { $0 }
        let inputBuffer = engine.device.makeBuffer(bytes: flatFeatures, length: flatFeatures.count * 4)!
        let outputBuffer = engine.device.makeBuffer(length: seqLen / 4 * audioTower.config.outputProjDims * 4)!
        
        let forwardStart = Date()
        try audioTower.forward(inputBuffer: inputBuffer, seqLen: seqLen, outputBuffer: outputBuffer)
        let forwardTime = Date().timeIntervalSince(forwardStart) * 1000
        
        let ptr = outputBuffer.contents().assumingMemoryBound(to: Float.self)
        let output = Array(UnsafeBufferPointer(start: ptr, count: seqLen / 4 * audioTower.config.outputProjDims))
        
        print("  ✓ Forward pass completed in \(forwardTime) ms")
        print("  Output shape: [\(seqLen / 4), \(audioTower.config.outputProjDims)]")
        print("  Range: [\(output.min() ?? 0), \(output.max() ?? 0)]")
        print("  NaN count: \(output.filter { $0.isNaN }.count)")
        
        XCTAssertFalse(output.contains { $0.isNaN }, "Audio output should not have NaN")
        
        print("\n═══════════════════════════════════════")
        print("✓ E4B audio test passed")
        print("  Load: \(loadTime) ms, Forward: \(forwardTime) ms")
        print("═══════════════════════════════════════\n")
    }
    
    func test12BAudioLoad() throws {
        print("\n═══════════════════════════════════════")
        print("  12B Audio Load Test (AudioTower12B)")
        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 AudioTower12B...")
        let loadStart = Date()
        let audioTower = try AudioTower12B.load(modelDir: modelDir, engine: engine)
        let loadTime = Date().timeIntervalSince(loadStart) * 1000
        
        print("  ✓ Loaded in \(loadTime) ms")
        print("  Type: AudioTower12B (projection only)")
        print("  Config: inDim=\(audioTower.config.audioDim), outDim=\(audioTower.config.outputDim)")
        
        print("\nStep 3: Generate synthetic audio embeddings (640-dim)...")
        let seqLen = 100
        let audioDim = audioTower.config.audioDim
        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 * audioTower.config.outputDim * 4)!
        
        let forwardStart = Date()
        try audioTower.forward(inputBuffer: inputBuffer, seqLen: seqLen, outputBuffer: outputBuffer)
        let forwardTime = Date().timeIntervalSince(forwardStart) * 1000
        
        let ptr = outputBuffer.contents().assumingMemoryBound(to: Float.self)
        let output = Array(UnsafeBufferPointer(start: ptr, count: seqLen * audioTower.config.outputDim))
        
        print("  ✓ Forward pass completed in \(forwardTime) ms")
        print("  Output shape: [\(seqLen), \(audioTower.config.outputDim)]")
        print("  Range: [\(output.min() ?? 0), \(output.max() ?? 0)]")
        print("  NaN count: \(output.filter { $0.isNaN }.count)")
        
        XCTAssertFalse(output.contains { $0.isNaN }, "Audio output should not have NaN")
        
        print("\n═══════════════════════════════════════")
        print("✓ 12B audio test passed")
        print("  Load: \(loadTime) ms, Forward: \(forwardTime) ms")
        print("═══════════════════════════════════════\n")
    }
}