# functions that implement analysis and synthesis of sounds using the Harmonic plus Residual Model
# (for example usage check the models_interface directory)

import math
import numpy as np
from scipy.signal import blackmanharris, triang
from scipy.fftpack import fft, ifft, fftshift

from . import dft, harmonic, sine
from .. import utils

def fromAudio(x, fs, w, N, H, t, minSineDur, nH, minf0, maxf0, f0et, harmDevSlope):
	"""Analysis of a sound using the harmonic plus residual model
	x: input sound, fs: sampling rate, w: analysis window; N: FFT size, t: threshold in negative dB,
	minSineDur: minimum duration of sinusoidal tracks
	nH: maximum number of harmonics; minf0: minimum fundamental frequency in sound
	maxf0: maximum fundamental frequency in sound; f0et: maximum error accepted in f0 detection algorithm
	harmDevSlope: allowed deviation of harmonic tracks, higher harmonics have higher allowed deviation
	returns hfreq, hmag, hphase: harmonic frequencies, magnitude and phases; xr: residual signal
	"""

	# perform harmonic analysis
	hfreq, hmag, hphase = harmonic.fromAudio(x, fs, w, N, H, t, nH, minf0, maxf0, f0et, harmDevSlope, minSineDur)
	Ns = 512
	xr = utils.sineSubtraction(x, Ns, H, hfreq, hmag, hphase, fs)    	# subtract sinusoids from original sound
	return hfreq, hmag, hphase, xr

def toAudio(hfreq, hmag, hphase, xr, N, H, fs):
	"""
	Synthesis of a sound using the sinusoidal plus residual model
	tfreq, tmag, tphase: sinusoidal frequencies, amplitudes and phases; stocEnv: stochastic envelope
	N: synthesis FFT size; H: hop size, fs: sampling rate
	returns y: output sound, yh: harmonic component
	"""

	yh = sine.toAudio(hfreq, hmag, hphase, N, H, fs)          # synthesize sinusoids
	y = yh[:min(yh.size, xr.size)]+xr[:min(yh.size, xr.size)]      # sum sinusoids and residual components
	return y, yh

def reconstruct(x, fs, w, N, t, nH, minf0, maxf0, f0et):
	"""
	Analysis/synthesis of a sound using the harmonic plus residual model
	x: input sound, fs: sampling rate, w: analysis window,
	N: FFT size (minimum 512), t: threshold in negative dB,
	nH: maximum number of harmonics, minf0: minimum f0 frequency in Hz,
	maxf0: maximim f0 frequency in Hz,
	f0et: error threshold in the f0 detection (ex: 5),
	maxhd: max. relative deviation in harmonic detection (ex: .2)
	returns y: output sound, yh: harmonic component, xr: residual component
	"""

	hN = N/2                                                      # size of positive spectrum
	hM1 = int(math.floor((w.size+1)/2))                           # half analysis window size by rounding
	hM2 = int(math.floor(w.size/2))                               # half analysis window size by floor
	Ns = 512                                                      # FFT size for synthesis (even)
	H = Ns/4                                                      # Hop size used for analysis and synthesis
	hNs = Ns/2
	pin = max(hNs, hM1)                                           # initialize sound pointer in middle of analysis window
	pend = x.size - max(hNs, hM1)                                 # last sample to start a frame
	fftbuffer = np.zeros(N)                                       # initialize buffer for FFT
	yhw = np.zeros(Ns)                                            # initialize output sound frame
	xrw = np.zeros(Ns)                                            # initialize output sound frame
	yh = np.zeros(x.size)                                         # initialize output array
	xr = np.zeros(x.size)                                         # initialize output array
	w = w / sum(w)                                                # normalize analysis window
	sw = np.zeros(Ns)
	ow = triang(2*H)                                              # overlapping window
	sw[hNs-H:hNs+H] = ow
	bh = blackmanharris(Ns)                                       # synthesis window
	bh = bh / sum(bh)                                             # normalize synthesis window
	wr = bh                                                       # window for residual
	sw[hNs-H:hNs+H] = sw[hNs-H:hNs+H] / bh[hNs-H:hNs+H]
	hfreqp = []
	f0t = 0
	f0stable = 0
	while pin<pend:
	#-----analysis-----
		x1 = x[pin-hM1:pin+hM2]                                     # select frame
		mX, pX = dft.fromAudio(x1, w, N)                              # compute dft
		ploc = utils.peakDetection(mX, t)                              # find peaks
		iploc, ipmag, ipphase = utils.peakInterp(mX, pX, ploc)         # refine peak values
		ipfreq = fs * iploc/N                                       # convert locations to Hz
		f0t = utils.f0Twm(ipfreq, ipmag, f0et, minf0, maxf0, f0stable) # find f0
		if ((f0stable==0)&(f0t>0)) \
			or ((f0stable>0)&(np.abs(f0stable-f0t)<f0stable/5.0)):
			f0stable = f0t                                            # consider a stable f0 if it is close to the previous one
		else:
			f0stable = 0
		hfreq, hmag, hphase = harmonic.findHarmonics(ipfreq, ipmag, ipphase, f0t, nH, hfreqp, fs) # find harmonics
		hfreqp = hfreq
		ri = pin-hNs-1                                             # input sound pointer for residual analysis
		xw2 = x[ri:ri+Ns]*wr                                       # window the input sound
		fftbuffer = np.zeros(Ns)                                   # reset buffer
		fftbuffer[:hNs] = xw2[hNs:]                                # zero-phase window in fftbuffer
		fftbuffer[hNs:] = xw2[:hNs]
		X2 = fft(fftbuffer)                                        # compute FFT of input signal for residual analysis
		#-----synthesis-----
		Yh = utils.genSpecSines(hfreq, hmag, hphase, Ns, fs)          # generate sines
		Xr = X2-Yh                                                 # get the residual complex spectrum
		fftbuffer = np.zeros(Ns)
		fftbuffer = np.real(ifft(Yh))                              # inverse FFT of harmonic spectrum
		yhw[:hNs-1] = fftbuffer[hNs+1:]                            # undo zero-phase window
		yhw[hNs-1:] = fftbuffer[:hNs+1]
		fftbuffer = np.zeros(Ns)
		fftbuffer = np.real(ifft(Xr))                              # inverse FFT of residual spectrum
		xrw[:hNs-1] = fftbuffer[hNs+1:]                            # undo zero-phase window
		xrw[hNs-1:] = fftbuffer[:hNs+1]
		yh[ri:ri+Ns] += sw*yhw                                     # overlap-add for sines
		xr[ri:ri+Ns] += sw*xrw                                     # overlap-add for residual
		pin += H                                                   # advance sound pointer
	y = yh+xr                                                    # sum of harmonic and residual components
	return y, yh, xr
