//
//  FFT.swift
//
// See: https://stackoverflow.com/a/60120843

import Foundation
import Accelerate

public class FFT {
    
    private var fftSetup:vDSP.FFT<DSPDoubleSplitComplex>
    
    private var length: Int
    
    private var forwardInputRealPtr:UnsafeMutableBufferPointer<Double>
    private var forwardInputImagPtr:UnsafeMutableBufferPointer<Double>
    private var forwardOutputRealPtr:UnsafeMutableBufferPointer<Double>
    private var forwardOutputImagPtr:UnsafeMutableBufferPointer<Double>
    private var magnitudes:UnsafeMutableBufferPointer<Double>
    
    init(length: Int) {
        self.length = length
        
        // The power of two of two times the length of the input.
        // Do not forget this factor 2.
        let log2n = vDSP_Length(ceil(log2(Double(length * 2))))
        // Create the instance of the FFT class which allow computing FFT of complex vector with length
        // up to `length`.
        self.fftSetup = vDSP.FFT(log2n: log2n, radix: .radix2, ofType: DSPDoubleSplitComplex.self)!
        
        // --- Input / Output arrays
        self.forwardInputRealPtr = UnsafeMutableBufferPointer<Double>.allocate(capacity: self.length)
        self.forwardInputImagPtr = UnsafeMutableBufferPointer<Double>.allocate(capacity: self.length)
        self.forwardOutputRealPtr = UnsafeMutableBufferPointer<Double>.allocate(capacity: self.length)
        self.forwardOutputImagPtr = UnsafeMutableBufferPointer<Double>.allocate(capacity: self.length)
        self.magnitudes = UnsafeMutableBufferPointer<Double>.allocate(capacity: self.length)
    }
    
    deinit {
        forwardInputRealPtr.deallocate()
        forwardInputImagPtr.deallocate()
        forwardOutputRealPtr.deallocate()
        forwardOutputImagPtr.deallocate()
        magnitudes.deallocate()
    }
    
    public func compute(_ signal: [Double]) -> [Double] {
        var result = [Double](repeating: 0, count: Int(length))
        
        _ = self.forwardInputRealPtr.initialize(from: signal)
        forwardInputImagPtr.initialize(repeating: 0.0)
        
        // Input
        let forwardInput = DSPDoubleSplitComplex(
            realp: self.forwardInputRealPtr.baseAddress!,
            imagp: self.forwardInputImagPtr.baseAddress!)
        
        // Output
        var forwardOutput = DSPDoubleSplitComplex(
            realp: self.forwardOutputRealPtr.baseAddress!,
            imagp: self.forwardOutputImagPtr.baseAddress!)
        
        fftSetup.transform(
            input: forwardInput,
            output: &forwardOutput,
            direction: .forward)
        
        vDSP.absolute(forwardOutput, result: &self.magnitudes)
        
        for i in 0..<self.length/2 {
            result[i] = self.magnitudes[i]
        }
        
        return result
    }
}