
# -*- coding: utf-8 -*_
"""
fbe.py
=========================================================================
FBE module provides several utilities and signal parametrization methods. 
"""


__author__ = '<NAME>'

import spectrum

import numpy as np

from typing import Tuple

from scipy.io import wavfile

import scipy.signal

import os

import soundfile as sf

class sin2cos2:

    """
    Class for computing signal windowing function with sin(x)^2 and cos(x)^2 tails.

    :param frame: the frame length in samples
    :type frame: int
    :param overlap: the size of the overlaping part of the window (the length of the tails on both sides)
    :type overlap: int

    :return: nothing

    """

    def __init__(self, frame : int = 512, overlap : int = 50):

        self._win = np.zeros((frame,))

        self._frame = frame

        self._overlap = overlap

        self._compute_window()

    def _compute_window(self):

        for i in range(self._overlap):

            self._win[i] = np.sin(2*np.pi/(4*(self._overlap+2))*(i+1))**2

        for i in range(self._overlap,self._frame-self._overlap):

            self._win[i] = 1

        for i in range(self._frame-self._overlap,self._frame):

            self._win[i] = np.cos(2*np.pi/(4*(self._overlap+2))*(i-self._frame+self._overlap+1))**2

    def window(self):

        """
        Method returning the vector of window's values.

        :return: the window

        :rtype: numpy array of length frame

        """

        return self._win

class fbe:

    """
    Versatile class computing various speech signal representations, mostly based on AR modelling and Mel Frequency
    Filterbanks. 

    :param frame_zero_adding: required length of the sequence after zero adding operation, defaults to None, which indicates no zero adding
    :type frame_zero_adding: int

    :param frame: frame length in samples
    :type frame: int

    :param sr: sampling frequency in Hz
    :type sr: float

    :param preem_alfa: the preemphasis coefficient
    :type preem_alfa: float

    :param freq_range: frequency range in which the mel frequency filterbanks should be computed
    :type freq_range: np.ndarray two elemements vector of floats

    :param filts_num: number of mel frequency triangular filters in the filterbank
    :type filts_num: int

    :param window: the windowing function
    :type window: np.ndarray, numpy vector of floats, defaults to None, which causes using of rectangular window

    :param ar_order: the AR model order
    :type ar_order: int

    :param cepstral_lifter: the cepstral lifter in MFCC computation
    :type cepstral_lifter: int

    :param num_ceps: number of cepstra
    :type num_ceps: int

    :returns: nothing

    .. note:: PSD is abbreviation for power spectral density in the whole documentation. AR is abbreviation for
                autoregressive in the whole documentation.

    """

    def __init__(self, frame_zero_adding=None, frame=512, sr=16000, preem_alfa=0.95, overlap=0,
                 freq_range=[20., 8000.], filts_num=23, num_gfs=70, spl_of_max_amplitude=88,
                 window=None, ar_order=16, cepstral_lifter=22, num_ceps=13):

        if overlap==0 or overlap > frame/2:
            overlap = frame/2

        if window is None:
            window = np.ones((frame,))

        if frame != len(window):
            print("ERROR in fbe, frame and window lengths do not match, program exits ...")
            sys.exit(1)

        self.sr = sr  # sampling frequency in Hz

        self.frame = frame  # number of samples in the frame

        self.num_ceps = num_ceps

        if not frame_zero_adding is None:
            self._nfft = frame_zero_adding  # fft length, sets the self._nfft atribute
        else:
            self._nfft = frame

        self.preem_alfa = preem_alfa  # preemphasis coefficient

        self.freq_range = freq_range  # frequency range in Hz

        self.filts_num = filts_num  # number of triangular filterbank channels

        self.K = int(self._nfft / 2.) + 1  # length of the unique part of the FFT

        self.f_min = 0

        self.f_max = float(sr) / 2.

        self.f_low = self.freq_range[0]

        self.f_high = self.freq_range[1]

        # matrices

        self._tfb = self._tfb()  # compute the mel-frequency triangular filterbank, sets the H atribute

        self._pinv_tfb = self._pinv_tfb()

        self._wgh_mat = self._wgh_mat()

        self._inv_wgh_mat = self._inv_wgh_mat()

        # window

        self._window = window

        self._ar_order = ar_order

        # compute cepstral lifter

        L = cepstral_lifter
        N = num_ceps

        self.cepstral_lifter =  1+0.5*L*np.sin(np.pi*np.asarray(range(N))/float(L))

        # dct matrix

        self.dctmat = np.zeros((self.num_ceps,self.filts_num))

        for i in range(self.num_ceps):

            for j in range(self.filts_num):

                self.dctmat[i,j] = np.sqrt(2./self.filts_num) * np.cos(np.pi*i/self.filts_num*(j+.5))

        self.lst_elm = ['fr','frwin','fft','mag','ang','psd','senmatpsd','senpsd','lpc','var_lpc','armag','arpsd','fbe',\
                        'fbekaldi','arfbe','wgh','arwgh','sfbe','sarfbe','smag','spsd','sarmag','sarpsd','sfbewgh',\
                        'smagwgh','spsdwgh','senmatspsdwgh','senspsdwgh','sarfbewgh','sarpsdwgh','psdspl']

        self.results = {}

        self._reset()

    def _reset(self): 

        """ 
        Resets the cache. 

        """

        for e in self.lst_elm:

            self.results[e] = None


    def get_frame_len(self) -> int:

        """Returns the frame length in samples

        :return: the frame length

        :rtype: int

        """

        return self.frame

    def get_tfb(self) -> np.ndarray:

        """Gets the triangular mel frequency filterbank.

        :return: the filter matrix containing in each row a single filter

        :rtype: np.ndarray, numpy array with filts_num rows

        """

        return self._tfb

    def get_wgh(self) -> np.ndarray:

        """Gets the weighting matrix, which is a square of the product of pseudo inverses of the Jacobian of the linear
        magnitude spectrum filter banks transform.

        :return: the weighting matrix

        :rtype: numpy array with dimension filts_num x filts_num

        """
        return self._wgh_mat

    def get_inv_wgh(self) -> np.ndarray:

        """
        Gets pseudo inverse of the weighting matrix.

        :returns: the pseudo inverse of the weighting matrix

        :rtype: np.ndarray, numpy array with dimension filts_num x filts_num


        """

        return self._inv_wgh_mat

    def get_pinv_tfb(self) -> np.ndarray:

        """
        Gets the pseudoinverse of the filterbanks matrix.

        :returns: the pseudo inverse of the weighting matrix

        :rtype: np.ndarray, numpy array with dimension filts_num x filts_num


        """

        return self._pinv_tfb

    def window(self) -> np.ndarray:

        """
        Gets the signal windowing function.

        :returns: the windowing function

        :rtype: np.ndarray, numpy array with dimension 1 x frame

        """

        return self._window

    def _tfb(self):
        """
        Computes the mel frequency triangular filterbank.

        """

        # filter cutoff frequencies (Hz) for all filters, size 1x(M+2)

        aux = np.linspace(0, self.filts_num + 1, self.filts_num + 2)

        c = self._mel2hz(
            (self._hz2mel(self.f_low) + aux * (self._hz2mel(self.f_high) - self._hz2mel(self.f_low)) / float(self.filts_num + 1)))

        f = np.linspace(self.f_min, self.f_max, self.K)

        H = np.zeros((self.filts_num, self.K))

        for m in range(self.filts_num):

            a = list(f >= c[m])

            b = list(f <= c[m + 1])

            k = np.array([a[i] and b[i] for i in range(len(f))])

            H[m, k] = (f[k] - c[m]) / (c[m + 1] - c[m])

            a = list(f >= c[m + 1])

            b = list(f <= c[m + 2])

            k = np.array([a[i] and b[i] for i in range(len(f))])

            H[m, k] = (c[m + 2] - f[k]) / (c[m + 2] - c[m + 1])

        return H

    def _proj_symmat_pd(self,A : np.ndarray,rcond : float) -> np.ndarray:
        """Projecting matrix A onto space of positive definite matrices.

        :param A: matrix to be projected
        :type A: np.ndarray
        :param rcond: reciprocal condition number of the resulting matrix
        :type rcond: float

        :return: the projected matrix
        :rtype: np.ndarray
        """

        A = .5*(A.T+A)

        w, v = np.linalg.eigh(A)

        w[w<rcond*np.max(w)] = rcond*np.max(w)

        f = (np.sqrt(w) * v).T

        f = f.T.dot(f)

        return f

    def _wgh_mat(self):

        """
        The weighting matrix.

        """

        W = np.dot(self._pinv_tfb.T, self._pinv_tfb)

        W = self._proj_symmat_pd(W,10.e-3)

        f = np.linalg.cholesky(W).T

        return f

    def _inv_wgh_mat(self):

        """
        The inverse of the weighting matrix

        """

        return np.linalg.pinv(self._wgh_mat)

    def _pinv_tfb(self):

        """
        The pseudoinverse of the mel frequency triangular filter banks matrix

        """

        return np.linalg.pinv(self._tfb)

    def _nextpow2(self, i):

        """
        The next nearest power of 2.

        """

        n = 1

        while n < i: n *= 2

        self._nfft = n

    def _hz2mel(self, hz):

        """
        Hertz to mel frequency scale

        """

        mel = 1127. * np.log(1 + hz / 700.)

        return mel

    def _mel2hz(self, mel):

        """
        Mel to frequency scale

        """

        hz = 700. * np.exp(mel / 1127.) - 700.

        return hz


    def idx2freq(self, i : int) -> float:

        """Converts frequency index to frequency in Hz

        :param i: frequency index
        :type i: int

        :return: frequency in Hz
        :rtype: float

        """

        f = float(i)/self._nfft*self.sr

        return f

    def freq2idx(self,f : float) -> int:

        """Converts frequency in Hz to the frequency index

        :param f: frequency in Herz
        :type f: float

        :return: frequency index
        :rtype: int


        """

        idx = int(np.round(f*float(self._nfft)/self.sr))

        return idx


    def set_frm(self,fr : np.ndarray):

        """
        Sets the frame of the signal - this is then used to compute the all signal representations

        :param fr: signal frame
        :type fr: np.ndarray, numpy vector of floats

        :return: nothing

        """

        self._reset()

        self.results['fr'] = fr

    def set_wgh(self,wgh : np.ndarray):

        """
        Set compact spectrum

        :param wgh: the compact specturm with filt_num elements
        :type wgh: numpy vector of floats

        :returns: nothing

        """

        self._reset()

        self.results['wgh'] = wgh

    def set_arwgh(self,arwgh : np.ndarray):

        """
        Set AR compact spectrum

        :param arwgh: the compact autoregresive specturm with filt_num elements
        :type arwgh: np.ndarray, numpy vector of floats

        :returns: nothing

        """

        self._reset()

        self.results['arwgh'] = arwgh

    def set_fbe(self,fbe : np.ndarray):

        """
        Set filterbank energies

        :param fbe: the filter bank energies (vector with filt_num elements)
        :type fbe: np.ndarray, numpy vector of floats

        :returns: nothing

        """

        self._reset()

        self.results['fbe'] = fbe


    def set_mag(self,mag : np.ndarray):

        """Set magnitude spectrum

        :param mag: the magnitude spectrum
        :type mag: np.ndarray, numpy vector of floats

        :returns: nothing


        """

        self._reset()

        self.results['mag'] = mag

    def set_psd(self,psd : np.ndarray):

        """
        Set power density spectrum

        :param psd: the power density spectrum
        :type psd: np.ndarray, numpy vector of floats

        :returns: nothing

        """

        self._reset()

        self.results['psd'] = psd

    def fr(self) -> np.ndarray:

        """
        Gets frame

        :returns: the frame
        :rtype: np.ndarray, numpy vector of floats

        """

        if self.results['fr'] is None:

            print("Frame not given (emtpy vector), program exits ...")

            sys.exit(1)

        else:

            return self.results['fr']

    def fr_win(self) -> np.ndarray:

        """
        Gets windowed frame

        :returns: the windowed frame
        :rtype: np.ndarray, numpy vector of floats

        """

        if self.results['frwin'] is None:

            self.results['frwin'] = np.zeros((self._nfft,))

            self.results['frwin'][:self.frame] = self.fr() * self.window()

        else:

            pass

        return self.results['frwin']

    def fft(self) -> np.ndarray:

        """
        Gets FFT

        :returns: the fft of the, possibly zero added, signal frame
        :rtype: np.ndarray, numpy vector of complex floats

        """

        if self.results['fft'] is None:

            self.results['fft'] = np.fft.fft(self.fr_win())

        else:

            pass

        return self.results['fft']

    def mag(self) -> np.ndarray:

        """
        Gets magnitude spectrum

        :returns: the magnitude of the, possibly zero added, signal frame
        :rtype: np.ndarray, numpy vector of floats

        """

        if self.results['mag'] is None:

            self.results['mag'] = np.abs(self.fft())[ : self.K]

        else:

            pass

        return self.results['mag']

    def ang(self) -> np.ndarray:

        """
        Gets angular spectrum.

        :returns: the angular spectrum of the, possibly zero added, signal frame
        :rtype: np.ndarray, numpy vector of floats

        """

        if self.results['ang'] is None:

            self.results['ang'] = np.angle( self.fft() )

        else:

            pass

        return self.results['ang']

    def psd(self) -> np.ndarray:

        """
        Gets power density spectrum

        :returns: the PSD of the, possibly zero added, signal frame
        :rtype: np.ndarray, numpy vector of floats


        """

        if self.results['psd'] is None:

            self.results['psd'] = self.mag()**2.

        else:

            pass

        return self.results['psd']

    def lpc(self) -> np.ndarray:

        """
        Gets LPC coefficients aka STP coefficients or AR coefficients

        :return: LPC with the leading 1
        :rtype: np.ndarray

        """

        if self.results['lpc'] is None:

            _lpc, self.results['var_lpc'], k = spectrum.aryule(self.fr_win(), self._ar_order)

            self.results['lpc'] = np.concatenate((np.array([1]),_lpc))

        else:

            pass

        return self.results['lpc']

    def var_lpc(self) -> np.ndarray:

        """
        Gets variance of the short term residual spectrum

        :return: short term residual variance
        :rtype: np.ndarray

        """

        if self.results['var_lpc'] is None:

            self.results['lpc'], self.results['var_lpc'], k = spectrum.aryule(self.fr_win(), self._ar_order)

        else:

            pass

        return self.results['var_lpc']

    def set_ar(self,a,var):
        """Setting AR coefficients and STP residual variance

        :param a: AR coefficients with leading one
        :param var: variance of the short term residual
        """

        """Sets the AR coefficients"""

        self._reset()

        self.results['var_lpc'] = var

        self.results['lpc'] = a

    def armag(self) -> np.ndarray:

        """
        Gets AR magnitude spectrum

        :return: AR magnitude spectrum
        :rtype: np.ndarray, numpy vector of floats of length _nfft/2+1

        """

        if self.results['armag'] is None:

            p = len(self.lpc())-1

            aux = np.concatenate([self.lpc(),np.zeros((self._nfft-p-1,))],axis=0)

            fftaux = np.abs(np.fft.fft(aux))

            std = np.sqrt(self.var_lpc()*self._nfft)

            self.results['armag'] = np.real(std/fftaux[ : self.K])

        else:

            pass

        return self.results['armag']

    def arpsd(self) -> np.ndarray:

        """
        Gets AR power density spectrum

        :return: the AR PSD
        :rtype: np.ndarray

        """


        if self.results['arpsd'] is None:

            self.results['arpsd'] = self.armag() ** 2.

        else:

            pass

        return self.results['arpsd']

    def fbe(self) -> np.ndarray:

        """
        Gets filter banks outputs based on magnitude spectrum

        :return: filter bank filtered magnitude spectrum
        :rtype: np.ndarray, numpy vector of floats of length filt_num

        """

        if self.results['fbe'] is None:

            self.results['fbe'] = np.dot(self.get_tfb(),self.mag())

        else:

            pass

        return self.results['fbe']


    def sfbe2mfcc(self) -> np.ndarray:

        """
        Converts smoothed filter banks energies to MFCC coefficients

        :return: MFCC coefficients
        :rtype: np.ndarray, numpy vector of floats (size num_cep)

        """

        fbe = self.sfbe()

        logfbe = np.log(fbe)

        mfcc = np.dot(self.dctmat,logfbe) #scipy.fftpack.dct(logfbe,n=self.num_ceps,norm='ortho')

        # liftering

        cmfcc = self.cepstral_lifter*mfcc

        return cmfcc

    def fbe2mfcc(self) -> np.ndarray:

        """
        Converts filter banks energies to MFCC coefficients

        :return: MFCC coefficients
        :rtype: np.ndarray, numpy vector of floats (size num_cep)

        """

        fbe = self.fbe()

        logfbe = np.log(fbe)

        mfcc = np.dot(self.dctmat,logfbe) #scipy.fftpack.dct(logfbe,n=self.num_ceps,norm='ortho')

        # here comes liftering

        cmfcc = self.cepstral_lifter*mfcc

        return cmfcc


    def arfbe(self) -> np.ndarray:

        """
        AR magnitude spectrum to filter banks energies

        :return: filter bank filtered AR magnitude spectrum
        :rtype: np.ndarray, numpy vector of floats (size num_filt)

        """

        if self.results['arfbe'] is None:

            self.results['arfbe'] = np.dot(self.get_tfb(),self.armag())

        else:

            pass

        return self.results['arfbe']

    def wgh(self) -> np.ndarray:
        """
        Weighted filter bank energies

        :return: the magnitude compact spectrum
        :rtype: np.ndarray, numpy vector of floats (size num_filt)

        """

        if self.results['wgh'] is None:

            self.results['wgh'] = np.dot(self.get_wgh(),self.fbe())

        else:

            pass

        return self.results['wgh']

    def arwgh(self) -> np.ndarray:
        """
        AR weighted filter bank energies

        :return: the AR magnitude compact spectrum
        :rtype: np.ndarray, numpy vector of floats (size num_filt)

        """

        if self.results['arwgh'] is None:

            self.results['arwgh']  = np.real(np.dot(self.get_wgh(),self.arfbe()))

        else:

            pass

        return self.results['arwgh']

    def smag(self) -> np.ndarray:
        """
        Smoothed magnitude spectrum

        :return: magnitude spectrum computed from filter bank energies
        :rtype: np.ndarray, numpy vector of floats (size _nfft/2+1)

        """

        if self.results['smag'] is None:

            self.results['smag'] = np.dot(self.get_pinv_tfb(), self.fbe())

        else:

            pass

        return self.results['smag']

    def spsd(self) -> np.ndarray:

        """
        Smoothed power density spectrum

        :return: PSD computed from filter bank energies
        :rtype: np.ndarray, numpy vector of floats(size _nfft/2+1)

        """

        if self.results['spsd'] is None:

            self.results['spsd'] = self.smag()**2.

        else:

            pass

        return self.results['spsd']


    def sarmag(self)->np.ndarray:

        """
        Smoothed AR magnitude spectrum

        :return: smoothed (from arfbe) AR magnitude spectrum (size _nfft/2+1)
        :rtype: np.ndarray, numpy vector of floats

        """

        if self.results['sarmag'] is None:

            self.results['sarmag'] = np.dot(self.get_pinv_tfb(), self.arfbe())

        else:

            pass

        return self.results['sarmag']

    def sarpsd(self) -> np.ndarray:

        """
        Smoothed AR PSD

        :return: smoothed (from arfbe) AR PSD (size _nfft/2+1)
        :rtype: np.ndarray, numpy vector of floats

        """

        if self.results['sarpsd'] is None:

            self.results['sarpsd']  = self.sarmag() ** 2.

        else:

            pass

        return self.results['sarpsd']

    def preemphasis(self, signal : np.ndarray) -> np.ndarray:

        """Perform preemphasis on the input signal.

        :param signal: The signal to filter.
        :type signal: np.ndarray, numpy vector of floats
        :param coeff: The preemphasis coefficient. 0 is no filter, default is 0.95.
        :type coeff: float
        :returns: the filtered signal.
        :rtype: numpy vector of floats
        """

        return np.asarray(np.append(signal[0], signal[1:] - self.preem_alfa * signal[:-1]))

    def sfbe(self) -> np.ndarray:

        """
        Smoothed filter bank energies

        :return: smoothed, from compact spectrum, filter bank energies
        :rtype: np.ndarray, numpy vector of floats (size num_filt)

        """

        if self.results['sfbe'] is None:

            self.results['sfbe'] = np.dot(self.get_inv_wgh(), self.wgh())

        else:

            pass

        return self.results['sfbe']

    def sarfbe(self) -> np.ndarray:

        """
        Smoothed AR filter bank energies

        :return smoothed, from compact AR spectrum, filter bank energies
        :rtype: np.ndarray, numpy vector of floats (size num_filt)

        """

        if self.results['sarfbe'] is None:

            self.results['sarfbe'] = np.dot(self.get_inv_wgh(), self.arwgh())

        else:

            pass

        return self.results['sarfbe']

    def smagwgh(self) -> np.ndarray:

        """
        Smoothed magnitude spectrum

        :return: computed from compact spectum magnitude spectrum
        :rtype: np.ndarray, numpy vector of floats (size _nfft/2+1)

        """

        if self.results['smagwgh'] is None:

            self.results['smagwgh'] = np.dot(self.get_pinv_tfb(), self.sfbe())

        else:

            pass

        return self.results['smagwgh']

    def sarmagwgh(self) -> np.ndarray:

        """
        Smoothed AR magnitude spectrum

        :return: computed from AR compact spectrum magnitude spectrum
        :rtype: np.ndarray, numpy vector of floats (size _nfft/2+1)

        """

        if self.results['sarmagwgh'] is None:

            self.results['sarmagwgh'] = np.dot(self.get_pinv_tfb(), self.sarfbe())

        else:

            pass

        return self.results['sarmagwgh']

    def spsdwgh(self) -> np.ndarray:

        """
        Smoothed PSD

        :return: computed from compact spectrum PSD
        :rtype: np.ndarray, numpy vector of floats (size _nfft/2+1)

        """

        if self.results['spsdwgh'] is None:

            self.results['spsdwgh'] = self.smagwgh() ** 2.

        else:

            pass

        return self.results['spsdwgh']

    def sarpsdwgh(self) -> np.ndarray:

        """
        Smoothed AR PSD

        :return: PSD computed from AR compact spectra
        :rtype: np.ndarray, numpy vector of floats (size _nfft/2+1)

        """

        if self.results['sarpsdwgh'] is None:

            self.results['sarpsdwgh'] = self.sarmagwgh() ** 2.

        else:

            pass

        return self.results['sarpsdwgh']

    def psd2wgh(self,psd : np.ndarray) -> np.ndarray:

        """
        PSD -> weighted compact spectrum

        :param psd: the PSD
        :type psd: np.ndarray, numpy vector of floats (size _nfft/2+1)

        :return: compact spectrum based on PSD
        :rtype: np.ndarray, numpy vector of floats (size num_filt)

        """

        mag = np.sqrt(psd)

        fbe = np.dot(self.get_tfb(),mag)

        return np.dot(self.get_wgh(),fbe)

    def psd2ar(self,psd : np.ndarray,LpcOrder : int) -> Tuple[np.ndarray, float]:
        """

        Converting PSD into LPC coefficients and excitation variance

        :param psd: left half of the PSD
        :type psd: numpy vector of floats (size _nfft/2+1)
        :param LpcOrder: AR model order
        :type LpcOrder: int

        :return: * (`vector of floats`) direct form AR coeff. with leading 1
                 * (`float`) the variance of the short term residual
        """

        D = len(psd)

        B = np.concatenate([psd,psd[D-2:0:-1]])

        xc = np.real(np.fft.ifft(B))

        xc = xc[:LpcOrder+1]

        a, var, k = spectrum.LEVINSON(xc)

        a = np.concatenate([[1],a])

        var = var/(2*D-2)

        return a, var


def synth(mag_enh : np.ndarray, angular_r : np.ndarray, an_win : np.ndarray):

    """
    Signal synthesis based on magnitude and angular spectra

    :param mag_enh: enhanced speech magnitude spectrum
    :type psd_enh: np.ndarray, numpy vector of floats
    :param angular_r: angular noisy signal frame spectrum
    :type angular_r: np.ndarray, numpy vector of floats
    :param an_win: windowing function
    :type an_win: np.ndarray numpy vector of floats

    :return: time domain enhanced signal frame (windowed)
    :rtype: numpy vector of floats

    """

    # X = np.sqrt( psd_enh )

    X = mag_enh

    X[-1] = 0

    enhDft = np.concatenate( (X, X[-2:0:-1]) ) * np.exp( 1j * angular_r )

    an_win = np.sqrt( an_win )

    enh_fs = an_win*np.real( np.fft.ifft( enhDft ) )

    return enh_fs

def enhance_mag( mag_r, psd_n, psd_s):

    """
    The Wiener filter in frequency domain

    :param mag_r: noisy, unprocessed signal frame magnitude spectrum
    :type psd_r: numpy vector of floats
    :param psd_n: noise, estimated noise frame PSD
    :type psd_n: numpy vector of floats
    :param psd_s: speech, estimated speech frame PSD
    :type psd_s: numpy vector of floats

    :return: enhanced speech PSD
    :rtype: numpy vector of floats

    """

    psd_r_smag = psd_s + psd_n

    mag_enh = np.maximum(psd_s, 1.e-6) / np.maximum(psd_r_smag, 1.e-4) * mag_r

    mag_enh = np.maximum(mag_enh,0.001*mag_r)

    mag_enh[-1] = 0 # filter out the most high frequency peak

    return mag_enh

def enhance( mag_r, psd_n, psd_s, angular_r, an_win ):

    """
    The Wiener filter returning the time frequency signal frame

    :param mag_r: noisy, unprocessed signal frame magnitude spectrum
    :type mag_r: numpy vector of floats
    :param psd_n: noise, estimated noise frame PSD
    :type psd_n: numpy vector of floats
    :param psd_s: speech, estimated speech frame PSD
    :type psd_s: numpy vector of floats
    :param angular_r: angular noisy signal frame spectrum
    :type angular_r: numpy vector of floats
    :param an_win: windowing function
    :type an_win: numpy vector of floats

    :return: time domain enhanced signal frame (windowed)
    :rtype: numpy vector of floats

    """

    mag_enh = enhance_mag(mag_r, psd_n, psd_s)

    enh_fs = synth(mag_enh,angular_r,an_win)

    return enh_fs #, psd_enh

def enhance1 ( mag_r, psd_n, psd_s, angular_r, an_win):
    """
    Perceptual enhancement (SETAP p. 245)

    :param psd_r: noisy, unprocessed signal frame PSD
    :type psd_r: numpy vector of floats
    :param psd_n: noise, estimated noise frame PSD
    :type psd_n: numpy vector of floats
    :param psd_s: speech, estimated speech frame PSD
    :type psd_s: numpy vector of floats
    :param angular_r: angular noisy signal frame spectrum
    :type angular_r: numpy vector of floats
    :param an_win: windowing function
    :type an_win: numpy vector of floats

    :return: time domain enhanced signal frame (windowed)
    :rtype: numpy vector of floats

    """

    e = psd_s/np.maximum(psd_n,1.e-4)+1.e-6

    g = mag_r**2/np.maximum(psd_n,1.e-4)+1.e-6

    v = e*g/(1+e)

    aux = np.maximum(.5*v,1.e-2)

    if np.sum(np.isnan(aux)) > 0:

        print('nan found')

        raw_input()

    gain = np.sqrt(v)/(np.sqrt(np.pi)*g)*1/scipy.special.i0e(aux)

    if np.sum(np.isnan(gain)) > 0:

        print (v)

        print (np.sqrt(v))

        print (gain)

        print (scipy.special.iv(0,aux))

        print ('nan found')

        raw_input()

    # plt.plot(gain)
    #
    # plt.show()

    mag_enh = gain*mag_r

    enh_fs = synth(mag_enh,angular_r,an_win)

    return enh_fs #, psd_enh

def enhance2(psd_r, psd_n, mt, angular_r, an_win, ro=0.05): # note that instead of psd_s we pass the masking treshold (mt)

    """
    Perceptual Wiener filter

    :param psd_r: PSD of the noisy signal frame
    :type psd_r: numpy vector of floats
    :param psd_n: PSD of noise (may be approximate, smoothed etc.)
    :type psd_n: numpy vector of floats
    :param mt: PSD masking threshold
    :type mt: numpy vector of floats
    :param angular_r: angular noisy signal frame spectrum
    :type angular_r: numpy vector of floats
    :param an_win: windowing function
    :type an_win: numpy vector of floats
    :param ro: (0.-1.) larger causes less signal modification
    :type ro: float

    :return: time domain enhanced signal frame (windowed)
    :rtype: numpy vector of floats

    """

    gain = np.minimum(np.sqrt(mt/psd_n)+ro,1)

    psd_enh = gain*psd_r

    enh_fs = synth(psd_enh,angular_r,an_win)

    return enh_fs #, psd_enh

def load_wav(file_ : str, target_sampling_rate : int = 16000) -> Tuple[np.ndarray,int,int]:
    """Loads sound in a variety of formats supported by libsndfile. Resamples the input sound to target_sampling_rate.

    :param file_: sound file path
    :type file_: str
    :param target_sampling_rate: target sampling rate
    :type target_sampling_rate: int
    :returns: * the resampled sound to target sampling rate
              * the original sound sampling rate
              * the target sampling rate
    :rtype: Tuple[np.ndarray,int,int]
    """

    if not os.path.exists(file_):
        raise ValueError(f"The file {file_} does not exist.")

    s, sr = sf.read(file_, dtype=np.float32)

    print(f"Start resmpling")

    if int(sr) != int(target_sampling_rate):
    
        s_resampled = scipy.signal.resample(s,np.uint32(np.round(target_sampling_rate/sr*s.shape[0])))

    else:

        s_resampled = s

    print(f"Resampling done")

    np.random.seed(1000) # causes the dither is same on each run

    s_resampled += np.random.randn(*s_resampled.shape)*1.e-6 # add dither to improve numerical behaviour

    return s_resampled, int(sr), target_sampling_rate

def save_wav(signal : np.ndarray, sampling_rate : int, file_name : str):

    if not os.path.exists(os.path.dirname(file_name)):
        try:
            os.makedirs(os.path.dirname(file_name))
        except OSError:
            print(f'Can not create the directory {os.path.dirname(file_name)}')
        else:
            print(f'Path {os.path.dirname(file_name)} succesfully created.')

    sf.write(file_name, signal, sampling_rate)



if __name__ == "__main__":

    pass


# -*- coding: utf-8 -*_
"""
fbe.py
=========================================================================
FBE module provides several utilities and signal parametrization methods. 
"""


__author__ = '<NAME>'

import spectrum

import numpy as np

from typing import Tuple

from scipy.io import wavfile

import scipy.signal

import os

import soundfile as sf

class sin2cos2:

    """
    Class for computing signal windowing function with sin(x)^2 and cos(x)^2 tails.

    :param frame: the frame length in samples
    :type frame: int
    :param overlap: the size of the overlaping part of the window (the length of the tails on both sides)
    :type overlap: int

    :return: nothing

    """

    def __init__(self, frame : int = 512, overlap : int = 50):

        self._win = np.zeros((frame,))

        self._frame = frame

        self._overlap = overlap

        self._compute_window()

    def _compute_window(self):

        for i in range(self._overlap):

            self._win[i] = np.sin(2*np.pi/(4*(self._overlap+2))*(i+1))**2

        for i in range(self._overlap,self._frame-self._overlap):

            self._win[i] = 1

        for i in range(self._frame-self._overlap,self._frame):

            self._win[i] = np.cos(2*np.pi/(4*(self._overlap+2))*(i-self._frame+self._overlap+1))**2

    def window(self):

        """
        Method returning the vector of window's values.

        :return: the window

        :rtype: numpy array of length frame

        """

        return self._win

class fbe:

    """
    Versatile class computing various speech signal representations, mostly based on AR modelling and Mel Frequency
    Filterbanks. 

    :param frame_zero_adding: required length of the sequence after zero adding operation, defaults to None, which indicates no zero adding
    :type frame_zero_adding: int

    :param frame: frame length in samples
    :type frame: int

    :param sr: sampling frequency in Hz
    :type sr: float

    :param preem_alfa: the preemphasis coefficient
    :type preem_alfa: float

    :param freq_range: frequency range in which the mel frequency filterbanks should be computed
    :type freq_range: np.ndarray two elemements vector of floats

    :param filts_num: number of mel frequency triangular filters in the filterbank
    :type filts_num: int

    :param window: the windowing function
    :type window: np.ndarray, numpy vector of floats, defaults to None, which causes using of rectangular window

    :param ar_order: the AR model order
    :type ar_order: int

    :param cepstral_lifter: the cepstral lifter in MFCC computation
    :type cepstral_lifter: int

    :param num_ceps: number of cepstra
    :type num_ceps: int

    :returns: nothing

    .. note:: PSD is abbreviation for power spectral density in the whole documentation. AR is abbreviation for
                autoregressive in the whole documentation.

    """

    def __init__(self, frame_zero_adding=None, frame=512, sr=16000, preem_alfa=0.95, overlap=0,
                 freq_range=[20., 8000.], filts_num=23, num_gfs=70, spl_of_max_amplitude=88,
                 window=None, ar_order=16, cepstral_lifter=22, num_ceps=13):

        if overlap==0 or overlap > frame/2:
            overlap = frame/2

        if window is None:
            window = np.ones((frame,))

        if frame != len(window):
            print("ERROR in fbe, frame and window lengths do not match, program exits ...")
            sys.exit(1)

        self.sr = sr  # sampling frequency in Hz

        self.frame = frame  # number of samples in the frame

        self.num_ceps = num_ceps

        if not frame_zero_adding is None:
            self._nfft = frame_zero_adding  # fft length, sets the self._nfft atribute
        else:
            self._nfft = frame

        self.preem_alfa = preem_alfa  # preemphasis coefficient

        self.freq_range = freq_range  # frequency range in Hz

        self.filts_num = filts_num  # number of triangular filterbank channels

        self.K = int(self._nfft / 2.) + 1  # length of the unique part of the FFT

        self.f_min = 0

        self.f_max = float(sr) / 2.

        self.f_low = self.freq_range[0]

        self.f_high = self.freq_range[1]

        # matrices

        self._tfb = self._tfb()  # compute the mel-frequency triangular filterbank, sets the H atribute

        self._pinv_tfb = self._pinv_tfb()

        self._wgh_mat = self._wgh_mat()

        self._inv_wgh_mat = self._inv_wgh_mat()

        # window

        self._window = window

        self._ar_order = ar_order

        # compute cepstral lifter

        L = cepstral_lifter
        N = num_ceps

        self.cepstral_lifter =  1+0.5*L*np.sin(np.pi*np.asarray(range(N))/float(L))

        # dct matrix

        self.dctmat = np.zeros((self.num_ceps,self.filts_num))

        for i in range(self.num_ceps):

            for j in range(self.filts_num):

                self.dctmat[i,j] = np.sqrt(2./self.filts_num) * np.cos(np.pi*i/self.filts_num*(j+.5))

        self.lst_elm = ['fr','frwin','fft','mag','ang','psd','senmatpsd','senpsd','lpc','var_lpc','armag','arpsd','fbe',\
                        'fbekaldi','arfbe','wgh','arwgh','sfbe','sarfbe','smag','spsd','sarmag','sarpsd','sfbewgh',\
                        'smagwgh','spsdwgh','senmatspsdwgh','senspsdwgh','sarfbewgh','sarpsdwgh','psdspl']

        self.results = {}

        self._reset()

    def _reset(self): 

        """ 
        Resets the cache. 

        """

        for e in self.lst_elm:

            self.results[e] = None


    def get_frame_len(self) -> int:

        """Returns the frame length in samples

        :return: the frame length

        :rtype: int

        """

        return self.frame

    def get_tfb(self) -> np.ndarray:

        """Gets the triangular mel frequency filterbank.

        :return: the filter matrix containing in each row a single filter

        :rtype: np.ndarray, numpy array with filts_num rows

        """

        return self._tfb

    def get_wgh(self) -> np.ndarray:

        """Gets the weighting matrix, which is a square of the product of pseudo inverses of the Jacobian of the linear
        magnitude spectrum filter banks transform.

        :return: the weighting matrix

        :rtype: numpy array with dimension filts_num x filts_num

        """
        return self._wgh_mat

    def get_inv_wgh(self) -> np.ndarray:

        """
        Gets pseudo inverse of the weighting matrix.

        :returns: the pseudo inverse of the weighting matrix

        :rtype: np.ndarray, numpy array with dimension filts_num x filts_num


        """

        return self._inv_wgh_mat

    def get_pinv_tfb(self) -> np.ndarray:

        """
        Gets the pseudoinverse of the filterbanks matrix.

        :returns: the pseudo inverse of the weighting matrix

        :rtype: np.ndarray, numpy array with dimension filts_num x filts_num


        """

        return self._pinv_tfb

    def window(self) -> np.ndarray:

        """
        Gets the signal windowing function.

        :returns: the windowing function

        :rtype: np.ndarray, numpy array with dimension 1 x frame

        """

        return self._window

    def _tfb(self):
        """
        Computes the mel frequency triangular filterbank.

        """

        # filter cutoff frequencies (Hz) for all filters, size 1x(M+2)

        aux = np.linspace(0, self.filts_num + 1, self.filts_num + 2)

        c = self._mel2hz(
            (self._hz2mel(self.f_low) + aux * (self._hz2mel(self.f_high) - self._hz2mel(self.f_low)) / float(self.filts_num + 1)))

        f = np.linspace(self.f_min, self.f_max, self.K)

        H = np.zeros((self.filts_num, self.K))

        for m in range(self.filts_num):

            a = list(f >= c[m])

            b = list(f <= c[m + 1])

            k = np.array([a[i] and b[i] for i in range(len(f))])

            H[m, k] = (f[k] - c[m]) / (c[m + 1] - c[m])

            a = list(f >= c[m + 1])

            b = list(f <= c[m + 2])

            k = np.array([a[i] and b[i] for i in range(len(f))])

            H[m, k] = (c[m + 2] - f[k]) / (c[m + 2] - c[m + 1])

        return H

    def _proj_symmat_pd(self,A : np.ndarray,rcond : float) -> np.ndarray:
        """Projecting matrix A onto space of positive definite matrices.

        :param A: matrix to be projected
        :type A: np.ndarray
        :param rcond: reciprocal condition number of the resulting matrix
        :type rcond: float

        :return: the projected matrix
        :rtype: np.ndarray
        """

        A = .5*(A.T+A)

        w, v = np.linalg.eigh(A)

        w[w<rcond*np.max(w)] = rcond*np.max(w)

        f = (np.sqrt(w) * v).T

        f = f.T.dot(f)

        return f

    def _wgh_mat(self):

        """
        The weighting matrix.

        """

        W = np.dot(self._pinv_tfb.T, self._pinv_tfb)

        W = self._proj_symmat_pd(W,10.e-3)

        f = np.linalg.cholesky(W).T

        return f

    def _inv_wgh_mat(self):

        """
        The inverse of the weighting matrix

        """

        return np.linalg.pinv(self._wgh_mat)

    def _pinv_tfb(self):

        """
        The pseudoinverse of the mel frequency triangular filter banks matrix

        """

        return np.linalg.pinv(self._tfb)

    def _nextpow2(self, i):

        """
        The next nearest power of 2.

        """

        n = 1

        while n < i: n *= 2

        self._nfft = n

    def _hz2mel(self, hz):

        """
        Hertz to mel frequency scale

        """

        mel = 1127. * np.log(1 + hz / 700.)

        return mel

    def _mel2hz(self, mel):

        """
        Mel to frequency scale

        """

        hz = 700. * np.exp(mel / 1127.) - 700.

        return hz


    def idx2freq(self, i : int) -> float:

        """Converts frequency index to frequency in Hz

        :param i: frequency index
        :type i: int

        :return: frequency in Hz
        :rtype: float

        """

        f = float(i)/self._nfft*self.sr

        return f

    def freq2idx(self,f : float) -> int:

        """Converts frequency in Hz to the frequency index

        :param f: frequency in Herz
        :type f: float

        :return: frequency index
        :rtype: int


        """

        idx = int(np.round(f*float(self._nfft)/self.sr))

        return idx


    def set_frm(self,fr : np.ndarray):

        """
        Sets the frame of the signal - this is then used to compute the all signal representations

        :param fr: signal frame
        :type fr: np.ndarray, numpy vector of floats

        :return: nothing

        """

        self._reset()

        self.results['fr'] = fr

    def set_wgh(self,wgh : np.ndarray):

        """
        Set compact spectrum

        :param wgh: the compact specturm with filt_num elements
        :type wgh: numpy vector of floats

        :returns: nothing

        """

        self._reset()

        self.results['wgh'] = wgh

    def set_arwgh(self,arwgh : np.ndarray):

        """
        Set AR compact spectrum

        :param arwgh: the compact autoregresive specturm with filt_num elements
        :type arwgh: np.ndarray, numpy vector of floats

        :returns: nothing

        """

        self._reset()

        self.results['arwgh'] = arwgh

    def set_fbe(self,fbe : np.ndarray):

        """
        Set filterbank energies

        :param fbe: the filter bank energies (vector with filt_num elements)
        :type fbe: np.ndarray, numpy vector of floats

        :returns: nothing

        """

        self._reset()

        self.results['fbe'] = fbe


    def set_mag(self,mag : np.ndarray):

        """Set magnitude spectrum

        :param mag: the magnitude spectrum
        :type mag: np.ndarray, numpy vector of floats

        :returns: nothing


        """

        self._reset()

        self.results['mag'] = mag

    def set_psd(self,psd : np.ndarray):

        """
        Set power density spectrum

        :param psd: the power density spectrum
        :type psd: np.ndarray, numpy vector of floats

        :returns: nothing

        """

        self._reset()

        self.results['psd'] = psd

    def fr(self) -> np.ndarray:

        """
        Gets frame

        :returns: the frame
        :rtype: np.ndarray, numpy vector of floats

        """

        if self.results['fr'] is None:

            print("Frame not given (emtpy vector), program exits ...")

            sys.exit(1)

        else:

            return self.results['fr']

    def fr_win(self) -> np.ndarray:

        """
        Gets windowed frame

        :returns: the windowed frame
        :rtype: np.ndarray, numpy vector of floats

        """

        if self.results['frwin'] is None:

            self.results['frwin'] = np.zeros((self._nfft,))

            self.results['frwin'][:self.frame] = self.fr() * self.window()

        else:

            pass

        return self.results['frwin']

    def fft(self) -> np.ndarray:

        """
        Gets FFT

        :returns: the fft of the, possibly zero added, signal frame
        :rtype: np.ndarray, numpy vector of complex floats

        """

        if self.results['fft'] is None:

            self.results['fft'] = np.fft.fft(self.fr_win())

        else:

            pass

        return self.results['fft']

    def mag(self) -> np.ndarray:

        """
        Gets magnitude spectrum

        :returns: the magnitude of the, possibly zero added, signal frame
        :rtype: np.ndarray, numpy vector of floats

        """

        if self.results['mag'] is None:

            self.results['mag'] = np.abs(self.fft())[ : self.K]

        else:

            pass

        return self.results['mag']

    def ang(self) -> np.ndarray:

        """
        Gets angular spectrum.

        :returns: the angular spectrum of the, possibly zero added, signal frame
        :rtype: np.ndarray, numpy vector of floats

        """

        if self.results['ang'] is None:

            self.results['ang'] = np.angle( self.fft() )

        else:

            pass

        return self.results['ang']

    def psd(self) -> np.ndarray:

        """
        Gets power density spectrum

        :returns: the PSD of the, possibly zero added, signal frame
        :rtype: np.ndarray, numpy vector of floats


        """

        if self.results['psd'] is None:

            self.results['psd'] = self.mag()**2.

        else:

            pass

        return self.results['psd']

    def lpc(self) -> np.ndarray:

        """
        Gets LPC coefficients aka STP coefficients or AR coefficients

        :return: LPC with the leading 1
        :rtype: np.ndarray

        """

        if self.results['lpc'] is None:

            _lpc, self.results['var_lpc'], k = spectrum.aryule(self.fr_win(), self._ar_order)

            self.results['lpc'] = np.concatenate((np.array([1]),_lpc))

        else:

            pass

        return self.results['lpc']

    def var_lpc(self) -> np.ndarray:

        """
        Gets variance of the short term residual spectrum

        :return: short term residual variance
        :rtype: np.ndarray

        """

        if self.results['var_lpc'] is None:

            self.results['lpc'], self.results['var_lpc'], k = spectrum.aryule(self.fr_win(), self._ar_order)

        else:

            pass

        return self.results['var_lpc']

    def set_ar(self,a,var):
        """Setting AR coefficients and STP residual variance

        :param a: AR coefficients with leading one
        :param var: variance of the short term residual
        """

        """Sets the AR coefficients"""

        self._reset()

        self.results['var_lpc'] = var

        self.results['lpc'] = a

    def armag(self) -> np.ndarray:

        """
        Gets AR magnitude spectrum

        :return: AR magnitude spectrum
        :rtype: np.ndarray, numpy vector of floats of length _nfft/2+1

        """

        if self.results['armag'] is None:

            p = len(self.lpc())-1

            aux = np.concatenate([self.lpc(),np.zeros((self._nfft-p-1,))],axis=0)

            fftaux = np.abs(np.fft.fft(aux))

            std = np.sqrt(self.var_lpc()*self._nfft)

            self.results['armag'] = np.real(std/fftaux[ : self.K])

        else:

            pass

        return self.results['armag']

    def arpsd(self) -> np.ndarray:

        """
        Gets AR power density spectrum

        :return: the AR PSD
        :rtype: np.ndarray

        """


        if self.results['arpsd'] is None:

            self.results['arpsd'] = self.armag() ** 2.

        else:

            pass

        return self.results['arpsd']

    def fbe(self) -> np.ndarray:

        """
        Gets filter banks outputs based on magnitude spectrum

        :return: filter bank filtered magnitude spectrum
        :rtype: np.ndarray, numpy vector of floats of length filt_num

        """

        if self.results['fbe'] is None:

            self.results['fbe'] = np.dot(self.get_tfb(),self.mag())

        else:

            pass

        return self.results['fbe']


    def sfbe2mfcc(self) -> np.ndarray:

        """
        Converts smoothed filter banks energies to MFCC coefficients

        :return: MFCC coefficients
        :rtype: np.ndarray, numpy vector of floats (size num_cep)

        """

        fbe = self.sfbe()

        logfbe = np.log(fbe)

        mfcc = np.dot(self.dctmat,logfbe) #scipy.fftpack.dct(logfbe,n=self.num_ceps,norm='ortho')

        # liftering

        cmfcc = self.cepstral_lifter*mfcc

        return cmfcc

    def fbe2mfcc(self) -> np.ndarray:

        """
        Converts filter banks energies to MFCC coefficients

        :return: MFCC coefficients
        :rtype: np.ndarray, numpy vector of floats (size num_cep)

        """

        fbe = self.fbe()

        logfbe = np.log(fbe)

        mfcc = np.dot(self.dctmat,logfbe) #scipy.fftpack.dct(logfbe,n=self.num_ceps,norm='ortho')

        # here comes liftering

        cmfcc = self.cepstral_lifter*mfcc

        return cmfcc


    def arfbe(self) -> np.ndarray:

        """
        AR magnitude spectrum to filter banks energies

        :return: filter bank filtered AR magnitude spectrum
        :rtype: np.ndarray, numpy vector of floats (size num_filt)

        """

        if self.results['arfbe'] is None:

            self.results['arfbe'] = np.dot(self.get_tfb(),self.armag())

        else:

            pass

        return self.results['arfbe']

    def wgh(self) -> np.ndarray:
        """
        Weighted filter bank energies

        :return: the magnitude compact spectrum
        :rtype: np.ndarray, numpy vector of floats (size num_filt)

        """

        if self.results['wgh'] is None:

            self.results['wgh'] = np.dot(self.get_wgh(),self.fbe())

        else:

            pass

        return self.results['wgh']

    def arwgh(self) -> np.ndarray:
        """
        AR weighted filter bank energies

        :return: the AR magnitude compact spectrum
        :rtype: np.ndarray, numpy vector of floats (size num_filt)

        """

        if self.results['arwgh'] is None:

            self.results['arwgh']  = np.real(np.dot(self.get_wgh(),self.arfbe()))

        else:

            pass

        return self.results['arwgh']

    def smag(self) -> np.ndarray:
        """
        Smoothed magnitude spectrum

        :return: magnitude spectrum computed from filter bank energies
        :rtype: np.ndarray, numpy vector of floats (size _nfft/2+1)

        """

        if self.results['smag'] is None:

            self.results['smag'] = np.dot(self.get_pinv_tfb(), self.fbe())

        else:

            pass

        return self.results['smag']

    def spsd(self) -> np.ndarray:

        """
        Smoothed power density spectrum

        :return: PSD computed from filter bank energies
        :rtype: np.ndarray, numpy vector of floats(size _nfft/2+1)

        """

        if self.results['spsd'] is None:

            self.results['spsd'] = self.smag()**2.

        else:

            pass

        return self.results['spsd']


    def sarmag(self)->np.ndarray:

        """
        Smoothed AR magnitude spectrum

        :return: smoothed (from arfbe) AR magnitude spectrum (size _nfft/2+1)
        :rtype: np.ndarray, numpy vector of floats

        """

        if self.results['sarmag'] is None:

            self.results['sarmag'] = np.dot(self.get_pinv_tfb(), self.arfbe())

        else:

            pass

        return self.results['sarmag']

    def sarpsd(self) -> np.ndarray:

        """
        Smoothed AR PSD

        :return: smoothed (from arfbe) AR PSD (size _nfft/2+1)
        :rtype: np.ndarray, numpy vector of floats

        """

        if self.results['sarpsd'] is None:

            self.results['sarpsd']  = self.sarmag() ** 2.

        else:

            pass

        return self.results['sarpsd']

    def preemphasis(self, signal : np.ndarray) -> np.ndarray:

        """Perform preemphasis on the input signal.

        :param signal: The signal to filter.
        :type signal: np.ndarray, numpy vector of floats
        :param coeff: The preemphasis coefficient. 0 is no filter, default is 0.95.
        :type coeff: float
        :returns: the filtered signal.
        :rtype: numpy vector of floats
        """

        return np.asarray(np.append(signal[0], signal[1:] - self.preem_alfa * signal[:-1]))

    def sfbe(self) -> np.ndarray:

        """
        Smoothed filter bank energies

        :return: smoothed, from compact spectrum, filter bank energies
        :rtype: np.ndarray, numpy vector of floats (size num_filt)

        """

        if self.results['sfbe'] is None:

            self.results['sfbe'] = np.dot(self.get_inv_wgh(), self.wgh())

        else:

            pass

        return self.results['sfbe']

    def sarfbe(self) -> np.ndarray:

        """
        Smoothed AR filter bank energies

        :return smoothed, from compact AR spectrum, filter bank energies
        :rtype: np.ndarray, numpy vector of floats (size num_filt)

        """

        if self.results['sarfbe'] is None:

            self.results['sarfbe'] = np.dot(self.get_inv_wgh(), self.arwgh())

        else:

            pass

        return self.results['sarfbe']

    def smagwgh(self) -> np.ndarray:

        """
        Smoothed magnitude spectrum

        :return: computed from compact spectum magnitude spectrum
        :rtype: np.ndarray, numpy vector of floats (size _nfft/2+1)

        """

        if self.results['smagwgh'] is None:

            self.results['smagwgh'] = np.dot(self.get_pinv_tfb(), self.sfbe())

        else:

            pass

        return self.results['smagwgh']

    def sarmagwgh(self) -> np.ndarray:

        """
        Smoothed AR magnitude spectrum

        :return: computed from AR compact spectrum magnitude spectrum
        :rtype: np.ndarray, numpy vector of floats (size _nfft/2+1)

        """

        if self.results['sarmagwgh'] is None:

            self.results['sarmagwgh'] = np.dot(self.get_pinv_tfb(), self.sarfbe())

        else:

            pass

        return self.results['sarmagwgh']

    def spsdwgh(self) -> np.ndarray:

        """
        Smoothed PSD

        :return: computed from compact spectrum PSD
        :rtype: np.ndarray, numpy vector of floats (size _nfft/2+1)

        """

        if self.results['spsdwgh'] is None:

            self.results['spsdwgh'] = self.smagwgh() ** 2.

        else:

            pass

        return self.results['spsdwgh']

    def sarpsdwgh(self) -> np.ndarray:

        """
        Smoothed AR PSD

        :return: PSD computed from AR compact spectra
        :rtype: np.ndarray, numpy vector of floats (size _nfft/2+1)

        """

        if self.results['sarpsdwgh'] is None:

            self.results['sarpsdwgh'] = self.sarmagwgh() ** 2.

        else:

            pass

        return self.results['sarpsdwgh']

    def psd2wgh(self,psd : np.ndarray) -> np.ndarray:

        """
        PSD -> weighted compact spectrum

        :param psd: the PSD
        :type psd: np.ndarray, numpy vector of floats (size _nfft/2+1)

        :return: compact spectrum based on PSD
        :rtype: np.ndarray, numpy vector of floats (size num_filt)

        """

        mag = np.sqrt(psd)

        fbe = np.dot(self.get_tfb(),mag)

        return np.dot(self.get_wgh(),fbe)

    def psd2ar(self,psd : np.ndarray,LpcOrder : int) -> Tuple[np.ndarray, float]:
        """

        Converting PSD into LPC coefficients and excitation variance

        :param psd: left half of the PSD
        :type psd: numpy vector of floats (size _nfft/2+1)
        :param LpcOrder: AR model order
        :type LpcOrder: int

        :return: * (`vector of floats`) direct form AR coeff. with leading 1
                 * (`float`) the variance of the short term residual
        """

        D = len(psd)

        B = np.concatenate([psd,psd[D-2:0:-1]])

        xc = np.real(np.fft.ifft(B))

        xc = xc[:LpcOrder+1]

        a, var, k = spectrum.LEVINSON(xc)

        a = np.concatenate([[1],a])

        var = var/(2*D-2)

        return a, var


def synth(mag_enh : np.ndarray, angular_r : np.ndarray, an_win : np.ndarray):

    """
    Signal synthesis based on magnitude and angular spectra

    :param mag_enh: enhanced speech magnitude spectrum
    :type psd_enh: np.ndarray, numpy vector of floats
    :param angular_r: angular noisy signal frame spectrum
    :type angular_r: np.ndarray, numpy vector of floats
    :param an_win: windowing function
    :type an_win: np.ndarray numpy vector of floats

    :return: time domain enhanced signal frame (windowed)
    :rtype: numpy vector of floats

    """

    # X = np.sqrt( psd_enh )

    X = mag_enh

    X[-1] = 0

    enhDft = np.concatenate( (X, X[-2:0:-1]) ) * np.exp( 1j * angular_r )

    an_win = np.sqrt( an_win )

    enh_fs = an_win*np.real( np.fft.ifft( enhDft ) )

    return enh_fs

def enhance_mag( mag_r, psd_n, psd_s):

    """
    The Wiener filter in frequency domain

    :param mag_r: noisy, unprocessed signal frame magnitude spectrum
    :type psd_r: numpy vector of floats
    :param psd_n: noise, estimated noise frame PSD
    :type psd_n: numpy vector of floats
    :param psd_s: speech, estimated speech frame PSD
    :type psd_s: numpy vector of floats

    :return: enhanced speech PSD
    :rtype: numpy vector of floats

    """

    psd_r_smag = psd_s + psd_n

    mag_enh = np.maximum(psd_s, 1.e-6) / np.maximum(psd_r_smag, 1.e-4) * mag_r

    mag_enh = np.maximum(mag_enh,0.001*mag_r)

    mag_enh[-1] = 0 # filter out the most high frequency peak

    return mag_enh

def enhance( mag_r, psd_n, psd_s, angular_r, an_win ):

    """
    The Wiener filter returning the time frequency signal frame

    :param mag_r: noisy, unprocessed signal frame magnitude spectrum
    :type mag_r: numpy vector of floats
    :param psd_n: noise, estimated noise frame PSD
    :type psd_n: numpy vector of floats
    :param psd_s: speech, estimated speech frame PSD
    :type psd_s: numpy vector of floats
    :param angular_r: angular noisy signal frame spectrum
    :type angular_r: numpy vector of floats
    :param an_win: windowing function
    :type an_win: numpy vector of floats

    :return: time domain enhanced signal frame (windowed)
    :rtype: numpy vector of floats

    """

    mag_enh = enhance_mag(mag_r, psd_n, psd_s)

    enh_fs = synth(mag_enh,angular_r,an_win)

    return enh_fs #, psd_enh

def enhance1 ( mag_r, psd_n, psd_s, angular_r, an_win):
    """
    Perceptual enhancement (SETAP p. 245)

    :param psd_r: noisy, unprocessed signal frame PSD
    :type psd_r: numpy vector of floats
    :param psd_n: noise, estimated noise frame PSD
    :type psd_n: numpy vector of floats
    :param psd_s: speech, estimated speech frame PSD
    :type psd_s: numpy vector of floats
    :param angular_r: angular noisy signal frame spectrum
    :type angular_r: numpy vector of floats
    :param an_win: windowing function
    :type an_win: numpy vector of floats

    :return: time domain enhanced signal frame (windowed)
    :rtype: numpy vector of floats

    """

    e = psd_s/np.maximum(psd_n,1.e-4)+1.e-6

    g = mag_r**2/np.maximum(psd_n,1.e-4)+1.e-6

    v = e*g/(1+e)

    aux = np.maximum(.5*v,1.e-2)

    if np.sum(np.isnan(aux)) > 0:

        print('nan found')

        raw_input()

    gain = np.sqrt(v)/(np.sqrt(np.pi)*g)*1/scipy.special.i0e(aux)

    if np.sum(np.isnan(gain)) > 0:

        print (v)

        print (np.sqrt(v))

        print (gain)

        print (scipy.special.iv(0,aux))

        print ('nan found')

        raw_input()

    # plt.plot(gain)
    #
    # plt.show()

    mag_enh = gain*mag_r

    enh_fs = synth(mag_enh,angular_r,an_win)

    return enh_fs #, psd_enh

def enhance2(psd_r, psd_n, mt, angular_r, an_win, ro=0.05): # note that instead of psd_s we pass the masking treshold (mt)

    """
    Perceptual Wiener filter

    :param psd_r: PSD of the noisy signal frame
    :type psd_r: numpy vector of floats
    :param psd_n: PSD of noise (may be approximate, smoothed etc.)
    :type psd_n: numpy vector of floats
    :param mt: PSD masking threshold
    :type mt: numpy vector of floats
    :param angular_r: angular noisy signal frame spectrum
    :type angular_r: numpy vector of floats
    :param an_win: windowing function
    :type an_win: numpy vector of floats
    :param ro: (0.-1.) larger causes less signal modification
    :type ro: float

    :return: time domain enhanced signal frame (windowed)
    :rtype: numpy vector of floats

    """

    gain = np.minimum(np.sqrt(mt/psd_n)+ro,1)

    psd_enh = gain*psd_r

    enh_fs = synth(psd_enh,angular_r,an_win)

    return enh_fs #, psd_enh

def load_wav(file_ : str, target_sampling_rate : int = 16000) -> Tuple[np.ndarray,int,int]:
    """Loads sound in a variety of formats supported by libsndfile. Resamples the input sound to target_sampling_rate.

    :param file_: sound file path
    :type file_: str
    :param target_sampling_rate: target sampling rate
    :type target_sampling_rate: int
    :returns: * the resampled sound to target sampling rate
              * the original sound sampling rate
              * the target sampling rate
    :rtype: Tuple[np.ndarray,int,int]
    """

    if not os.path.exists(file_):
        raise ValueError(f"The file {file_} does not exist.")

    s, sr = sf.read(file_, dtype=np.float32)

    print(f"Start resmpling")

    if int(sr) != int(target_sampling_rate):
    
        s_resampled = scipy.signal.resample(s,np.uint32(np.round(target_sampling_rate/sr*s.shape[0])))

    else:

        s_resampled = s

    print(f"Resampling done")

    np.random.seed(1000) # causes the dither is same on each run

    s_resampled += np.random.randn(*s_resampled.shape)*1.e-6 # add dither to improve numerical behaviour

    return s_resampled, int(sr), target_sampling_rate

def save_wav(signal : np.ndarray, sampling_rate : int, file_name : str):

    if not os.path.exists(os.path.dirname(file_name)):
        try:
            os.makedirs(os.path.dirname(file_name))
        except OSError:
            print(f'Can not create the directory {os.path.dirname(file_name)}')
        else:
            print(f'Path {os.path.dirname(file_name)} succesfully created.')

    sf.write(file_name, signal, sampling_rate)



if __name__ == "__main__":

    pass

# -*- coding: utf-8 -*_ fbe.py ========================================================================= FBE module provides several utilities and signal parametrization methods. Class for computing signal windowing function with sin(x)^2 and cos(x)^2 tails.      :param frame: the frame length in samples     :type frame: int     :param overlap: the size of the overlaping part of the window (the length of the tails on both sides)     :type overlap: int      :return: nothing Method returning the vector of window's values.          :return: the window          :rtype: numpy array of length frame Versatile class computing various speech signal representations, mostly based on AR modelling and Mel Frequency     Filterbanks.       :param frame_zero_adding: required length of the sequence after zero adding operation, defaults to None, which indicates no zero adding     :type frame_zero_adding: int      :param frame: frame length in samples     :type frame: int      :param sr: sampling frequency in Hz     :type sr: float      :param preem_alfa: the preemphasis coefficient     :type preem_alfa: float      :param freq_range: frequency range in which the mel frequency filterbanks should be computed     :type freq_range: np.ndarray two elemements vector of floats      :param filts_num: number of mel frequency triangular filters in the filterbank     :type filts_num: int      :param window: the windowing function     :type window: np.ndarray, numpy vector of floats, defaults to None, which causes using of rectangular window      :param ar_order: the AR model order     :type ar_order: int      :param cepstral_lifter: the cepstral lifter in MFCC computation     :type cepstral_lifter: int      :param num_ceps: number of cepstra     :type num_ceps: int      :returns: nothing      .. note:: PSD is abbreviation for power spectral density in the whole documentation. AR is abbreviation for                 autoregressive in the whole documentation. # sampling frequency in Hz # number of samples in the frame # fft length, sets the self._nfft atribute # preemphasis coefficient # frequency range in Hz # number of triangular filterbank channels # length of the unique part of the FFT # matrices # compute the mel-frequency triangular filterbank, sets the H atribute # window # compute cepstral lifter # dct matrix Resets the cache. Returns the frame length in samples          :return: the frame length          :rtype: int Gets the triangular mel frequency filterbank.          :return: the filter matrix containing in each row a single filter          :rtype: np.ndarray, numpy array with filts_num rows Gets the weighting matrix, which is a square of the product of pseudo inverses of the Jacobian of the linear         magnitude spectrum filter banks transform.          :return: the weighting matrix          :rtype: numpy array with dimension filts_num x filts_num Gets pseudo inverse of the weighting matrix.          :returns: the pseudo inverse of the weighting matrix          :rtype: np.ndarray, numpy array with dimension filts_num x filts_num Gets the pseudoinverse of the filterbanks matrix.          :returns: the pseudo inverse of the weighting matrix          :rtype: np.ndarray, numpy array with dimension filts_num x filts_num Gets the signal windowing function.          :returns: the windowing function          :rtype: np.ndarray, numpy array with dimension 1 x frame Computes the mel frequency triangular filterbank. # filter cutoff frequencies (Hz) for all filters, size 1x(M+2) Projecting matrix A onto space of positive definite matrices.          :param A: matrix to be projected         :type A: np.ndarray         :param rcond: reciprocal condition number of the resulting matrix         :type rcond: float          :return: the projected matrix         :rtype: np.ndarray The weighting matrix. The inverse of the weighting matrix The pseudoinverse of the mel frequency triangular filter banks matrix The next nearest power of 2. Hertz to mel frequency scale Mel to frequency scale Converts frequency index to frequency in Hz          :param i: frequency index         :type i: int          :return: frequency in Hz         :rtype: float Converts frequency in Hz to the frequency index          :param f: frequency in Herz         :type f: float          :return: frequency index         :rtype: int Sets the frame of the signal - this is then used to compute the all signal representations          :param fr: signal frame         :type fr: np.ndarray, numpy vector of floats          :return: nothing Set compact spectrum          :param wgh: the compact specturm with filt_num elements         :type wgh: numpy vector of floats          :returns: nothing Set AR compact spectrum          :param arwgh: the compact autoregresive specturm with filt_num elements         :type arwgh: np.ndarray, numpy vector of floats          :returns: nothing Set filterbank energies          :param fbe: the filter bank energies (vector with filt_num elements)         :type fbe: np.ndarray, numpy vector of floats          :returns: nothing Set magnitude spectrum          :param mag: the magnitude spectrum         :type mag: np.ndarray, numpy vector of floats          :returns: nothing Set power density spectrum          :param psd: the power density spectrum         :type psd: np.ndarray, numpy vector of floats          :returns: nothing Gets frame          :returns: the frame         :rtype: np.ndarray, numpy vector of floats Gets windowed frame          :returns: the windowed frame         :rtype: np.ndarray, numpy vector of floats Gets FFT          :returns: the fft of the, possibly zero added, signal frame         :rtype: np.ndarray, numpy vector of complex floats Gets magnitude spectrum          :returns: the magnitude of the, possibly zero added, signal frame         :rtype: np.ndarray, numpy vector of floats Gets angular spectrum.          :returns: the angular spectrum of the, possibly zero added, signal frame         :rtype: np.ndarray, numpy vector of floats Gets power density spectrum          :returns: the PSD of the, possibly zero added, signal frame         :rtype: np.ndarray, numpy vector of floats Gets LPC coefficients aka STP coefficients or AR coefficients          :return: LPC with the leading 1         :rtype: np.ndarray Gets variance of the short term residual spectrum          :return: short term residual variance         :rtype: np.ndarray Setting AR coefficients and STP residual variance          :param a: AR coefficients with leading one         :param var: variance of the short term residual Sets the AR coefficients Gets AR magnitude spectrum          :return: AR magnitude spectrum         :rtype: np.ndarray, numpy vector of floats of length _nfft/2+1 Gets AR power density spectrum          :return: the AR PSD         :rtype: np.ndarray Gets filter banks outputs based on magnitude spectrum          :return: filter bank filtered magnitude spectrum         :rtype: np.ndarray, numpy vector of floats of length filt_num Converts smoothed filter banks energies to MFCC coefficients          :return: MFCC coefficients         :rtype: np.ndarray, numpy vector of floats (size num_cep) #scipy.fftpack.dct(logfbe,n=self.num_ceps,norm='ortho') # liftering Converts filter banks energies to MFCC coefficients          :return: MFCC coefficients         :rtype: np.ndarray, numpy vector of floats (size num_cep) #scipy.fftpack.dct(logfbe,n=self.num_ceps,norm='ortho') # here comes liftering AR magnitude spectrum to filter banks energies          :return: filter bank filtered AR magnitude spectrum         :rtype: np.ndarray, numpy vector of floats (size num_filt) Weighted filter bank energies          :return: the magnitude compact spectrum         :rtype: np.ndarray, numpy vector of floats (size num_filt) AR weighted filter bank energies          :return: the AR magnitude compact spectrum         :rtype: np.ndarray, numpy vector of floats (size num_filt) Smoothed magnitude spectrum          :return: magnitude spectrum computed from filter bank energies         :rtype: np.ndarray, numpy vector of floats (size _nfft/2+1) Smoothed power density spectrum          :return: PSD computed from filter bank energies         :rtype: np.ndarray, numpy vector of floats(size _nfft/2+1) Smoothed AR magnitude spectrum          :return: smoothed (from arfbe) AR magnitude spectrum (size _nfft/2+1)         :rtype: np.ndarray, numpy vector of floats Smoothed AR PSD          :return: smoothed (from arfbe) AR PSD (size _nfft/2+1)         :rtype: np.ndarray, numpy vector of floats Perform preemphasis on the input signal.          :param signal: The signal to filter.         :type signal: np.ndarray, numpy vector of floats         :param coeff: The preemphasis coefficient. 0 is no filter, default is 0.95.         :type coeff: float         :returns: the filtered signal.         :rtype: numpy vector of floats Smoothed filter bank energies          :return: smoothed, from compact spectrum, filter bank energies         :rtype: np.ndarray, numpy vector of floats (size num_filt) Smoothed AR filter bank energies          :return smoothed, from compact AR spectrum, filter bank energies         :rtype: np.ndarray, numpy vector of floats (size num_filt) Smoothed magnitude spectrum          :return: computed from compact spectum magnitude spectrum         :rtype: np.ndarray, numpy vector of floats (size _nfft/2+1) Smoothed AR magnitude spectrum          :return: computed from AR compact spectrum magnitude spectrum         :rtype: np.ndarray, numpy vector of floats (size _nfft/2+1) Smoothed PSD          :return: computed from compact spectrum PSD         :rtype: np.ndarray, numpy vector of floats (size _nfft/2+1) Smoothed AR PSD          :return: PSD computed from AR compact spectra         :rtype: np.ndarray, numpy vector of floats (size _nfft/2+1) PSD -> weighted compact spectrum          :param psd: the PSD         :type psd: np.ndarray, numpy vector of floats (size _nfft/2+1)          :return: compact spectrum based on PSD         :rtype: np.ndarray, numpy vector of floats (size num_filt) Converting PSD into LPC coefficients and excitation variance          :param psd: left half of the PSD         :type psd: numpy vector of floats (size _nfft/2+1)         :param LpcOrder: AR model order         :type LpcOrder: int          :return: * (`vector of floats`) direct form AR coeff. with leading 1                  * (`float`) the variance of the short term residual Signal synthesis based on magnitude and angular spectra      :param mag_enh: enhanced speech magnitude spectrum     :type psd_enh: np.ndarray, numpy vector of floats     :param angular_r: angular noisy signal frame spectrum     :type angular_r: np.ndarray, numpy vector of floats     :param an_win: windowing function     :type an_win: np.ndarray numpy vector of floats      :return: time domain enhanced signal frame (windowed)     :rtype: numpy vector of floats # X = np.sqrt( psd_enh ) The Wiener filter in frequency domain      :param mag_r: noisy, unprocessed signal frame magnitude spectrum     :type psd_r: numpy vector of floats     :param psd_n: noise, estimated noise frame PSD     :type psd_n: numpy vector of floats     :param psd_s: speech, estimated speech frame PSD     :type psd_s: numpy vector of floats      :return: enhanced speech PSD     :rtype: numpy vector of floats # filter out the most high frequency peak The Wiener filter returning the time frequency signal frame      :param mag_r: noisy, unprocessed signal frame magnitude spectrum     :type mag_r: numpy vector of floats     :param psd_n: noise, estimated noise frame PSD     :type psd_n: numpy vector of floats     :param psd_s: speech, estimated speech frame PSD     :type psd_s: numpy vector of floats     :param angular_r: angular noisy signal frame spectrum     :type angular_r: numpy vector of floats     :param an_win: windowing function     :type an_win: numpy vector of floats      :return: time domain enhanced signal frame (windowed)     :rtype: numpy vector of floats #, psd_enh Perceptual enhancement (SETAP p. 245)      :param psd_r: noisy, unprocessed signal frame PSD     :type psd_r: numpy vector of floats     :param psd_n: noise, estimated noise frame PSD     :type psd_n: numpy vector of floats     :param psd_s: speech, estimated speech frame PSD     :type psd_s: numpy vector of floats     :param angular_r: angular noisy signal frame spectrum     :type angular_r: numpy vector of floats     :param an_win: windowing function     :type an_win: numpy vector of floats      :return: time domain enhanced signal frame (windowed)     :rtype: numpy vector of floats # plt.plot(gain) # # plt.show() #, psd_enh # note that instead of psd_s we pass the masking treshold (mt) Perceptual Wiener filter      :param psd_r: PSD of the noisy signal frame     :type psd_r: numpy vector of floats     :param psd_n: PSD of noise (may be approximate, smoothed etc.)     :type psd_n: numpy vector of floats     :param mt: PSD masking threshold     :type mt: numpy vector of floats     :param angular_r: angular noisy signal frame spectrum     :type angular_r: numpy vector of floats     :param an_win: windowing function     :type an_win: numpy vector of floats     :param ro: (0.-1.) larger causes less signal modification     :type ro: float      :return: time domain enhanced signal frame (windowed)     :rtype: numpy vector of floats #, psd_enh Loads sound in a variety of formats supported by libsndfile. Resamples the input sound to target_sampling_rate.      :param file_: sound file path     :type file_: str     :param target_sampling_rate: target sampling rate     :type target_sampling_rate: int     :returns: * the resampled sound to target sampling rate               * the original sound sampling rate               * the target sampling rate     :rtype: Tuple[np.ndarray,int,int] # causes the dither is same on each run # add dither to improve numerical behaviour