# Training Procedure for MRNN
# Author: Raymond Xia (yangyanx@andrew.cmu.edu)

# MRNN core and configuration
from config import *
import theano
from mrnn import MRNNTheano, MRNNEngine
from rats_util import SampleGenerator, MRNNLogger

# Audio utilities
from batch_io import AudioTransformer
from util import *
from spectrogram import *

# Other general utilities
import argparse
import subprocess
import sys
import os
import time
import numpy as np
import itertools
from datetime import datetime

# Configurations
cfg_rats  = config['rats']
floatX = theano.config.floatX # config to float32 for GPU processing
flatten = lambda l: [item for sublist in l for item in sublist]
# for getting magnitude spectrogram from spectrogram
if cfg_rats['ZERO_PHASE_STFT']:
    magspec = lambda x: magphase(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])[0].astype(floatX)
else:
    magspec = lambda x: magphase(stft_old(x,
            int(cfg_rats['WINDOW_LENGTH']*cfg_rats['SAMPLE_RATE']),
            cfg_rats['HOP_FRACTION'],
            truncate=True))[0].astype(floatX)

def train_rats_2015(outdir,start_with=None,wiener=False):
    # Train MRNN using LDC2015S02 data

    # Create a MRNN logger
    logger = MRNNLogger(outdir)
    print "Logger created."

    # build MRNN engine
    print "Start building MRNN."
    #engine =None
    engine = MRNNEngine(cfg_rats['FREQ_DIM'],cfg_rats['HIDDEN_DIM'],
                        cfg_rats['BPTT_TRUNCATE'], param_path=start_with,
                        wiener=wiener)
    print "MRNN successfully built."

    # collect features and do training
    epoch_size = cfg_rats['NEPOCH']
    minibatch_size = cfg_rats['MINIBATCH_SIZE'] # number of audio segments to be used each batch
    example_size = cfg_rats['EXAMPLE_SIZE']
    type_ratio = cfg_rats['TYPE_RATIO']

    # Learning starts here
    for epoch in xrange(epoch_size): # do 1 epoch
        #print "Epoch %d: Loading data from %s ..." % (epoch,cfg_rats['RATS_A'])
        # Each epoch will have a different order
        TRAINDIR = [cfg_rats['RATS_A'],cfg_rats['RATS_D'],cfg_rats['RATS_H']]
        TRAIN = AudioTransformer(TRAINDIR,
                                cfg_rats['SAMPLE_RATE'], mono=True,
                                random_order=True, verbose=True)
        ll = [] # contains loss for each epoch
        #"""
        for audio_batch in grouper(TRAIN,10): # process 10 audio files at a time
            #xSG = SampleGenerator(x,xpath,minibatch_size,type_ratio,maxbatch=1)
            xSGs = map(lambda tup: SampleGenerator(tup[0],tup[1],example_size,
                                    type_ratio=cfg_rats['TYPE_RATIO'],maxbatch=1),audio_batch)
            for data_batch in itertools.izip(*xSGs): # grouped in batch already
                MAG_Y,MAG_X,PHZ_X = map(np.concatenate, zip(*data_batch))
                assert MAG_Y.shape == MAG_X.shape == PHZ_X.shape
                if wiener:
                    train_loss = engine.train_with_sgd(MAG_X,MAG_Y,P_train=PHZ_X,
                                      learning_rate=cfg_rats['LEARNING_RATE'],
                                      decay=cfg_rats['RMSPROP_DECAY'])
                else:
                    train_loss = engine.train_with_sgd(MAG_X,MAG_Y,
                                      learning_rate=cfg_rats['LEARNING_RATE'],
                                      decay=cfg_rats['RMSPROP_DECAY'])
                ll.append(train_loss)
                # print out loss after one pass
                dt = datetime.now().isoformat()
                print("\n%s (Epoch:%d, Examples:%d)" % (dt,epoch,engine.num_seen))
                print("--------------------------------------------------")
                print("Loss: %f" % train_loss)
                sys.stdout.flush()

        # Save parameters after one entire pass of all batches
        dt = datetime.now().isoformat()
        ts = datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
        logger.write_epoch_loss(ll,engine.num_seen,ts)
        MODEL_OUTPUT_FILE = "MRNN-%s-%s-%s-%s.dat" % (ts, engine.feat_dim, engine.hidden_dim,engine.num_seen)
        print("Saving parameters to %s..." % MODEL_OUTPUT_FILE)
        engine.save_params(os.path.join(outdir,MODEL_OUTPUT_FILE))

# An iterator grouper for chunk processing
def grouper(iterable,n):
    it = iter(iterable)
    while True:
       chunk = tuple(itertools.islice(it, n))
       if not chunk:
           return
       yield chunk


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('-m',
        help='Specify initial model path', required=False)
    parser.add_argument('-o',
        help='Specify model saving path', required=True)
    args = parser.parse_args()

    start_with = None
    if args.m: start_with = args.m
    train_rats_2015(args.o,start_with=start_with,wiener=cfg_rats['WIENER'])
