# Path: settings.py
# FLAGS = flags.FLAGS
# OUT_DIR = '/common/homes/students/rothfuss/Documents/selected_trainings/9_20bn_vae_no_OF/08-06-18_10-21'
# DUMP_DIR = "/common/homes/students/rothfuss/Documents/selected_trainings/9_20bn_vae_no_OF/08-06-18_10-21"
# TF_RECORDS_DIR = "/localhome/rothfuss/data/20bn-something/tf_records_train"
# MODE = 'valid_mode'
# VALID_MODE = 'data_frame' #'data_frame gif'
# NUM_IMAGES = 15
# NUM_DEPTH = 3
# WIDTH = 128
# HEIGHT = 128
# NUM_THREADS_QUEUERUNNER = 32 # specifies the number of pre-processing threads
# PRETRAINED_MODEL = "/common/homes/students/rothfuss/Documents/selected_trainings/9_20bn_vae_no_OF/08-06-18_10-21"
# EXCLUDE_FROM_RESTORING = None
# FINE_TUNING_WEIGHTS_LIST = None
# INPUT_DIR = ""
# MEMORY_PATH = ""
# LOSS_FUNCTIONS = ['mse', 'gdl', 'mse_gdl', 'vae']
# MODES = ["train_mode", "valid_mode", "feeding_mode"]
# VALID_MODES = ['count_trainable_weights', 'vector', 'gif', 'similarity', 'data_frame', 'psnr', 'memory_prep', 'measure_test_time']
# 
# Path: models/loss_functions.py
# def gradient_difference_loss(true, pred, alpha=2.0):
# def difference_gradient(image, vertical=True):
# def mean_squared_error(true, pred):
# def peak_signal_to_noise_ratio(true, pred):
# def kl_penalty(mu, sigma):
# def decoder_loss(frames_gen, frames_original, loss_fun):
# def decoder_psnr(frames_gen, frames_original):
# def composite_loss(original_frames, frames_pred, frames_reconst, loss_fun='mse',
#                    encoder_length=5, decoder_future_length=5,
#                    decoder_reconst_length=5, mu_latent=None, sigm_latent=None):


import tensorflow as tf
import math
import numpy as np
import data_prep.model_input as input
from pprint import pprint
from settings import FLAGS, model
from models import loss_functions
    print("Constructing TrainModel")

    with tf.variable_scope(scope_name, reuse=None) as training_scope:
      Model.__init__(self)
      self.scope = training_scope

      tower_grads = []
      tower_losses = []
      for i in range(FLAGS.num_gpus):
        train_batch, _, _ = input.create_batch(FLAGS.tf_records_dir, 'train', FLAGS.batch_size,
                                               int(math.ceil(
                                                 FLAGS.num_iterations / (FLAGS.batch_size * 20))),
                                               False)
        train_batch = tf.cast(train_batch, tf.float32)
        with tf.device('/gpu:%d' % i):
          with tf.name_scope('%s_%d' % ('tower', i)):
            tower_loss, _, _, _ = tower_operations(train_batch[:,FLAGS.image_range_start:,:,:,:], train=True)
            tower_losses.append(tower_loss)

            # Reuse variables for the next tower.
            tf.get_variable_scope().reuse_variables()

            if FLAGS.fine_tuning_weights_list is not None:
              train_vars = []
              for scope_i in FLAGS.fine_tuning_weights_list:
                train_vars += tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=scope_i)
              pprint('Finetuning. Training only specified weights: %s' % (FLAGS.fine_tuning_weights_list))
              grads = self.opt.compute_gradients(tower_loss, var_list=train_vars)
            else:
              grads = self.opt.compute_gradients(tower_loss)

# Path: settings.py
# FLAGS = flags.FLAGS
# OUT_DIR = '/common/homes/students/rothfuss/Documents/selected_trainings/9_20bn_vae_no_OF/08-06-18_10-21'
# DUMP_DIR = "/common/homes/students/rothfuss/Documents/selected_trainings/9_20bn_vae_no_OF/08-06-18_10-21"
# TF_RECORDS_DIR = "/localhome/rothfuss/data/20bn-something/tf_records_train"
# MODE = 'valid_mode'
# VALID_MODE = 'data_frame' #'data_frame gif'
# NUM_IMAGES = 15
# NUM_DEPTH = 3
# WIDTH = 128
# HEIGHT = 128
# NUM_THREADS_QUEUERUNNER = 32 # specifies the number of pre-processing threads
# PRETRAINED_MODEL = "/common/homes/students/rothfuss/Documents/selected_trainings/9_20bn_vae_no_OF/08-06-18_10-21"
# EXCLUDE_FROM_RESTORING = None
# FINE_TUNING_WEIGHTS_LIST = None
# INPUT_DIR = ""
# MEMORY_PATH = ""
# LOSS_FUNCTIONS = ['mse', 'gdl', 'mse_gdl', 'vae']
# MODES = ["train_mode", "valid_mode", "feeding_mode"]
# VALID_MODES = ['count_trainable_weights', 'vector', 'gif', 'similarity', 'data_frame', 'psnr', 'memory_prep', 'measure_test_time']
# 
# Path: models/loss_functions.py
# def gradient_difference_loss(true, pred, alpha=2.0):
# def difference_gradient(image, vertical=True):
# def mean_squared_error(true, pred):
# def peak_signal_to_noise_ratio(true, pred):
# def kl_penalty(mu, sigma):
# def decoder_loss(frames_gen, frames_original, loss_fun):
# def decoder_psnr(frames_gen, frames_original):
# def composite_loss(original_frames, frames_pred, frames_reconst, loss_fun='mse',
#                    encoder_length=5, decoder_future_length=5,
#                    decoder_reconst_length=5, mu_latent=None, sigm_latent=None):


# Path: core/Model.py
import tensorflow as tf
import math
import numpy as np
import data_prep.model_input as input
from pprint import pprint
from settings import FLAGS, model
from models import loss_functions

    print("Constructing TrainModel")

    with tf.variable_scope(scope_name, reuse=None) as training_scope:
      Model.__init__(self)
      self.scope = training_scope

      tower_grads = []
      tower_losses = []
      for i in range(FLAGS.num_gpus):
        train_batch, _, _ = input.create_batch(FLAGS.tf_records_dir, 'train', FLAGS.batch_size,
                                               int(math.ceil(
                                                 FLAGS.num_iterations / (FLAGS.batch_size * 20))),
                                               False)
        train_batch = tf.cast(train_batch, tf.float32)
        with tf.device('/gpu:%d' % i):
          with tf.name_scope('%s_%d' % ('tower', i)):
            tower_loss, _, _, _ = tower_operations(train_batch[:,FLAGS.image_range_start:,:,:,:], train=True)
            tower_losses.append(tower_loss)

            # Reuse variables for the next tower.
            tf.get_variable_scope().reuse_variables()

            if FLAGS.fine_tuning_weights_list is not None:
              train_vars = []
              for scope_i in FLAGS.fine_tuning_weights_list:
                train_vars += tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope=scope_i)
              pprint('Finetuning. Training only specified weights: %s' % (FLAGS.fine_tuning_weights_list))
              grads = self.opt.compute_gradients(tower_loss, var_list=train_vars)
            else:
              grads = self.opt.compute_gradients(tower_loss)

            tower_grads.append(grads)