import Foundation

public enum AudioSampleType {
    case syntheticUniform
    case syntheticSine
    case syntheticSpeechLike
    case syntheticNatural
}

public struct AudioSample {
    public let type: AudioSampleType
    public let name: String
    public let melFeatures: [[Float]]
    public let seqLen: Int
    public let nMels: Int
    public let description: String
}

public final class AudioSampleGenerator {
    public let nMels: Int = 128
    public let seqLen: Int = 100
    
    public init() {}
    
    public func generate(type: AudioSampleType) -> AudioSample {
        switch type {
        case .syntheticUniform:
            return generateUniform()
        case .syntheticSine:
            return generateSine()
        case .syntheticSpeechLike:
            return generateSpeechLike()
        case .syntheticNatural:
            return generateNatural()
        }
    }
    
    public func generateAll() -> [AudioSample] {
        return [
            generateUniform(),
            generateSine(),
            generateSpeechLike(),
            generateNatural()
        ]
    }
    
    private func generateUniform() -> AudioSample {
        var mel: [[Float]] = []
        for _ in 0..<seqLen {
            var frame: [Float] = []
            for _ in 0..<nMels {
                frame.append(Float.random(in: -1.0...1.0) * 0.5)
            }
            mel.append(frame)
        }
        return AudioSample(
            type: .syntheticUniform,
            name: "synthetic_uniform",
            melFeatures: mel,
            seqLen: seqLen,
            nMels: nMels,
            description: "Uniform random distribution [-0.5, 0.5]"
        )
    }
    
    private func generateSine() -> AudioSample {
        var mel: [[Float]] = []
        let frequencies: [Float] = [100.0, 200.0, 400.0, 800.0, 1600.0]
        for t in 0..<seqLen {
            var frame: [Float] = []
            for m in 0..<nMels {
                let freqIdx = m % frequencies.count
                let freq = frequencies[freqIdx]
                let phase = Float(t) / Float(seqLen)
                let value = sin(2.0 * Float.pi * freq * phase / 100.0) * 0.5
                let melScale = Float(m) / Float(nMels)
                frame.append(value * melScale + Float.random(in: -0.05...0.05))
            }
            mel.append(frame)
        }
        return AudioSample(
            type: .syntheticSine,
            name: "synthetic_sine",
            melFeatures: mel,
            seqLen: seqLen,
            nMels: nMels,
            description: "Multi-frequency sine wave simulation"
        )
    }
    
    private func generateSpeechLike() -> AudioSample {
        var mel: [[Float]] = []
        let fundamentalFreq: Float = 150.0
        let harmonics = 8
        for t in 0..<seqLen {
            var frame: [Float] = []
            let envelope = sin(Float.pi * Float(t) / Float(seqLen))
            let voicing: Float = t > 10 && t < seqLen - 10 ? 1.0 : 0.1
            for m in 0..<nMels {
                let melFreq = Float(m) * 80.0
                var value: Float = 0
                for h in 1..<harmonics {
                    let harmonicFreq = fundamentalFreq * Float(h)
                    if abs(melFreq - harmonicFreq) < 40.0 {
                        let strength = 1.0 / Float(h)
                        value += sin(2.0 * Float.pi * harmonicFreq * Float(t) / 16000.0) * strength
                    }
                }
                value *= envelope * voicing * Float(0.3)
                value += Float.random(in: -0.02...0.02)
                frame.append(value)
            }
            mel.append(frame)
        }
        return AudioSample(
            type: .syntheticSpeechLike,
            name: "synthetic_speech_like",
            melFeatures: mel,
            seqLen: seqLen,
            nMels: nMels,
            description: "Speech-like with fundamental + harmonics + envelope"
        )
    }
    
    private func generateNatural() -> AudioSample {
        var mel: [[Float]] = []
        for t in 0..<seqLen {
            var frame: [Float] = []
            let speechActivity = Float.random(in: 0.3...1.0)
            let noiseFloor: Float = 0.05
            for m in 0..<nMels {
                let melBand = Float(m) / Float(nMels)
                let spectralTilt = pow(melBand, -0.5)
                let baseValue = spectralTilt * speechActivity * 0.4
                let noise = Float.random(in: -noiseFloor...noiseFloor)
                let temporalMod = sin(2.0 * Float.pi * Float(t) / 20.0) * 0.1
                frame.append(baseValue + noise + temporalMod)
            }
            mel.append(frame)
        }
        return AudioSample(
            type: .syntheticNatural,
            name: "synthetic_natural",
            melFeatures: mel,
            seqLen: seqLen,
            nMels: nMels,
            description: "Natural audio statistics (spectral tilt + noise)"
        )
    }
    
    public func flattenMel(mel: [[Float]]) -> [Float] {
        return mel.flatMap { $0 }
    }
}

public struct AudioTestResult {
    public let modelName: String
    public let sampleName: String
    public let outputShape: (Int, Int)
    public let min: Float
    public let max: Float
    public let mean: Float
    public let std: Float
    public let nanCount: Int
    public let infCount: Int
    public let forwardTimeMs: Double
    public let memoryPeakMB: Double
    public let passed: Bool
}