from __future__ import print_function
import os

#os.environ['TF_CPP_MIN_LOG_LEVEL']='0' #info
#os.environ['TF_CPP_MIN_LOG_LEVEL']='1' #warning
os.environ['TF_CPP_MIN_LOG_LEVEL']='2' #error
#os.environ['TF_CPP_MIN_LOG_LEVEL']='3' #fatal
import tensorflow as tf
#tf.logging.set_verbosity(tf.logging.FATAL)


import sys
import math 

import time

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.tri import triplot
from matplotlib.pyplot import savefig

from keras.models import Sequential
from keras.layers import Dense, LSTM, Activation, Masking, Convolution1D, MaxPooling1D
from keras.layers.wrappers import TimeDistributed 
#from keras.initializations import normal, identity
from keras.optimizers import SGD, RMSprop
#from keras.callbacks import EarlyStopping, ModelCheckpoint, TensorBoard

from keras.preprocessing.sequence import pad_sequences
import keras.backend.tensorflow_backend as K

import keras

from keras.backend.tensorflow_backend import set_session

from utils.profiler import profile
  
#from train.weighted_accuracy import WeightedAccuracy


#
#  MimoTrainer: a class to patterns_train multiple-input multiple-output (MIMO) data using Keras and Tensorflow
#
class MimoTrainer:
  def __init__(self, train_batch_size=10, block_seq_base_length=6000, out_dec_factor=1):
    #self.train_batch_size = 248  
    #self.train_batch_size = 62  
    #self.train_batch_size = 6  
    #self.train_batch_size = 50 
    self.train_batch_size = train_batch_size  
    self.test_batch_size = 1   
    #self.max_epochs = 10
    self.max_epochs = 256
    #self.max_epochs = 1
    #self.look_ahead = 30   # 30 samples x 4s/sample = 120s  
    self.look_ahead = 0   # 30 samples x 4s/sample = 120s  
    #self.look_ahead = 1000   # 30 samples x 4s/sample = 120s  
    self.learning_rate = 1e-1
    #self.learning_rate = 1e-2
    self.lr_decay = 1e-6
    #self.early_stopping_epochs = 20
    #self.lstm_n_cell=[50, 50, 50] # Trainable params: , 355s/epoch
    self.weights_filebase = 'best_weights'
    self.weights_file_ext = '.hdf5'
    self.weights_filename = self.weights_filebase + self.weights_file_ext 
    #self.dropout=0.0 
    #self.recurrent_dropout=0.0
#    self.dropout=0.1 
#    self.recurrent_dropout=0.1
    #self.device_id = '/cpu:0'
    self.n_trainTestCycles = 1    
    #self.n_trainTestCycles = 1
    #self.input_gaussian_noise_sigma = 0.6    
    self.input_gaussian_noise_sigma = 0.0    
    #self.redirect_std()
#    self.dropout=0.5 
#    self.recurrent_dropout=0.5
    #self.train_proportion = 0.9
    #self.train_proportion = 0.975
    self.train_proportion = 0.9924
    self.block_seq_base_length = block_seq_base_length
    #self.block_seq_base_length = 6000
    #self.block_seq_base_length = 3000
    #self.block_seq_base_length = 1500
    #self.block_seq_base_length = 750
    #self.block_seq_base_length = 375
    self.outdir_base = os.getcwd()
    self.outdir = self.outdir_base
    self.out_dec_factor = out_dec_factor 
#        
#Redirecting stdout and stderr within a program is a bad idea 
#(only the local sys.stdout and stderr are affected)
# Use the command-line method instead
#   def redirect_std(self):
#     f = open(self.outdir + '/' + self.__class__.__name__ + '_log.txt', 'w')
#     sys.stdout = sys.stderr = f    

  def set_outbase(self, outbase):
    self.outdir_base = outbase
    if not os.path.exists(self.outdir_base):
      os.makedirs(self.outdir_base) 
    print('Created output base dir: ' + self.outdir_base)
    os.chdir(self.outdir_base)

    
  def set_outdir(self, do_generate_out_dir=False):
    if not do_generate_out_dir:
      return
    t = time.localtime()
    timestamp = '%4d%02d%02d_%02d%02d%02d' % (t.tm_year, t.tm_mon, t.tm_mday, 
                                         t.tm_hour, t.tm_min, t.tm_sec)
    self.outdir = self.outdir_base + '/' + timestamp
    if not os.path.exists(self.outdir):
      os.makedirs(self.outdir) 
    print('Output sent to: ' + self.outdir)
    os.chdir(self.outdir)

    
  def plot_training_metrics(self, hist, tag=None):
      if tag is None:
        tag = ''

      print('Plotting Training metrics')
          
      lossLine, = plt.plot(hist['loss'], 'r', label="Train", linewidth = 1)
      vlossLine, = plt.plot(hist['val_loss'], 'b', label="Validation", linewidth = 1)
      
      first_legend = plt.legend(handles=[lossLine], loc=1)    
      ax = plt.gca().add_artist(first_legend)
      plt.legend(handles=[vlossLine], loc=7)
      plt.ylabel('Loss')
      plt.xlabel('Epoch')
        
      savefig('TrainingMetrics' + tag +'.pdf', bbox_inches='tight')
      plt.gca().cla()
      #plt.show()
      
  def set_device(self, deviceId):
    self.device_id = deviceId

  def create_model(self, modelMimo, batch_size, max_length=None):    
    return 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.device_id):
#     
#       print('Creating Model')
#       model = Sequential()
#       #model.add(Masking(mask_value=0., batch_input_shape=(batch_size, max_length, modelMimo.n_input_dim)))
#       
#       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(MaxPooling1D(pool_size=8))
# #       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(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(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))
#        
#       # didn't have a final Dense layer with CURRENNT
#       #model.add(Dense(modelMimo.n_labels, init='he_normal',
#     #                                  activation='softmax'))
#     # Can't use Dense because it doesn't support maksing
# #       model.add(TimeDistributed(Dense(modelMimo.n_labels+1, kernel_initializer='he_normal',
# #                                       activation='softmax'))) 
#       model.add(Dense(modelMimo.n_labels+1, kernel_initializer='he_normal',
#                                       activation='softmax')) 
#       
#       print (model.summary())
#       #model = make_parallel(model, 2)
#       
#       #print('Model created')
#       #os.system("nvidia-smi")
# 
#       
#       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 do_train(self, graph_model, trainMimo, validationMimo, batch_size, tag=None):
    hist = graph_model.do_train(trainMimo, validationMimo, batch_size, self.weights_filename, tag)
    self.plot_training_metrics(hist, tag)
    
#     if tag is None:
#       tag = '';
# 
#     print('Train input shape')
#     print(trainMimo.block_input.shape)
#     print('Train output shape')
#     print(trainMimo.block_target.shape)
#     
#     print('Validation input shape')
#     print(validationMimo.block_input.shape)
#     print('Validation output shape')
#     print(validationMimo.block_target.shape)
#                   
#     # setup callbacks 
#     early_stopping = EarlyStopping(monitor='val_loss', 
#                                    patience=self.early_stopping_epochs)
#     validation_weighted_acc = WeightedAccuracy('Validation', 
#                                                 validationMimo.block_input, 
#                                                 validationMimo.block_target, 
#                                                 validationMimo.sample_weights, 
#                                                 batch_size)
#     model_checkpoint = ModelCheckpoint(self.weights_filename,
#                                       monitor='val_loss', verbose=0, 
#                                       save_best_only=True, 
#                                       save_weights_only=False, 
#                                       mode='auto', period=1)
# #     tensor_board = TensorBoard(log_dir= self.outdir_base + './tensorboard', 
# #                                histogram_freq=1, 
# #                                write_graph=True, write_images=True)
#       
#     print('Training')
#     hist = graph_model.model.fit(trainMimo.block_input,
#               trainMimo.block_target,
#               batch_size,
#               verbose=1,
#               epochs=self.max_epochs,
#               shuffle=True,
#               validation_data=(validationMimo.block_input, 
#                                validationMimo.block_target,
#                                validationMimo.sample_weights),
#               sample_weight = trainMimo.sample_weights,
#               callbacks=[early_stopping, 
#                          #train_weighted_acc, 
#                          validation_weighted_acc,
#                          model_checkpoint,
#                          #tensor_board 
#                          ])
#             
#     print(hist.history)
#     minValI = min(range(len(hist.history['val_loss'])), key=hist.history['val_loss'].__getitem__)
#     minVal = hist.history['val_loss'][minValI]
#     print('Minimum validation loss of %.4f at epoch %d' % (minVal, minValI+1))
#     #sys.stdout.flush()
#     
#     self.plot_training_metrics(hist.history, tag)
    

  def shift_left(self, prediction, n_samples):
    last_samples = prediction[:, np.newaxis, prediction.shape[1]-1, :];
    last_samples_padded = np.tile(last_samples, (1, n_samples, 1));
    prediction = np.concatenate((prediction, last_samples_padded), axis=1)
    prediction = np.delete(prediction, np.s_[0:n_samples:1], axis=1) 
         
    return prediction
  

  def do_classification(self, test_mimo):
    test_mimo.do_classification()
    
  
  def do_test(self, graph_model, test_mimo, test_batch_size, do_evaluate=True, 
              tag=None, weights_filename=None, csvOutputFile=None):
    if weights_filename is None:
      weights_filename=self.weights_filebase + self.weights_file_ext
    if tag is None:
      tag = '';
    print("loading weights from " + weights_filename)
    graph_model.model.load_weights(weights_filename)
      
    print('input shape')
    print(test_mimo.block_input.shape)
    #print('output shape')
    #print(test_mimo.block_target.shape)
    
    print('Predicting: ' + test_mimo.tag)
    #test_predicted_output = model.predict(test_mimo.block_input, test_batch_size)
    block_output = graph_model.model.predict(test_mimo.block_input, test_batch_size)
    test_mimo.set_block_output(block_output)
    
    ## DEBUG
    #test_mimo.pad_last_output()
    ## DEBUG
    
    if do_evaluate:
      test_mimo.plot_target_and_output(tag)
      
      print('Evaluating: ' + test_mimo.tag)
      metrics = graph_model.model.evaluate(test_mimo.block_input, test_mimo.block_target, 
                     batch_size=test_batch_size, verbose=1, sample_weight=test_mimo.sample_weights)
      print(metrics)
      
#       class_report_str = classification_report(test_mimo.block_target, 
#                                                test_mimo.block_output, 
#                                                labels=range(0, test_mimo.n_labels), 
#                                                target_names=test_mimo.labels, 
#                                                sample_weight=test_mimo.sample_weights, 
#                                                digits=2)    
#        
#       print(class_report_str)
#       
      
       
    # undo the output lag induced by training
    #self.plot_test_results(test_mimo.tag + tag + 'BeforeShift', test_mimo.block_target, test_predicted_output);
    #test_predicted_output = self.shift_left(test_predicted_output, self.look_ahead)
    #self.plot_test_results(test_mimo.tag + tag + 'AfterShift', test_mimo.block_target, test_predicted_output);
    
    if csvOutputFile is not None:
      test_mimo.write_predictions(csvOutputFile)
  
    sys.stdout.flush()
    

  #
  #  MimoTrainer.do_optimal_padding
  #
  #  For training purposes, the batch size must be equal for both validation and
  #  training sets. The usual scenario is that the training set size is an 
  #  integer multiple of the validation set size.  do_optimal_padding searches for 
  #  a combination of training or validation set sizes that minimizes first the
  #  number of batches, then the number of total padded sequences. This may not 
  #  be the optimal criterion for all scenarios.
  #    
  def do_optimal_padding(self, trainMimo, validationMimo, max_batch_size=None):
    new_train_n_sequences = 0
    new_val_n_sequences = 0
    min_pad_size = math.inf
    max_pad_size = min(trainMimo.n_sequences(), validationMimo.n_sequences()) - 1
#     if max_batch_size not is None
#       max_pad_size = min(max_batch_size)
    padGrid = np.zeros((max_pad_size, max_pad_size))
    train_pad = 0
    validation_pad = 0
    #opt_n_batches = 0
    opt_n_batches = math.inf
    
    print('Input # training sequences: ' + str(trainMimo.n_sequences())) 
    print('Input # validation sequences: ' + str(validationMimo.n_sequences())) 
      
    for i in range(max_pad_size): 
      for j in range(max_pad_size):
        temp_train_n_sequences = trainMimo.n_sequences() + i 
        temp_val_n_sequences = validationMimo.n_sequences() + j 
        if temp_train_n_sequences >= temp_val_n_sequences:
          n_pad = trainMimo.calculate_pad_size(temp_train_n_sequences, temp_val_n_sequences)
          if n_pad == 0:
            n_batches = temp_train_n_sequences // temp_val_n_sequences
            #pad_size = i*n_batches + j
            pad_size = i + j
            padGrid[i, j] = pad_size;
            if n_batches < opt_n_batches and pad_size < min_pad_size:
            #if n_batches > opt_n_batches:
              min_pad_size = pad_size
              new_train_n_sequences = temp_train_n_sequences
              new_val_n_sequences = temp_val_n_sequences
              train_pad = i
              validation_pad = j
              opt_n_batches = n_batches
              print('#batches:' + str(n_batches) + ' i=' + str(i) + ' (train_n_seq=' + str(new_train_n_sequences) + ') j=' + str(j) + ' val_n_seq=(' + str(new_val_n_sequences) + ') total pad_size:' + str(min_pad_size))
              
              
#     plt.contour(range(max_pad_size), range(max_pad_size), padGrid)     
#     plt.plot(train_pad, validation_pad, 'rx') 
#     plt.xlabel('Train pad')
#     plt.ylabel('Validation pad')
#     savefig('PadSearch.pdf', bbox_inches='tight')              
#     plt.gca().cla()

    print("pad search grid:")
    print(padGrid)
    trainMimo.pad_number_of_sequences(new_train_n_sequences)
    validationMimo.pad_number_of_sequences(new_val_n_sequences)

    print('Optimized # training sequences: ' + str(trainMimo.n_sequences())) 
    print('Optimized # validation sequences: ' + str(validationMimo.n_sequences())) 
    print('Optimized # batches: ' + str(opt_n_batches)) 

    return trainMimo, validationMimo
        
  def prepare_data(self, mimo, batch_size):
    mimo = mimo.split_long(self.block_seq_base_length)
    mimo.generate_block_sequences()
    
    print('Sequence: ' + mimo.tag)
    print('Initial # sequences: ' + str(mimo.n_sequences()))
    print('Batch size: ' + str(batch_size))   
    mimo.pad_number_of_sequences(batch_size)
    print('Padded # sequences: ' + str(mimo.n_sequences())) 
    return mimo
    

  def partition_and_train(self, all_mimo):
    train_len = math.floor(all_mimo.n_sequences() * self.train_proportion)  
    train_mimo, validation_mimo = all_mimo.filter_by_index(range(train_len), 'patterns_train', 'validation')
        
    train_mimo = train_mimo.split_long(self.block_seq_base_length)
    validation_mimo = validation_mimo.split_long(self.block_seq_base_length)
    
    train_mimo.generate_block_sequences()
    validation_mimo.generate_block_sequences()
    
    train_mimo.plot_input()
    #train_mimo.add_input_gaussian_noise(self.input_gaussian_noise_sigma)
    #train_mimo.plot_input('TrainInputWithNoise')

    validation_mimo.plot_input()
    
    train_mimo.plot_target() 
    validation_mimo.plot_target()    

    #train_mimo, validation_mimo = self.do_optimal_padding(train_mimo, validation_mimo)
    #self.train_batch_size = validation_mimo.n_sequences()
    
    print('Initial # training sequences: ' + str(train_mimo.n_sequences())) 
    print('Initial # validation sequences: ' + str(validation_mimo.n_sequences()))
    
    print('Batch size: ' + str(self.train_batch_size))
   
    train_mimo.pad_number_of_sequences(self.train_batch_size)
    validation_mimo.pad_number_of_sequences(self.train_batch_size)
    
    print('Padded # training sequences: ' + str(train_mimo.n_sequences())) 
    print('Padded # validation sequences: ' + str(validation_mimo.n_sequences())) 
    

    trainModel = self.create_model(train_mimo, self.train_batch_size, train_mimo.max_length)
    cycleTag = '';
    for i in range(self.n_trainTestCycles):
      if self.n_trainTestCycles>1:
        cycleTag = '_Cycle' + str(i);
      self.weights_filename = self.weights_filebase + cycleTag + self.weights_file_ext

      self.do_train(trainModel, train_mimo, validation_mimo, self.train_batch_size, tag=cycleTag)
  
      self.do_test(trainModel, train_mimo, self.train_batch_size, tag=cycleTag, csvOutputFile=self.outdir+'/trainOutput')
      self.do_test(trainModel, validation_mimo, self.train_batch_size, tag=cycleTag, csvOutputFile=self.outdir+'/validationOutput')
    
    
  def shut_down(self):
    K.clear_session()
    
        
    
if __name__ == "__main__":
  mt = MimoTrainer(train_batch_size = int(sys.argv[6]))
  mt.outdir_base = sys.argv[4]
  mt.set_outdir()
  mt.set_device(sys.argv[1])
  if sys.argv[2] == 'PartitionAndTrain':
    mt.partition_and_train(sys.argv[3], int(sys.argv[5]))
  else:
    print('Usage: MimoTrainer [device] ["PartitionAndTrain"] [input_dir] [file_limit] [train_batch_size]')
   
  