from __future__ import print_function
import os
import sys

from train.mimo_trainer import MimoTrainer
#from train.mimo_experiment import MimoExperiment

import keras
from keras.models import Sequential
from keras.layers import Dense, LSTM, Convolution1D, MaxPooling1D, BatchNormalization
from keras.layers.noise import GaussianNoise

from keras.optimizers import SGD, RMSprop
import keras.backend.tensorflow_backend as K
import tensorflow as tf

import numpy as np
import shutil

from utils import my_numpy_utils as my_np

from data_io.mimo_wfdb_adaptor import MimoWfdbAdaptor
from mimo_sequence.mimo_sequence import MimoSequence 

from ch2017_train.ch2017_graph_model import Ch2017GraphModel 


#
#  mimo_trainer: a class to train multiple-input multiple-output (MIMO) data using Keras and Tensorflow
#
block_seq_base_length=6000
out_dec_factor=18
#out_dec_factor=9000
class Ch2017MimoTrainer(MimoTrainer):
  def __init__(self, train_batch_size=None, block_seq_base_length=block_seq_base_length, out_dec_factor=out_dec_factor):
    super(Ch2017MimoTrainer, self).__init__(train_batch_size, block_seq_base_length, out_dec_factor)    
    self.mimo_wfdb_adaptor = MimoWfdbAdaptor()

#    self.experiment_tag = 'Conv-BN-Pool-LSTM-LSTM-LSTM-Dense'      

  #  self.experiment_tag = 'Conv-BN-Pool-LSTM-DO-LSTM-DO-LSTM-DO-Dense'      
    self.dropout=0.1 
    self.recurrent_dropout=0.1

    #self.experiment_tag = 'Conv-BN-Pool-LSTM-DO-BN-LSTM-DO-BN-LSTM-DO-BN-SGD-Dense_10xCrossValidation'   
    #self.experiment_tag = 'Conv-BN-Pool-LSTM50-DO-BN-LSTM50-DO-BN-LSTM50-DO-BN-Dense_10xCrossValidation'   
    self.experiment_tag = 'Conv-BN-Pool-LSTM100-DO-BN-LSTM100-DO-BN-LSTM100-DO-BN-Dense_10xCrossValidation'   
    
    self.train_proportion = 0.90
    #self.max_epochs = 1
    self.max_epochs = 400
    self.learning_rate = 1e-2

   

  def create_model(self, model_mimo, batch_size, max_length=None, device_id = '/cpu:0'):
    print("Keras version: " + keras.__version__)
    print("Tensorflow version: " + tf.__version__)
    
    graph_model = Ch2017GraphModel(model_mimo, batch_size, max_length=None, device_id = device_id)
    
    graph_model.create_model()
    return graph_model
    
        
    
#     print("Keras version: " + keras.__version__)
#     print("Tensorflow version: " + tf.__version__)
#     
#     
#     config = tf.ConfigProto()
#     #config.gpu_options.per_process_gpu_memory_fraction = 0.3
#     config.gpu_options.allow_growth=True
#     config.allow_soft_placement=True
#     config.log_device_placement=False
#     K.set_session(tf.Session(config=config))
#       
#     with tf.device(self.deviceId):
#     
#       print('Creating Model')
#       model = Sequential()
#       #model.add(Masking(mask_value=0., batch_input_shape=(batch_size, max_length, modelMimo.n_input_dim)))
#       
# #       model.add(GaussianNoise(self.input_gaussian_noise_sigma, batch_input_shape=(batch_size, max_length, modelMimo.n_input_dim)))
# #       model.add(BatchNormalization(momentum=0.99, epsilon=0.001, center=True, 
# #                                    scale=True, beta_initializer='zeros', gamma_initializer='ones', moving_mean_initializer='zeros', moving_variance_initializer='ones'))
#             
#       
#       model.add(Convolution1D(filters=10, kernel_size=10, padding='same', 
#                               activation='relu', batch_input_shape=(batch_size, max_length, modelMimo.n_input_dim)))
# #       model.add(Convolution1D(filters=10, kernel_size=10, padding='same', 
# #                               activation='relu'))
#        
#       # new for entry 2
#       model.add(BatchNormalization(momentum=0.99, epsilon=0.001, center=True, 
#                                    scale=True, beta_initializer='zeros', gamma_initializer='ones', moving_mean_initializer='zeros', moving_variance_initializer='ones'))
#             
#       #model.add(MaxPooling1D(pool_size=8))
#       model.add(MaxPooling1D(pool_size=18))
# #       model.add(Convolution1D(nb_filter=10, filter_length=10, border_mode='same', 
# #                               activation='relu',))
# #       model.add(MaxPooling1D(pool_length=2))
# #       model.add(Convolution1D(nb_filter=10, filter_length=10, border_mode='same', 
# #                               activation='relu'))
# #       model.add(MaxPooling1D(pool_length=2))
# 
#       #model.add(LSTM(self.lstm_n_cell[0], batch_input_shape=(batch_size, max_length, modelMimo.n_input_dim),
#       model.add(LSTM(self.lstm_n_cell[0],
#                       return_sequences=True,
#                       stateful=False,
#                       kernel_initializer='he_normal',
#                       activation='tanh',
#                       dropout = self.dropout, 
#                       recurrent_dropout = self.recurrent_dropout))
#       model.add(BatchNormalization(momentum=0.99, epsilon=0.001, center=True, 
#                                    scale=True, beta_initializer='zeros', gamma_initializer='ones', moving_mean_initializer='zeros', moving_variance_initializer='ones'))
#       model.add(LSTM(self.lstm_n_cell[1],
#                      return_sequences=True,
#                      stateful=False,
#                      kernel_initializer='he_normal',
#                       activation='tanh',
#                       dropout = self.dropout, 
#                       recurrent_dropout = self.recurrent_dropout))
#       model.add(BatchNormalization(momentum=0.99, epsilon=0.001, center=True, 
#                                    scale=True, beta_initializer='zeros', gamma_initializer='ones', moving_mean_initializer='zeros', moving_variance_initializer='ones'))
#       model.add(LSTM(self.lstm_n_cell[2],
#                      return_sequences=True,
#                      stateful=False,
#                      kernel_initializer='he_normal',
#                       activation='tanh',
#                       dropout = self.dropout, 
#                       recurrent_dropout = self.recurrent_dropout))
#       model.add(BatchNormalization(momentum=0.99, epsilon=0.001, center=True, 
#                                    scale=True, beta_initializer='zeros', gamma_initializer='ones', moving_mean_initializer='zeros', moving_variance_initializer='ones'))
#        
# #       model.add(Dense(modelMimo.n_labels+1, kernel_initializer='he_normal',
# #                                       activation='softmax')) 
#       model.add(Dense(modelMimo.n_labels, kernel_initializer='he_normal',
#                                       activation='softmax')) 
#       #model.add(GlobalAveragePooling1D())
#       #model.add(Dropout(0.5))
#       
#       print (model.summary())
#       print('dropout: {:5.0f}%'.format(self.dropout*100))
#       print('recurrent dropout: {:5.0f}%'.format(self.recurrent_dropout*100))
#       
# 
#       #opt = SGD(lr=self.learning_rate, decay=self.lr_decay, momentum=0.9, nesterov=True)
#       opt = RMSprop(lr=0.001, rho=0.9, epsilon=1e-08, decay=0.0)
#   
#       model.compile(loss='categorical_crossentropy',
#                     optimizer=opt,
#                     metrics=['accuracy'],
#                     sample_weight_mode='temporal')
#     
#     print('Model compiled')
#     #os.system("nvidia-smi")     
#     
#     return model   
  
  def read(self, wfdb_dir, file_limit=None):
    self.mimo_wfdb_adaptor.read_all(wfdb_dir, file_limit, pad_target=True, out_dec_factor=self.out_dec_factor)
    return self.mimo_wfdb_adaptor.mimo
  
  def do_classification(self, test_mimo):
    last_predictions = test_mimo.do_classification_by_first_target_last_output()    

    n_seq = len(test_mimo.sequence_lengths)
    file = open(test_mimo.tag + '_answers.txt', 'w')
    for i in range(0, n_seq):
      file.write('{0},{1}\n'.format(test_mimo.sequence_tags[i], test_mimo.labels[last_predictions[i]]))
    file.close()

  def do_test(self, test_model, test_mimo, batch_size, tag=None, 
                  weights_filename=None, csvOutputFile=None, do_evaluate=False):
    if weights_filename is None:
      weights_filename=self.weights_filebase + self.weights_file_ext
    if tag is None:
      tag = 'Test';
    
    #test_mimo.tag = tag
    test_mimo = self.prepare_data(test_mimo, batch_size) 
      
    super(Ch2017MimoTrainer, self).do_test(test_model, test_mimo, batch_size, tag=tag, weights_filename=weights_filename, csvOutputFile=csvOutputFile, do_evaluate=do_evaluate)
    
    if do_evaluate:
        self.do_classification(test_mimo)
    
    
    
  
if __name__ == "__main__":
  batch_size = int(sys.argv[6])
  mt = Ch2017MimoTrainer(train_batch_size = int(batch_size))
  #mt.outdir_base = sys.argv[4]
  mt.set_outbase(sys.argv[4])
  mt.set_outdir(do_generate_out_dir=False)
  shutil.copy('../../ch2017_mimo_trainer.py', '.')
  shutil.copy('../../ch2017_graph_model.py', '.')
  file = open(mt.experiment_tag, 'w')
  file.close()
  
  mt.set_device(sys.argv[1])
  if sys.argv[2] == 'PartitionAndTrain':
    all_mimo = mt.read(sys.argv[3], int(sys.argv[5]))
    mt.partition_and_train(all_mimo)
    mt.shut_down()    
  if sys.argv[2] == 'CrossValidation':
    all_mimo = mt.read(sys.argv[3], int(sys.argv[5]))
    #experiment = MimoExperiment()
    #experiment.stratified_cross_validation(mt, all_mimo, int(sys.argv[7]))
    #mt.cross_validation(all_mimo, int(sys.argv[7]))
    mt.shut_down()    
  if sys.argv[2] == 'Test':
    mt.mimo_wfdb_adaptor.target_filename = 'REFERENCE.csv'
    mt.block_seq_base_length = int(sys.argv[8])
    if mt.block_seq_base_length == -1:
      mt.block_seq_base_length = sys.maxsize
    test_mimo = mt.read(sys.argv[3], int(sys.argv[5]))
    test_model = mt.create_model(test_mimo, batch_size)
    mt.do_test(test_model, test_mimo, batch_size, weights_filename=sys.argv[7], csvOutputFile='Test.csv', do_evaluate=True)
    mt.shut_down()
  else:
    print('Usage: mimo_trainer [device] ["PartitionAndTrain"] [input_dir] [file_limit] [train_batch_size]')
   
  