import Accelerate
import CoreML
import Foundation
import SpeechCore

/// Unified audio tower (multifunction Core ML model, ANE):
/// mel bucket -> function tower_5s / tower_10s / tower_30s -> hidden
/// -> projector (CPU). Masks come from LMDecoder.maskGen (390-sized, sliced).
package final class AudioTower {
    private let modelURL: URL
    private let computeUnits: MLComputeUnits
    private var modelsByBucket: [Int: MLModel] = [:]
    private let store: AssetStore
    private static let functionByBucket = [500: "tower_5s", 1000: "tower_10s", 3000: "tower_30s"]
    package static let tokensByBucket = [500: 65, 1000: 130, 3000: 390]

    package init(unifiedModelURL: URL, store: AssetStore,
                computeUnits: MLComputeUnits = .cpuAndNeuralEngine) throws {
        self.computeUnits = computeUnits
        self.store = store
        if unifiedModelURL.pathExtension == "mlpackage" {
            modelURL = try MLModel.compileModel(at: unifiedModelURL)
        } else {
            modelURL = unifiedModelURL
        }
    }

    /// Loads (and caches) the function for a bucket. First load per bucket
    /// triggers the ANE compile; call warmUp() at app start.
    private func model(forBucket bucket: Int) throws -> MLModel {
        if let m = modelsByBucket[bucket] { return m }
        let config = MLModelConfiguration()
        config.computeUnits = computeUnits
        config.functionName = Self.functionByBucket[bucket]!
        let m = try MLModel(contentsOf: modelURL, configuration: config)
        modelsByBucket[bucket] = m
        return m
    }

    package func warmUp(buckets: [Int] = [500, 3000]) {
        for b in buckets { _ = try? model(forBucket: b) }
    }

    private static func multiArray(_ values: [Float], shape: [NSNumber]) throws -> MLMultiArray {
        let arr = try MLMultiArray(shape: shape, dataType: .float32)
        values.withUnsafeBufferPointer { src in
            arr.dataPointer.bindMemory(to: Float.self, capacity: values.count)
                .update(from: src.baseAddress!, count: values.count)
        }
        return arr
    }

    private static func floats(_ arr: MLMultiArray) -> [Float] {
        switch arr.dataType {
        case .float32:
            return arr.withUnsafeBufferPointer(ofType: Float.self) { Array($0) }
        case .float16:
            return arr.withUnsafeBufferPointer(ofType: Float16.self) { $0.map(Float.init) }
        default:
            return (0..<arr.count).map { arr[$0].floatValue }
        }
    }

    /// mel (128*bucketFrames) + 390-sized masks -> projected embeddings (N, 512)
    package func embed(mel: [Float], bucketFrames: Int, attnMask390: [Float],
                      validMask390: [Bool], sampleCount: Int) throws -> [[Float]] {
        let tokens = Self.tokensByBucket[bucketFrames]!

        // slice the 390x390 additive mask to tokens x tokens, fp16-safe fill
        var mask = [Float](repeating: 0, count: tokens * tokens)
        for q in 0..<tokens {
            for k in 0..<tokens {
                mask[q * tokens + k] = max(attnMask390[q * 390 + k], -3e4)
            }
        }

        let m = try model(forBucket: bucketFrames)
        let out = try m.prediction(from: MLDictionaryFeatureProvider(dictionary: [
            "audios": Self.multiArray(mel, shape: [1, 128, NSNumber(value: bucketFrames)]),
            "attn_mask": Self.multiArray(mask, shape: [1, 1, NSNumber(value: tokens), NSNumber(value: tokens)]),
        ]))
        let hidden = Self.floats(out.featureValue(for: "hidden")!.multiArrayValue!)  // tokens*1024
        let dim = 1024

        var valid: [[Float]] = []
        for i in 0..<tokens where validMask390[i] {
            valid.append(Array(hidden[(i * dim)..<((i + 1) * dim)]))
        }

        let target = Self.arkAudioTokenCount(sampleCount: sampleCount)
        if valid.count != target {
            valid = Self.adaptiveAvgPool(valid, outputSize: target)
        }

        return valid.map { row in
            let normed = Self.layerNorm(row, weight: store.projNormW, bias: store.projNormB)
            var out = store.projLinB
            cblas_sgemv(CblasRowMajor, CblasNoTrans, 512, 1024,
                        1, store.projLinW, 1024, normed, 1, 1, &out, 1)
            return out
        }
    }

    package static func arkAudioTokenCount(sampleCount: Int, hop: Int = 160, mergeFactor: Int = 4) -> Int {
        let melFrames = sampleCount / hop
        let downsampled = (melFrames + 1) / 2
        return max(downsampled / mergeFactor, 1)
    }

    static func adaptiveAvgPool(_ x: [[Float]], outputSize: Int) -> [[Float]] {
        let inputSize = x.count
        guard inputSize != outputSize, inputSize > 0 else { return x }
        let dim = x[0].count
        var out: [[Float]] = []
        out.reserveCapacity(outputSize)
        for i in 0..<outputSize {
            let start = i * inputSize / outputSize
            var end = (((i + 1) * inputSize) + outputSize - 1) / outputSize
            end = max(end, start + 1)
            var acc = [Float](repeating: 0, count: dim)
            for r in start..<end { vDSP.add(acc, x[r], result: &acc) }
            out.append(vDSP.divide(acc, Float(end - start)))
        }
        return out
    }

    static func layerNorm(_ x: [Float], weight: [Float], bias: [Float], eps: Float = 1e-5) -> [Float] {
        let mean = vDSP.mean(x)
        let centered = vDSP.add(-mean, x)
        let variance = vDSP.meanSquare(centered)
        let inv = 1.0 / sqrt(variance + eps)
        var out = vDSP.multiply(inv, centered)
        out = vDSP.multiply(out, weight)
        return vDSP.add(out, bias)
    }
}