'''
Created on May 4, 2017

@author: philip
'''
import tensorflow as tf

import keras
import keras.backend.tensorflow_backend as K
from keras.models import Sequential
from keras.layers import Dense, LSTM, Convolution1D, MaxPooling1D, BatchNormalization
from keras.optimizers import SGD, RMSprop

from train.graph_model import GraphModel


class Ch2017GraphModel(GraphModel):
  '''
  Patterns2017 graph models (Keras only for now)
  '''

  def __init__(self, model_mimo, batch_size, max_length=None, device_id='/cpu:0'):  
    super(Ch2017GraphModel, self).__init__()    
    
    self.model_mimo = model_mimo
    self.batch_size = batch_size
    self.max_length = max_length
    self.device_id = device_id

    #self.lstm_n_cell=[50, 50, 50] # Trainable params: , 355s/epoch
    self.lstm_n_cell=[100, 100, 100] # Trainable params: , 355s/epoch
    self.dropout=0.1 
    self.recurrent_dropout=0.1
    self.learning_rate = 1e-1
    self.lr_decay = 1e-6
    
    
    self.config()
    
  def config(self):
    print("Keras version: " + keras.__version__)
    print("Tensorflow version: " + tf.__version__)
    
    config = tf.ConfigProto()
    #self.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))
    return config
    
  def create_model(self):        
      
    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(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=(self.batch_size, self.max_length, self.model_mimo.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(self.model_mimo.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")     
        
    self.model = model
    return self
        
    
if __name__ == "__main__":
  mt = Ch2017GraphModel()
  