FFT.swift

//
//  FFTAnalyzer.swift
//  FFTVisualizer2018
//
//  Created by Mike Muszynski on 2/12/18.
//  Copyright © 2018 Mike Muszynski. All rights reserved.
//

import Cocoa
import AudioToolbox
import Accelerate

public enum FFTAnalyzerColorWeight {
    case average
    case relativeMaximum
}

public enum FFTAnalyzerColorScheme {
    case hsb, rgb
}

class FFTAnalyzer {
    private var setup: FFTSetup?
    var audioFile: AudioFileRepresentation
    
    init(audioFile: AudioFileRepresentation) {
        self.audioFile = audioFile
    }
    
    deinit {
        self.destroyFFT()
    }
    
    var bufferSize: Int = 1024 {
        didSet {
            self.setupFFT()
        }
    }
    
    private var bufferSizeInLog2: UInt {
        let bufferLog2Double = round(log2(Double(bufferSize)))
        let bufferLog2 = UInt(bufferLog2Double)
        return bufferLog2
    }
    
    func setupFFT() {
        self.setup = vDSP_create_fftsetup(bufferSizeInLog2, FFTRadix(kFFTRadix2))
    }
    
    func destroyFFT() {
        vDSP_destroy_fftsetup(self.setup)
        self.setup = nil
    }
        
    func fftPowerData(at startingSampleIndex: Int, halfWidth: Bool = true) throws -> FFTFrame {
        guard let fftSetup = self.setup else {
            print("setup fft first")
            throw NSError()
        }
        
        var outReal = [Float](repeatElement(0, count: bufferSize))
        var outImag = [Float](repeatElement(0, count: bufferSize))
        var out = DSPSplitComplex(realp: &outReal, imagp: &outImag)
        
        var buffer: [Float] = Array(repeatElement(0, count: bufferSize))
        let range = (startingSampleIndex..<min(startingSampleIndex+bufferSize, audioFile.rawData.count))
        
        buffer[0..<range.count] = audioFile.rawData[range]
        var bufferData = [Float](repeatElement(0, count: bufferSize))
        
        let windowSize = bufferSize
        var window = [Float](repeating: 0.0, count: windowSize)
        //vDSP_hann_window(&window, vDSP_Length(windowSize), Int32(vDSP_HANN_NORM))
        vDSP_hamm_window(&window, vDSP_Length(windowSize), 0)
        vDSP_vmul(buffer, 1, window, 1, &bufferData, 1, vDSP_Length(windowSize))
        
        let temp = UnsafePointer<Float>(bufferData)
        temp.withMemoryRebound(to: DSPComplex.self, capacity: bufferData.count) { (typeConvertedTransferBuffer) -> Void in
            vDSP_ctoz(typeConvertedTransferBuffer, 2, &out, 1, vDSP_Length(bufferSize / 2))
        }
        
        vDSP_fft_zip(fftSetup, &out, 1, bufferSizeInLog2, FFTDirection(FFT_FORWARD))
        
        let values = powerData(fromReal: &outReal, imaginary: &outImag, halfWidth: halfWidth)
        let frame = FFTFrame(sampleCount: bufferSize,
                             sampleRate: audioFile.sampleRate,
                             firstSampleOffset: startingSampleIndex,
                             fftData: values)
        return frame
    }
    
    private func powerData(fromReal real: inout [Float], imaginary: inout [Float], halfWidth: Bool) -> [Float] {
        var values = [Float]()
        for i in 0..<(halfWidth ? real.count/2 : real.count) {
            let realPart = real[i]
            let imagPart = imaginary[i]
            let value = sqrt(realPart * realPart + imagPart * imagPart)
            values.append(value)
        }
        return values
    }
    
    func maxFrequency(for fftData: [Float], fromSampleRate rate: Int) -> [Float] {
        let nyquist = Float(rate) / 2.0
        let binSize = nyquist / Float(fftData.count)
        
        let topThree = fftData.reduce([]) { (result, nextFloat) -> [Float] in
            var result = result
            if result.count < 3 {
                result.append(nextFloat)
            } else {
                if result.min()! < nextFloat {
                    result[0] = nextFloat
                }
            }
            return result.sorted()
        }
        
        let indices = topThree.map { (value) -> Float in
            return Float(fftData.index(of: value)!) * binSize
        }
        
        return indices
    }
    
}