from config import *
from util import *
from spectrogram import *
import theano
from batch_io import AudioTransformer
import argparse
from train import *

#theano.config.compute_test_value = 'warn'
#theano.config.optimizer='None'
cfg_rats = config['rats']
floatX = theano.config.floatX
def mrnn_denoise(engine, noisy_path, out_path=None, wiener=True,analysis_mode=False):

    if not analysis_mode:
        noisy_dataset = AudioTransformer(noisy_path,cfg_rats['SAMPLE_RATE'],
                                        verbose=True)
        for x,_ in noisy_dataset:
            if cfg_rats['ZERO_PHASE_STFT']:
                X = stft(x, cfg_rats['SAMPLE_RATE'],
                        window_length=cfg_rats['WINDOW_LENGTH'],
                        hop_fraction=cfg_rats['HOP_FRACTION'],
                        nfft=cfg_rats['FFT_SIZE'],
                        stft_truncate=True)[-1]
            else:
                X = stft_old(x,
                        int(cfg_rats['WINDOW_LENGTH']*cfg_rats['SAMPLE_RATE']),
                        cfg_rats['HOP_FRACTION'],
                        truncate=True)
            M,P = magphase(X)
            if theano.config.floatX == 'float32': # for GPU
                M = M.astype('float32')
                P = P.astype('complex64')
            elif theano.config.floatX == 'float64': # default
                pass
            else:
                raise ValueError("Use float32/float64 for high quality reconstruction.")

            if wiener: # enable wiener filter in MRNN
                M_de = engine.denoise(M,p_noisy=np.angle(P))
            else:
                M_de = engine.denoise(M)
            xr = istft(M_de*P,cfg_rats['SAMPLE_RATE'],len(x),
                        window_length=cfg_rats['WINDOW_LENGTH'],
                        hop_fraction=cfg_rats['HOP_FRACTION'],
                        stft_truncate=True)
            if out_path != None: noisy_dataset.write_to_file(xr,out_path)

    # output all MRNN hidden parameters
    # also assume single file input
    else:
        assert os.path.isfile(noisy_path)
        assert cfg_rats['ZERO_PHASE_STFT']
        x = list(AudioTransformer(noisy_path,cfg_rats['SAMPLE_RATE'],
                                        verbose=True))[0][0]
        X = stft(x, cfg_rats['SAMPLE_RATE'],
                window_length=cfg_rats['WINDOW_LENGTH'],
                hop_fraction=cfg_rats['HOP_FRACTION'],
                nfft=cfg_rats['FFT_SIZE'],
                stft_truncate=True)[-1]
        M,P = magphase(X)
        if theano.config.floatX == 'float32': # for GPU
            M = M.astype('float32')
            P = P.astype('complex64')
        elif theano.config.floatX == 'float64': # default
            pass
        else:
            raise ValueError("Use float32/float64 for high quality reconstruction.")
        M_clean,N,G1,G2,G3 = engine.denoise_output_all(M)
        xr = istft(M_clean*P,cfg_rats['SAMPLE_RATE'],
                    window_length=cfg_rats['WINDOW_LENGTH'],
                    hop_fraction=cfg_rats['HOP_FRACTION'],
                    stft_truncate=True)
        if out_path != None: noisy_dataset.write_to_file(xr,out_path)
        return M,M_clean,N,G1,G2,G3




if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-m',
        help='Specify model path', required=True)
    parser.add_argument('-i',
        help='Input path (could be file/directory)', required=True)
    parser.add_argument('-o',
        help='Output path', required=True)
    parser.add_argument('-w',
        help='Enable Wiener filtering.', required=False,action='store_true',
        default=False)
    args = parser.parse_args()

    engine = MRNNEngine(cfg_rats['FREQ_DIM'],cfg_rats['HIDDEN_DIM'],
                        cfg_rats['BPTT_TRUNCATE'],param_path=args.m,train=0,
                        wiener=args.w)
    mrnn_denoise(engine,args.i,out_path=args.o,wiener=args.w)
