'''
Created on May 4, 2017

@author: philip
'''
import os
import tensorflow as tf

import keras
import keras.backend.tensorflow_backend as K
from keras.models import Sequential, load_model
from keras.callbacks import EarlyStopping, ModelCheckpoint, TensorBoard
from train.weighted_accuracy import WeightedAccuracy

from tensorflow.core.framework import graph_pb2
from tensorflow.python.framework import graph_io
from tensorflow.python.tools import freeze_graph
from tensorflow.python.training import saver as saver_lib

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


  def __init__(self, device_id='/cpu:0'):      
    self.device_id = device_id
    self.model = []
    self.config()
    
    self.experiment_tag='default'    
    self.early_stopping_epochs = 20
    self.max_epochs = 400
#     self.weights_filebase = 'best_weights'
#     self.weights_file_ext = '.hdf5'
#     self.weights_filename = self.weights_filebase + self.weights_file_ext 
    
    
      
      
  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()
    
    self.model = model    
    return self.model
  
  
    
  def save_model_to_frozen_pb(self, model, save_dir, output_node_names, output_graph_name):
    K.clear_session()
    load_model(model)
    
    checkpoint_prefix = os.path.join(save_dir, "saved_checkpoint")
    checkpoint_state_name = "checkpoint_state"
    input_graph_name = "input_graph.pb"
    output_graph_name = "output_graph.pb"
    
    saver = saver_lib.Saver()
    checkpoint_path = saver.save(
        K.get_session(),
        checkpoint_prefix,
        global_step=0,
        latest_filename=checkpoint_state_name)
#    print(checkpoint_path) /saved_checkpoint-0
#    from tensorflow.python.tools import inspect_checkpoint
#    inspect_checkpoint.print_tensors_in_checkpoint_file(checkpoint_path, tensor_name=None, all_tensors=True)
    
    graph_io.write_graph(K.get_session().graph, save_dir, input_graph_name)
    input_graph_path = os.path.join(save_dir, input_graph_name)
    input_saver_def_path = ""
    input_binary = False
    restore_op_name = "save/restore_all"
    filename_tensor_name = "save/Const:0"
    output_graph_path = os.path.join(save_dir, output_graph_name)
    clear_devices = False
    freeze_graph.freeze_graph(input_graph_path, input_saver_def_path,
                              input_binary, checkpoint_path, output_node_names,
                              restore_op_name, filename_tensor_name,
                              output_graph_path, clear_devices, "")
    K.clear_session()
#    [print(n.name) for n in tf.get_default_graph().as_graph_def().node]
    return output_graph_path
  
  def load_frozen_pb(self, graph_pb):
    output_graph_def = graph_pb2.GraphDef()
    with open(graph_pb, "rb") as f:
        output_graph_def.ParseFromString(f.read())
        tf.import_graph_def(output_graph_def, name="")
        
  def run_inference_on_frozen(self, block_input):    
    y = K.get_session().graph.get_tensor_by_name('dense_1/truediv:0')
    y_out = K.get_session().run(y, feed_dict={'keras_learning_phase_1:0': 0, 'lstm_1_input:0': block_input})
    return y_out

  def save_array_to_chkpt(self, ndarray, tensor_name, prefix):
    with tf.Session() as sess:
        init = tf.constant(ndarray)
        tf.get_variable(tensor_name, initializer=init)
        init_op = tf.global_variables_initializer()
        sess.run(init_op)
        saver = saver_lib.Saver()
        saver.save(sess, prefix)
        
  def do_train(self, trainMimo, validationMimo, batch_size, weights_filename, tag=None):
    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(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 = self.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()
    
    return hist.history
    #self.plot_training_metrics(hist.history, tag)
  
  
if __name__ == "__main__":
  gm = GraphModel()
  gm.create_model()
  