# Magnitude Recurrent Neural Networks for Speech denoising
# Author: Raymond Xia (yangyanx@andrew.cmu.edu)
# Based on paper by Osako, Raj, Singh
import numpy as np
import theano as theano
import theano.tensor as T
from theano.gradient import grad_clip
import time
import operator
from dropout import Dropout
from theano.ifelse import ifelse
from theano.tensor.slinalg import solve
from theano.tensor import fft as Tfft

CONFIG_WIENER = {
    # STFT Parameters
    'SAMPLE_RATE': 16000, # in Hz
    'WINDOW_LENGTH': 0.032, # in seconds
    'HOP_FRACTION': 0.25,
    'FFT_SIZE': 512,
    # LPC Parameters
    'LPC_ORDER': 12,
}

floatX = theano.config.floatX
complexX = 'complex64' if floatX=='float32' else 'complex128'
WIENER_EXP_MAT = \
    T.constant(np.exp(-1j*(2*np.pi)/CONFIG_WIENER['FFT_SIZE']*np.outer(np.arange(
    CONFIG_WIENER['LPC_ORDER']+1),np.arange(CONFIG_WIENER['FFT_SIZE']/2+1))).astype(complexX))
class MRNNTheano:

    def __init__(self, freq_dim, hidden_dim=128, look_back=1,
                 bptt_truncate=-1, params=None, train=1, wiener=True):
        # Assign instance variables
        self.freq_dim = freq_dim
        self.hidden_dim = hidden_dim
        self.look_back = look_back
        self.bptt_truncate = bptt_truncate
        self.activate = T.nnet.sigmoid #T.nnet.relu
        self.dropout = Dropout(train)
        self.wiener = wiener

        # Initialize the network parameters for G1, G2, G3
        # Input-Hidden layer weights
        # 1 each for G2,G3; 5 for G1 as it looks back 5 frames
        if(params is not None):
            self.__check_params__(params)
            self.E = theano.shared(params['E'].astype(floatX), name='E', borrow=True)
            self.V = theano.shared(params['V'].astype(floatX), name='V', borrow=True)
            self.bE = theano.shared(params['bE'].astype(floatX), name='bE', borrow=True)
            self.bV = theano.shared(params['bV'].astype(floatX), name='bV', borrow=True)
            print '... restoring model weights/bias'
            print 'E: ', params['E'].shape
            print 'V: ', params['V'].shape
            print 'bE: ', params['bE'].shape
            print 'bV: ', params['bV'].shape
        else:
            # default values
            E = np.random.uniform(-np.sqrt(1./freq_dim), np.sqrt(1./freq_dim),
                (self.look_back+2,hidden_dim,4*freq_dim))
            # Hidden-Output layer weights
            V = np.random.uniform(-np.sqrt(1./hidden_dim), np.sqrt(1./hidden_dim),
                (self.look_back+2,freq_dim,hidden_dim))
            # Input-Hidden layer bias
            bE = np.zeros((self.look_back+2, hidden_dim))
            # Hidden-Output layer bias
            bV = np.zeros((self.look_back+2, freq_dim))
            # Theano: Created shared variables
            self.E = theano.shared(name='E',value=E.astype(floatX))
            self.V = theano.shared(name='V',value=V.astype(floatX))
            self.bE = theano.shared(name='bE',value=bE.astype(floatX))
            self.bV = theano.shared(name='bV',value=bV.astype(floatX))
            print '... creating model weights/bias'
            print 'E: ', E.shape
            print 'V: ', V.shape
            print 'bE: ', bE.shape
            print 'bV: ', bV.shape

        # SGD / rmsprop: Initialize parameters
        self.mE = theano.shared(name='mE',
                    value=np.zeros((self.look_back+2,hidden_dim,4*freq_dim)).astype(floatX))
        self.mV = theano.shared(name='mV',
                    value=np.zeros((self.look_back+2,freq_dim,hidden_dim)).astype(floatX))
        self.mbE = theano.shared(name='mbE',
                    value=np.zeros((self.look_back+2, hidden_dim)).astype(floatX))
        self.mbV = theano.shared(name='mbV',
                    value=np.zeros((self.look_back+2, freq_dim)).astype(floatX))

        # We store the Theano graph here
        self.theano = {}
        self.__theano_build__()

    def __check_params__(self,params):
        assert params['E'].shape[0] == self.look_back+2
        assert params['E'].shape[1] == self.hidden_dim
        assert params['E'].shape[2] == self.freq_dim * 4
        assert params['V'].shape[0] == self.look_back+2
        assert params['V'].shape[1] == self.freq_dim
        assert params['V'].shape[2] == self.hidden_dim
        assert params['bE'].shape[0] == self.look_back+2
        assert params['bE'].shape[1] == self.hidden_dim
        assert params['bV'].shape[0] == self.look_back+2
        assert params['bV'].shape[1] == self.freq_dim

    def get_params(self):
        params = {'E':self.E.get_value(),
                  'V':self.V.get_value(),
                  'bE':self.bE.get_value(),
                  'bV':self.bV.get_value()}
        return params


    def __theano_build__(self):

        # Assume tensor here for minibatch processing
        # Dimension: instance, time, frequency
        X = T.tensor3('Magnitude spec',dtype=floatX) # input magnitude STFT
        if self.wiener:
            P = T.tensor3('Phase spec',dtype=floatX)
        target = T.tensor3('target',dtype=floatX) # target magnitude STFT
        g2_init = theano.shared(
                np.zeros(self.freq_dim,dtype=floatX)+.9)
        g3_init =theano.shared(
                np.zeros(self.freq_dim,dtype=floatX)+.1)

        def forward_prop_5(x_t,
            y_tm1,y_tm2,y_tm3,y_tm4,y_tm5,
            n_tm1, n_tm2,n_tm3,n_tm4,n_tm5,
            g1_tm1,g1_tm2,g1_tm3,g1_tm4,g1_tm5,
            g2_tm1,g3_tm1):

            # Activation to hidden layer in G11-G15
            zg11 = self.activate(self.E[0].dot(T.concatenate(
                [g1_tm1, y_tm1, n_tm1, x_t])) + self.bE[0])
            g11_t = self.activate(self.V[0].dot(zg11) + self.bV[0])

            zg12 = self.activate(self.E[1].dot(T.concatenate(
                [g1_tm2, y_tm2, n_tm2, x_t])) + self.bE[1])
            g12_t = self.activate(self.V[1].dot(zg12) + self.bV[1])

            zg13 = self.activate(self.E[2].dot(T.concatenate(
                [g1_tm3, y_tm3, n_tm3, x_t])) + self.bE[2])
            g13_t = self.activate(self.V[2].dot(zg13) + self.bV[2])

            zg14 = self.activate(self.E[3].dot(T.concatenate(
                [g1_tm4, y_tm4, n_tm4, x_t])) + self.bE[3])
            g14_t = self.activate(self.V[3].dot(zg14) + self.bV[3])

            zg15 = self.activate(self.E[4].dot(T.concatenate(
                [g1_tm5, y_tm5, n_tm5, x_t])) + self.bE[4])
            g15_t = self.activate(self.V[4].dot(zg15) + self.bV[4])




            # Activation to hidden layer in G2
            zg2 = self.activate(self.E[5].dot(T.concatenate(
                [g2_tm1, y_tm1, n_tm1, x_t])) + self.bE[5])
            g2_t = self.activate(self.V[5].dot(zg2) + self.bV[5])

            n_t = (n_tm1*g11_t+n_tm2*g12_t+n_tm3*g13_t+
                n_tm4*g14_t+n_tm5*g15_t)/5. + x_t*g2_t

            # Activation to hidden layer in G3
            zg3 = self.activate(self.E[6].dot(T.concatenate(
                [g3_tm1, y_tm1, n_t, x_t])) + self.bE[6])
            g3_t = self.activate(self.V[6].dot(zg3) + self.bV[6])

            # Final output calculation
            y_t = x_t * g3_t

            return [y_t, n_t, g11_t, g2_t, g3_t]

        def forward_prop_1(x_t,y_tm1,n_tm1,g1_tm1,g2_tm1,g3_tm1):

            E = self.dropout.dropout_layer(self.E, .2)
            V = self.dropout.dropout_layer(self.V, .5)
            # Activation to hidden layer in G11
            zg11 = self.activate(E[0].dot(T.concatenate(
                [g1_tm1, y_tm1, n_tm1, x_t])) + self.bE[0])
            g11_t = T.nnet.sigmoid(V[0].dot(zg11) + self.bV[0])

            # Activation to hidden layer in G2
            zg2 = self.activate(E[1].dot(T.concatenate(
                [g2_tm1, y_tm1, n_tm1, x_t])) + self.bE[1])
            g2_t = T.nnet.sigmoid(V[1].dot(zg2) + self.bV[1])

            n_t = n_tm1*g11_t + x_t*g2_t

            # Activation to hidden layer in G3
            zg3 = self.activate(E[2].dot(T.concatenate(
                [g3_tm1, y_tm1, n_t, x_t])) + self.bE[2])
            g3_t = T.nnet.sigmoid(V[2].dot(zg3) + self.bV[2])

            # Final output calculation
            y_t = x_t * g3_t

            return [y_t, n_t, g11_t, g2_t, g3_t]

        # Estimate speech spectral envelope by LPC
        def lpc_env(Xm,Xp,Pnn):
            X_r = Xm*T.cos(Xp)
            X_i = Xm*T.sin(Xp)
            X = T.stack([X_r,X_i])
            x = T.roll(Tfft.irfft(X.dimshuffle(('x',1,0))),256,axis=-1)

            # Calculate autocorrelation by DFT
            xpad_fft = Tfft.rfft(T.concatenate([x,T.zeros_like(x)],axis=1))[0].dimshuffle((1,0))
            # time reversal corresponds to frequency conjugation for real signals
            x_xcorr_r = xpad_fft[0]**2+xpad_fft[1]**2 # real part
            x_xcorr_i = T.zeros_like(xpad_fft[1]) # imaginary part
            x_xcorr = Tfft.irfft(
                T.stack([x_xcorr_r,x_xcorr_i]).dimshuffle(('x',1,0)))[0,:CONFIG_WIENER['LPC_ORDER']+1]

            # Form toeplitz matrix
            toeplitz_init = T.eye(CONFIG_WIENER['LPC_ORDER'],m=CONFIG_WIENER['LPC_ORDER'])*x_xcorr[0]
            toeplitz_mat, _ = theano.scan(fn=lambda i,ps: \
                ps+T.eye(CONFIG_WIENER['LPC_ORDER'],m=CONFIG_WIENER['LPC_ORDER'],k=i+1)*x_xcorr[i+1]+\
                T.eye(CONFIG_WIENER['LPC_ORDER'],m=CONFIG_WIENER['LPC_ORDER'],k=-(i+1))*x_xcorr[i+1],
                 outputs_info=toeplitz_init,
                 sequences=T.arange(CONFIG_WIENER['LPC_ORDER']-1))
            toeplitz_mat = toeplitz_mat[-1]

            # And solve for LPC coefficients
            alpha = solve(toeplitz_mat,x_xcorr[1:])
            lpc_coeffs = T.concatenate([[1.],-alpha]).dimshuffle(('x',0))

            # Finally, estimate speech PSD
            Pxx = Xm**2
            Pss_est = T.cast(1 / T.abs_(lpc_coeffs.dot(WIENER_EXP_MAT))**2,floatX)
            gain = T.mean(Pxx-Pnn) / T.mean(Pss_est)
            gain = ifelse(T.gt(gain, 1e-16), gain, 1e-16)
            # use power-normalized gain factor and for MRNN
            #gain = ifelse(T.gt(gain, 1e-16), 1/T.sum(Pss_est), 1e-16)
            Pss = gain * Pss_est
            #H = Pss / (Pss+Pnn)
            #return theano.gradient.disconnected_grad(H.dimshuffle((1,)))
            return theano.gradient.disconnected_grad(Pss.dimshuffle((1,)))

        def forward_prop_1_wiener(x_t,p_t,y_tm1,n_tm1,g1_tm1,g3_tm1):

            E = self.dropout.dropout_layer(self.E, .2)
            V = self.dropout.dropout_layer(self.V, .5)
            # Activation to hidden layer in G11
            # This gate is intended to learn the noise PSD
            zg11 = self.activate(E[0].dot(T.concatenate(
                [g1_tm1, y_tm1, n_tm1, x_t])) + self.bE[0])
            # this is the noise mask
            g11_t = T.nnet.sigmoid(V[0].dot(zg11) + self.bV[0])

            n_t = n_tm1*g11_t # estimate noise magnitude
            Pnn = n_t**2
            Pss = lpc_env(x_t,p_t,Pnn)

            # Activation to hidden layer in G2
            #zg2 = self.activate(E[1].dot(T.concatenate(
            #    [g2_tm1, y_tm1, n_tm1, x_t])) + self.bE[1])
            #g2_t = T.nnet.sigmoid(V[1].dot(zg2) + self.bV[1])

            # Now this mask needs to be "regularized" by the LPC estimate of speech PSD
            #g2_t = g2_t * H

            #x_est = x_t * g2_t # apply LPC mask to get an estimated speech

            # Activation to hidden layer in G3
            zg3 = self.activate(E[2].dot(T.concatenate(
                [g3_tm1, y_tm1, Pnn, Pss])) + self.bE[2])
            g3_t = T.nnet.sigmoid(V[2].dot(zg3) + self.bV[2])

            # Final output calculation
            y_t = x_t * g3_t

            return [y_t, n_t, g11_t, g3_t]


        if self.look_back == 1:
            #n0 =
            y0 = theano.shared(
                    np.zeros(self.freq_dim,dtype=floatX))
            g1_init = theano.shared(
                    np.zeros(self.freq_dim,dtype=floatX)+.1)
            forward_prop_step = forward_prop_1
            taps = [-1]
        elif self.look_back == 5:
            n0 =  T.tile(T.sum(x[:5,:],axis=0,dtype=floatX) / 5.0,
                    (self.look_back,1))
            y0 = theano.shared(
                    np.zeros((self.look_back,self.freq_dim),dtype=floatX))
            g1_init = theano.shared(
                    np.zeros((self.look_back,self.freq_dim),dtype=floatX)+.1)
            forward_prop_step = forward_prop_5
            taps = [-1,-2,-3,-4,-5]
        else:
            raise ValueError('Forward propagation not implemented.')

        if self.wiener:
            assert self.look_back == 1 # don't want to do more complicated stuff than this
            forward_prop_step = forward_prop_1_wiener


        def denoise(x):
            """
            Denoise a single example drawn from a minibatch.
            x     : input spectrogram with dimension [T x F]
            """

            [y, n, g1, g2, g3], _ = theano.scan(
                forward_prop_step,
                sequences=x,
                truncate_gradient=self.bptt_truncate,
                outputs_info=[
                dict(initial=y0,taps=taps),
                dict(initial=T.sum(x[:5,:],axis=0,dtype=floatX) / 5.0,taps=taps),
                dict(initial=g1_init,taps=taps),
                dict(initial=g2_init),
                dict(initial=g3_init)])

            return y,n,g1,g2,g3

        def denoise_wiener(x,p):
            """
            Denoise a single example drawn from a minibatch.
            x     : magnitude spectrogram with dimension [T x F]
            p     : phase spectrogram with dimension [T x F]
            """

            [y, n, g1, g3], _ = theano.scan(
                forward_prop_step,
                sequences=[x,p],
                truncate_gradient=self.bptt_truncate,
                outputs_info=[
                dict(initial=y0,taps=taps),
                dict(initial=T.sum(x[:5,:],axis=0,dtype=floatX) / 5.0,taps=taps),
                dict(initial=g1_init,taps=taps),
                dict(initial=g3_init)])

            return y,n,g1,g3

        def get_cost(y, target):
            """
            Process a single example drawn from a minibatch.
            y     : denoised output
            target: target output
            """
            y_error = T.sum((y - target)**2)
            # Total cost (could add regularization here)
            cost = y_error

            return cost

        if self.wiener:
            # Pass phase into scan for reconstructing time-domain signal
            output_all, _ = theano.scan(lambda i, X,P: denoise_wiener(X[i],P[i]),
                                sequences=T.arange(X.shape[0]),
                                non_sequences=[X,P])
        else:
            output_all, _ = theano.scan(lambda i, X: denoise(X[i]),
                                sequences=T.arange(X.shape[0]),
                                non_sequences=[X])
        ys = output_all[0]

        costs, _ = theano.scan(lambda i, y,t: get_cost(y[i],t[i]),
                                   sequences=T.arange(ys.shape[0]),
                                   non_sequences=[ys,target])
        batch_cost = T.mean(costs)

        # Gradients from batch
        dE = T.grad(batch_cost, self.E)
        dbE = T.grad(batch_cost, self.bE)
        dV = T.grad(batch_cost, self.V)
        dbV = T.grad(batch_cost, self.bV)

        # SGD parameters
        learning_rate = T.scalar('learning_rate')
        decay = T.scalar('decay')

        # rmsprop cache updates
        mE = decay * self.mE + (1 - decay) * dE ** 2
        mV = decay * self.mV + (1 - decay) * dV ** 2
        mbE = decay * self.mbE + (1 - decay) * dbE ** 2
        mbV = decay * self.mbV + (1 - decay) * dbV ** 2

        # Assign functions
        if self.wiener:
            self.denoise = theano.function([X,P], ys)
            self.denoise_output_all = theano.function([X,P],output_all)
            self.ce_error = theano.function([X,P,target], batch_cost)
            self.bptt = theano.function([X,P,target], [dE, dbE, dV, dbV])
            self.sgd_step = theano.function(
                [X,P, target, learning_rate, theano.In(decay, value=0.9)],
                [batch_cost],
                updates=[(self.E, self.E - learning_rate * dE / T.sqrt(mE + 1e-6)),
                     (self.V, self.V - learning_rate * dV / T.sqrt(mV + 1e-6)),
                     (self.bE, self.bE - learning_rate * dbE / T.sqrt(mbE + 1e-6)),
                     (self.bV, self.bV - learning_rate * dbV / T.sqrt(mbV + 1e-6)),
                     (self.mE, mE),
                     (self.mV, mV),
                     (self.mbE, mbE),
                     (self.mbV, mbV)
                    ])
        else:
            self.denoise = theano.function([X], ys)
            self.denoise_output_all = theano.function([X],output_all)
            self.ce_error = theano.function([X, target], batch_cost)
            self.bptt = theano.function([X, target], [dE, dbE, dV, dbV])
            self.sgd_step = theano.function(
                [X, target, learning_rate, theano.In(decay, value=0.9)],
                [batch_cost],
                updates=[(self.E, self.E - learning_rate * dE / T.sqrt(mE + 1e-6)),
                     (self.V, self.V - learning_rate * dV / T.sqrt(mV + 1e-6)),
                     (self.bE, self.bE - learning_rate * dbE / T.sqrt(mbE + 1e-6)),
                     (self.bV, self.bV - learning_rate * dbV / T.sqrt(mbV + 1e-6)),
                     (self.mE, mE),
                     (self.mV, mV),
                     (self.mbE, mbE),
                     (self.mbV, mbV)
                    ])




class MRNNEngine(object):

    def __init__(self,feat_dim,hidden_dim,
                 bptt_truncate,train=1,param_path=None,wiener=True):
        self.feat_dim = feat_dim
        self.hidden_dim = hidden_dim
        self.bptt_truncate = bptt_truncate
        self.num_seen = 0 # remember number of examples passed
        self.wiener = wiener
        if param_path is not None: # load from file
            params = self.load_params(param_path)
            self.model = MRNNTheano(self.feat_dim,hidden_dim=self.hidden_dim,
                bptt_truncate=self.bptt_truncate,
                params=params,look_back=1,train=train,wiener=self.wiener)
        else:
            self.model = MRNNTheano(self.feat_dim,hidden_dim=self.hidden_dim,
                bptt_truncate=self.bptt_truncate,look_back=1,train=train,
                wiener=self.wiener)
    def train_with_sgd(self,X_train,T_train,P_train=None,
                        learning_rate=0.001,
                        decay=0.9):
        if P_train is None:
            assert not self.wiener
            assert X_train.shape == T_train.shape
            idx = np.random.permutation(X_train.shape[0]) # randomize feature order
            train_loss = self.model.sgd_step(X_train[[idx]],T_train[[idx]],learning_rate,decay)
            #train_loss = self.model.ce_error(X_train[[idx]],T_train[[idx]])
            self.num_seen += len(idx)
            return train_loss[0]
        else:
            assert self.wiener
            assert X_train.shape == T_train.shape == P_train.shape
            idx = np.random.permutation(X_train.shape[0]) # randomize feature order
            train_loss = self.model.sgd_step(X_train[[idx]],P_train[[idx]],
                                T_train[[idx]],learning_rate,decay)
            #train_loss = self.model.ce_error(X_train[[idx]],P_train[[idx]],T_train[[idx]])
            self.num_seen += len(idx)
            return train_loss[0]

    def denoise(self,x_noisy,p_noisy=None):
        # takes in tensor only
        if p_noisy is None:
            assert not self.wiener
            xr = self.model.denoise(np.array([x_noisy])) # takes in tensor only
        else:
            assert self.wiener
            xr = self.model.denoise(np.array([x_noisy]),np.array([p_noisy]))
        return xr[0]
    def denoise_output_all(self,x_noisy,p_noise=None):
        if p_noisy is None:
            assert not self.wiener
            xr,n,g1,g2,g3 = self.model.denoise_output_all(np.array([x_noisy]))
        else:
            assert self.wiener
            xr,n,g1,g2,g3 = self.model.denoise_output_all(np.array([x_noisy]),
                        np.array([p_noisy]))
        return xr[0],n[0],g1[0],g2[0],g3[0]

    def get_params(self):
        return self.model.get_params()

    def load_params(self, path):
        return np.load(path)

    def save_params(self, path):
        params = self.get_params()
        np.savez(path,E=params['E'],V=params['V'],\
            bE=params['bE'], bV=params['bV'])
