import XCTest
@testable import MarkBase

final class AudioTowerLoadTest: XCTestCase {
    func testAudioTowerLoad() throws {
        print("\n═══════════════════════════════════════")
        print("      AudioTower Loading Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        
        print("Step 1: Load AudioConfig...")
        let acfg = loadAudioConfig(modelDir: modelDir)
        print("  hiddenSize: \(acfg.hiddenSize)")
        print("  numHiddenLayers: \(acfg.numHiddenLayers)")
        print("  subsamplingConvChannels: \(acfg.subsamplingConvChannels)")
        
        print("\nStep 2: Load safetensors...")
        let reader = try SafeTensorsReader(path: modelDir + "/model.safetensors")
        let descriptors = reader.allDescriptors()
        let audioKeys = descriptors.filter { $0.name.hasPrefix("audio_tower.") }
        print("  Found \(audioKeys.count) audio tensors")
        
        print("\nStep 3: Create engine...")
        let engine = try MarkBaseEngine(autoCompile: true)
        
        print("\nStep 4: Load AudioTower...")
        do {
            let tower = try loadAudioTower(reader: reader, config: acfg, engine: engine)
            print("  ✓ AudioTower loaded: \(tower.config.numHiddenLayers) layers")
        } catch {
            print("  ✗ AudioTower failed: \(error)")
            throw error
        }
        
        print("\n═══════════════════════════════════════")
        print("✓ AudioTower loading test passed")
        print("═══════════════════════════════════════\n")
    }
    
    func testAudioForward() throws {
        print("\n═══════════════════════════════════════")
        print("      AudioTower Forward Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        
        print("Step 1: Load AudioTower...")
        let engine = try MarkBaseEngine(autoCompile: true)
        let acfg = loadAudioConfig(modelDir: modelDir)
        let reader = try SafeTensorsReader(path: modelDir + "/model.safetensors")
        let tower = try loadAudioTower(reader: reader, config: acfg, engine: engine)
        print("  ✓ Loaded: \(tower.config.numHiddenLayers) layers, hidden=\(tower.config.hiddenSize)")
        
        print("\nStep 2: Create fake mel spectrogram...")
        let seqLen = 98
        let nMels = 128
        var melFeatures: [Float] = Array(repeating: 0.1, count: seqLen * nMels)
        for i in 0..<melFeatures.count {
            melFeatures[i] = Float(i % 100) / 100.0
        }
        print("  Input shape: [\(seqLen), \(nMels)]")
        
        print("\nStep 3: Run forward pass...")
        let inputBuffer = engine.device.makeBuffer(bytes: melFeatures, length: melFeatures.count * 4)!
        let outputProjDims = tower.config.outputProjDims
        let outputBuffer = engine.device.makeBuffer(length: seqLen / 4 * outputProjDims * 4)!
        
        do {
            try tower.forward(inputBuffer: inputBuffer, seqLen: seqLen, outputBuffer: outputBuffer)
            print("  ✓ Forward pass completed")
        } catch {
            print("  ✗ Forward failed: \(error)")
            throw error
        }
        
        print("\nStep 4: Check output...")
        let ptr = outputBuffer.contents().assumingMemoryBound(to: Float.self)
        let outputLen = seqLen / 4 * outputProjDims
        var hasNaN = false
        var maxVal: Float = 0
        var minVal: Float = 0
        for i in 0..<outputLen {
            let v = ptr[i]
            if v.isNaN { hasNaN = true }
            if v > maxVal { maxVal = v }
            if v < minVal { minVal = v }
        }
        print("  Output shape: [\(seqLen/4), \(outputProjDims)]")
        print("  Range: [\(minVal), \(maxVal)]")
        print("  Has NaN: \(hasNaN)")
        
        XCTAssertFalse(hasNaN, "Output should not have NaN")
        
        print("\n═══════════════════════════════════════")
        print("✓ AudioTower forward test passed")
        print("═══════════════════════════════════════\n")
    }
}