# 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



class Model:
  def __init__(self):

    self.learning_rate = tf.placeholder_with_default(FLAGS.learning_rate, ())
    self.iter_num = tf.placeholder_with_default(FLAGS.num_iterations, ())
    self.summaries = []
    self.noise_std = tf.placeholder_with_default(FLAGS.noise_std, ())
    self.opt = tf.train.AdamOptimizer(self.learning_rate)
    self.model_name = model.__file__

    assert FLAGS.image_range_start + FLAGS.encoder_length + FLAGS.decoder_future_length <= FLAGS.overall_images_count and FLAGS.image_range_start >= 0, \
            "settings for encoder/decoder lengths along with starting range exceed number of available images"
    assert FLAGS.encoder_length >= FLAGS.decoder_reconst_length, "encoder must be at least as long as reconstructer"

  def add_image_summary(self, summary_prefix, frames, encoder_length, decoder_future_length, decoder_reconst_length):
    for i in range(decoder_future_length):
      self.summaries.append(tf.summary.image(summary_prefix + '_future_gen_' + str(i + 1),
                                             self.frames_pred[i], max_outputs=1))

# 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




class Model:
  def __init__(self):

    self.learning_rate = tf.placeholder_with_default(FLAGS.learning_rate, ())
    self.iter_num = tf.placeholder_with_default(FLAGS.num_iterations, ())
    self.summaries = []
    self.noise_std = tf.placeholder_with_default(FLAGS.noise_std, ())
    self.opt = tf.train.AdamOptimizer(self.learning_rate)
    self.model_name = model.__file__

    assert FLAGS.image_range_start + FLAGS.encoder_length + FLAGS.decoder_future_length <= FLAGS.overall_images_count and FLAGS.image_range_start >= 0, \
            "settings for encoder/decoder lengths along with starting range exceed number of available images"
    assert FLAGS.encoder_length >= FLAGS.decoder_reconst_length, "encoder must be at least as long as reconstructer"

  def add_image_summary(self, summary_prefix, frames, encoder_length, decoder_future_length, decoder_reconst_length):
    for i in range(decoder_future_length):
      self.summaries.append(tf.summary.image(summary_prefix + '_future_gen_' + str(i + 1),
                                             self.frames_pred[i], max_outputs=1))

      self.summaries.append(tf.summary.image(summary_prefix + '_future_orig_' + str(i + 1),