import Accelerate
import Foundation
import SpeechCore

/// 128-bin log-mel extractor matching HF WhisperFeatureExtractor
/// (n_fft=400, hop=160, periodic hann, reflect padding, slaney mel,
/// log10 -> clamp to max-8 -> (x+4)/4), then zero-padded to 3000 frames.
package final class MelExtractor {
    package static let nFFT = 400
    package static let hop = 160
    package static let nMels = 128
    package static let nFreq = 201  // nFFT/2 + 1
    package static let framesPadded = 3000
    package static let maxSamples = 480_000
    /// mel-frame buckets matching the multifunction tower (5s / 10s / 30s)
    package static let buckets = [500, 1000, 3000]

    package static func bucket(forFrames frames: Int) -> Int {
        for b in buckets where frames <= b { return b }
        return buckets.last!
    }

    private let window: [Float]           // 400
    private let melFilters: [Float]       // 201*128 row-major [freq][mel]
    private let dftCos: [Float]           // 400*201 [n][f]
    private let dftSin: [Float]           // 400*201

    package init(hannWindow: [Float], melFilters: [Float]) {
        precondition(hannWindow.count == Self.nFFT)
        precondition(melFilters.count == Self.nFreq * Self.nMels)
        self.window = hannWindow
        self.melFilters = melFilters

        var cosM = [Float](repeating: 0, count: Self.nFFT * Self.nFreq)
        var sinM = [Float](repeating: 0, count: Self.nFFT * Self.nFreq)
        for n in 0..<Self.nFFT {
            for f in 0..<Self.nFreq {
                let angle = 2.0 * Double.pi * Double(f) * Double(n) / Double(Self.nFFT)
                cosM[n * Self.nFreq + f] = Float(cos(angle))
                sinM[n * Self.nFreq + f] = Float(-sin(angle))
            }
        }
        self.dftCos = cosM
        self.dftSin = sinM
    }

    /// Returns (mel: 128*bucketFrames row-major [mel][frame], encLen, bucketFrames).
    /// The output is zero-padded to the smallest bucket that fits the audio.
    package func extract(_ samples: [Float]) -> (mel: [Float], encLen: Int, bucketFrames: Int) {
        var x = samples
        if x.count > Self.maxSamples { x = Array(x[0..<Self.maxSamples]) }
        let n = x.count
        let half = Self.nFFT / 2  // 200

        // reflect pad
        var padded = [Float](repeating: 0, count: n + Self.nFFT)
        for i in 0..<half { padded[i] = x[half - i] }
        for i in 0..<n { padded[half + i] = x[i] }
        for i in 0..<half { padded[half + n + i] = x[n - 2 - i] }

        // frame count: whisper computes 1 + n/hop frames then drops the last
        let tFrames = max(n / Self.hop, 1)

        // windowed frames (T,400)
        var frames = [Float](repeating: 0, count: tFrames * Self.nFFT)
        for t in 0..<tFrames {
            let src = t * Self.hop
            for i in 0..<Self.nFFT {
                frames[t * Self.nFFT + i] = padded[src + i] * window[i]
            }
        }

        // DFT via sgemm: R = F*C (T,201), I = F*S
        var re = [Float](repeating: 0, count: tFrames * Self.nFreq)
        var im = [Float](repeating: 0, count: tFrames * Self.nFreq)
        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                    Int32(tFrames), Int32(Self.nFreq), Int32(Self.nFFT),
                    1, frames, Int32(Self.nFFT), dftCos, Int32(Self.nFreq),
                    0, &re, Int32(Self.nFreq))
        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                    Int32(tFrames), Int32(Self.nFreq), Int32(Self.nFFT),
                    1, frames, Int32(Self.nFFT), dftSin, Int32(Self.nFreq),
                    0, &im, Int32(Self.nFreq))

        // power spectrum in-place into re
        vDSP.multiply(re, re, result: &re)
        vDSP.multiply(im, im, result: &im)
        vDSP.add(re, im, result: &re)

        // mel projection: (T,201)x(201,128) -> (T,128)
        var mel = [Float](repeating: 0, count: tFrames * Self.nMels)
        cblas_sgemm(CblasRowMajor, CblasNoTrans, CblasNoTrans,
                    Int32(tFrames), Int32(Self.nMels), Int32(Self.nFreq),
                    1, re, Int32(Self.nFreq), melFilters, Int32(Self.nMels),
                    0, &mel, Int32(Self.nMels))

        // log10(max(x,1e-10)); clamp to globalMax-8; (x+4)/4
        var clamped = vDSP.clip(mel, to: 1e-10...Float.greatestFiniteMagnitude)
        var logMel = [Float](repeating: 0, count: clamped.count)
        var count = Int32(clamped.count)
        vvlog10f(&logMel, &clamped, &count)
        let maxVal = vDSP.maximum(logMel)
        logMel = vDSP.clip(logMel, to: (maxVal - 8.0)...Float.greatestFiniteMagnitude)
        logMel = vDSP.add(4.0, logMel)
        logMel = vDSP.divide(logMel, 4.0)

        // transpose (T,128) -> (128, bucket) zero-padded
        let bucketFrames = Self.bucket(forFrames: tFrames)
        var out = [Float](repeating: 0, count: Self.nMels * bucketFrames)
        let tUse = min(tFrames, bucketFrames)
        for t in 0..<tUse {
            for m in 0..<Self.nMels {
                out[m * bucketFrames + t] = logMel[t * Self.nMels + m]
            }
        }
        // encoder_feature_len = min(ceil(n/hop), tFrames) per runtime
        let encLen = min(max((n + Self.hop - 1) / Self.hop, 1), tFrames)
        return (out, encLen, bucketFrames)
    }
}