#
# FFT overlap-add
#
# Take FFT of windowed chunks of input data and overlap-add them back to a
# resulting time-domain signal
#
# TODO: make it work for hopsizes other than N/2
#
#
# © 2020 Marc Groenewegen
#

import os
import sys
import numpy as np
from scipy.signal import get_window
from scipy.fftpack import fft,ifft
from scipy.io.wavfile import read
from scipy.io.wavfile import write

#
# WARNING: always look at the output file and turn down the volume
#          before you play it!
#

def overlapAdd(infile,outfile):

  INT16_FAC = (2**15)-1
  INT32_FAC = (2**31)-1
  INT64_FAC = (2**63)-1
  norm_fact = {'int16':INT16_FAC, 'int32':INT32_FAC, 'int64':INT64_FAC,'float32':1.0,'float64':1.0}

  # read input file and if necessary, convert it to mono
  (fs,insamples) = read(infile)

  # take only first (mono) channel
  insamples=insamples[:,0]

  # convert input to the desired data type for FFT
  insamples = np.float32(insamples)/norm_fact[insamples.dtype.name]

  N=2048
  hopsize=int(N/2)

  # window needs to conform to COLA criterion, Hanning does that, Hamming
  #  comes close and Blackman is not quite right for hopsize of N/2
  #  COLA: constant overlap-add

  w=get_window("hanning",N)

  xPos=0
  yResult=np.array([],dtype=np.float32) # array to hold the reconstructed signal
  prevResult=np.zeros(hopsize,dtype=np.float32) # array filled with zeros for first run

  # traverse through the input in hops
  while xPos+N < len(insamples):

    x = insamples[xPos:xPos+N] # take a frame of samples
    x = x * w # apply window

    Y=fft(x)

    # decompose spectrum into magnitude and phase
    YMag=abs(Y)
    YPhase=np.angle(Y)

    #
    # This is the place to modify the spectrum -if you want-
    #

    # reconstruct spectrum
    Y = YMag*(np.cos(YPhase) + 1j*np.sin(YPhase))

    # back to time domain
    y=ifft(Y)

    # calculate a partial buffer by taking the first part of the newly
    #  calculated IFFT and add that to the last part of the previous
    #  IFFT that was stored in prevResult
    partialResult = np.real(y[0:hopsize]) + prevResult

    # add the partial result at the end of the array
    yResult = np.concatenate((yResult,partialResult))

    prevResult = np.real(y[hopsize:N]) # remember second half of IFFT for next round

    xPos += hopsize # next hop

  # write the result
  write(outfile,fs,yResult)


if __name__ == '__main__':
  if len(sys.argv) < 2:
    print("Please enter input filename")
    sys.exit()
  infile=sys.argv[1]
  basename=os.path.splitext(sys.argv[1])[0]
  extension=os.path.splitext(sys.argv[1])[1]
  outfile=basename + "_out" + extension
  print("Writing to " + outfile)
  overlapAdd(infile,outfile)