import tensorflow as tf
import numpy as np

import os
import learning.utils
from learning.utils import utilities

from tensorflow.examples.tutorials.mnist import input_data
from learning.utils.input_format import DataSet


class RBM(object):
    def __init__(self, input_size, hidden_size, main_dir="rbm", model_name="rbm_model",
                 standard_derivation=0.01, input_to_binary=False, verbose=False, ):
        '''
        :param input_size: Number of features used in the input vector.
        :param hidden_size: Number of classes the output should be devided in.
        :param main_dir: Directory where trainingdata, modeldata and imagedata should be saved to.
        :param model_name: Name of the model.
        :param standard_derivation: Standard-derivation for the random initialised weights.
        :param input_to_binary: Determines if the network converts input into binary digits.
        :param verbose: If true it will print the result of loss_function.
        '''

        # Sets network-size
        self._input_size = input_size
        self._hidden_size = hidden_size

        # parameters to influence the learning process
        self._gibbs_sampling_steps = None
        self._batch_size = None
        self._momentum_factor = None
        self._learning_rate = None
        self._weight_decay_factor = None

        self._input_to_binary = input_to_binary
        self._standard_derivation = standard_derivation
        self._epochs = None

        # Tensorflow: variabales
        self._tf_w = None
        self._tf_bh = None
        self._tf_bv = None

        self._tf_delta_w = None
        self._tf_delta_bh = None
        self._tf_delta_bv = None

        self._tf_w_upd8 = None
        self._tf_bh_upd8 = None
        self._tf_bv_upd8 = None

        # Tesorflow: Placeholders
        self._tf_lr = None
        self._tf_momentum = None
        self._tf_weight_decay = None
        self._tf_gibbs_steps = None

        self._tf_input_data = None
        self._tf_hrand = None
        self._tf_vrand = None
        self._tf_desired_output = None

        # Varaibles to profile the training progress
        self._tf_loss_function = None
        self._tf_accuracy = None

        # Tensoflow: Manage tensoflow interactions
        self._tf_session = None
        self._tf_merged_summaries = None
        self._tf_summary_writer = None
        self._tf_saver = None

        # Profiling and Saving of Modellparameters
        self._main_dir = main_dir
        self._model_name = model_name
        self._model_dir, self._data_dir, self.summary_dir = self._create_data_directories()
        self._model_path = self._model_dir + self._model_name
        self._verbose = verbose

    def fit(self, train_set, restore_previous_model=None, start_epoche=0, epochs=300,
            gibbs_sampling_steps=1, learning_rate=0.01, momentum_factor=0.9, weight_decay_factor=0.0001, batch_size=10):
        '''

        :param train_set: Trainset formatted as input_format.
        :param restore_previous_model: If true restores learned model weights.
        :param start_epoche: Already passed epochs to keep the Tensorboard graph well formatted.
        :param epochs: Number of epochs the learning should last.
        :param gibbs_sampling_steps: Number of gibbs sampling steps the net should performs.
        :param learning_rate: Hyperparameter Learning rate parameter; will be devided by batch_size.
        :param momentum_factor: Specifies the momentumterm.
        :param weight_decay_factor:Specifies how strong the high weights are penetalized.
        :param batch_size: The trainingset is sliced in mini batches. One batch has the size of batch_size.
        :return: self.
        '''

        self._epochs = epochs
        self._learning_rate = learning_rate / batch_size
        self._momentum_factor = momentum_factor
        self._weight_decay_factor = weight_decay_factor * self._learning_rate
        self._gibbs_sampling_steps = gibbs_sampling_steps
        self._batch_size = batch_size

        self._build_model()
        self._create_summary_nodes()

        with tf.Session() as self._tf_session:
            self._initialize_tf_utilities_and_ops(restore_previous_model)
            self._train_model(train_set, start_epoche)
            self._tf_saver.save(self._tf_session, self._model_path)

        tf.reset_default_graph()

    def classify(self, input, input_to_binary=False, return_hstates=False):

        '''

        :param input: the input vectors to classify as input_format.
        :param input_to_binary: if true input is normalized to 0 or 1.
        :param return_hstates: if true returns only the values of 0 or 1.
        :return: binary vector or probabilty which indicates the corresponding class of the inputdata.
        '''

        self._create_placeholders()
        self._create_variables()

        if input_to_binary:
            self._tf_input_data = self._sample_prob(self._tf_input_data, self._tf_vrand)

        classification = self._sample_hidden_from_visible(self._tf_input_data)

        with tf.Session() as self._tf_session:
            init_op = tf.global_variables_initializer()
            self._tf_session.run(init_op)
            self._tf_saver = tf.train.Saver()
            self._tf_saver.restore(self._tf_session, self._model_path)

            hprobs, hstates = self._tf_session.run(classification, feed_dict=self._create_complete_dict(input))

        tf.reset_default_graph()

        if return_hstates:
            return hstates
        else:
            return hprobs

    def _initialize_tf_utilities_and_ops(self, restore_previous_model):

        '''

        :param restore_previous_model: if set to true the lastest learned weights are loaded to continue training.
        :return: self.
        '''
        init_op = tf.global_variables_initializer()

        self._tf_session.run(init_op)

        self._tf_merged_summaries = tf.summary.merge_all()
        self._tf_saver = tf.train.Saver()

        if restore_previous_model:
            self._tf_saver.restore(self._tf_session, self._model_path)

        self._tf_summary_writer = tf.summary.FileWriter(self.summary_dir, self._tf_session.graph)

    def _train_model(self, train_set, start_epoche):

        '''
            Devides the training process into the number of epochs. Every epoch run_train_step is called to iterate
            over the whole train_set.

        :param train_set: The input data formatted as input_format.
        :param start_epoche: If the training is performed in multiple steps then the start_epoch should be specified.
                    The start_epoch is the number of already performed trainingsteps.
        :return: self.
        '''

        for i in range(self._epochs):
            self._run_train_step(train_set)
            # print("finished train-step ", i)
            if i % 5 == 0:
                self._run_summary(train_set, i + start_epoche)

    def _run_train_step(self, train_set):

        '''
            Iterates over the whole training set.

        :param train_set: The input data formatted as input_format.
        :return: self.
        '''

        updates = [self._tf_w_upd8, self._tf_bh_upd8, self._tf_bv_upd8]

        iterations = (int)(train_set.num_examples / self._batch_size)

        for i in range(iterations):
            self._tf_session.run(updates, feed_dict=self._create_feed_dict(train_set))

    def _run_summary(self, train_set, epoch):

        '''
            After every 20 completed epochs the network performs a mergesummary call. The in create_summary_node
            called tensors are save to the logs file. Also the loss is calculated and also saved. The progress of the
            network can be monitored by tensorboard.

        :param train_set: [Update Needed] Can either operate the validation or the trainset.
        :param epoch: An integer to plot the summary data to an x-value.
        :return: self.
        '''

        loss, summary_str = self._tf_session.run([self._tf_loss_function, self._tf_merged_summaries],
                                                 feed_dict=self._create_feed_dict(train_set))

        self._tf_summary_writer.add_summary(summary_str, epoch)

        if self._verbose:
            print("[loss]: ", loss)

    def _create_feed_dict(self, data):

        '''
            Helper function, which slices the inputdata at every iteration and feeds the data in a dictionary.
            The dictionary is used to feed the tensorplaceholders.

        :param data: takes the whole train_set in form of input_format.
        :return: A dictionary with the data to feed to the network.
        '''

        x, y = data.next_batch(self._batch_size)
        return {
            self._tf_input_data: x,
            self._tf_hrand: np.random.rand(x.shape[0], self._hidden_size),
            self._tf_vrand: np.random.rand(x.shape[0], self._input_size),
            self._tf_lr: self._learning_rate,
            self._tf_momentum: self._momentum_factor,
            self._tf_weight_decay: self._weight_decay_factor,
            self._tf_gibbs_steps: self._gibbs_sampling_steps
        }

    def _create_complete_dict(self, data):

        '''
            This function works like create_feed_dict but it also feeds a placeholder(self._tf_desired_output) with the desired output
            of the network.

        :param data: Validation data as input_format.
        :return: A Dictionary with all data that needs to be feed to the network.
        '''

        x = data.images
        return {
            self._tf_input_data: x,
            self._tf_hrand: np.random.rand(x.shape[0], self._hidden_size),
            self._tf_vrand: np.random.rand(x.shape[0], self._input_size),
            self._tf_lr: self._learning_rate,
            self._tf_momentum: self._momentum_factor,
            self._tf_weight_decay: self._weight_decay_factor,
            self._tf_gibbs_steps: self._gibbs_sampling_steps
        }

    def _build_model(self):

        '''
            For Tensorflow we need to specify the graph we want to run beforehand. A need graph is build by every call
            of fit. the build model has every needed to train the network.
            It creates the placeholders and variables. Then it performs the gibbs sampling steps to update the weights.
            The steps are structured with tf.name_scope to format the graph displayed in tensorboard.
            The build function also includes a graph for the accuracy and loss function.

        :return: self.
        '''

        self._create_placeholders()

        self._create_variables()

        with tf.name_scope('learning_process'):
            with tf.name_scope('gibbs_sapmling_step'):
                hprobs0, hstates0, vprobs, hprobs1, hstates1 = self._gibbs_sampling_step(
                    self._sample_prob(self._tf_input_data, self._tf_vrand)
                )
            with tf.name_scope('_compute_positive_association'):
                positive = self._compute_positive_association(self._tf_input_data, hprobs0, hstates0)

            with tf.name_scope('gibbs_sampling_steps'):
                nn_input = vprobs

                for step in range(self._gibbs_sampling_steps - 1):
                    hprobs, hstates, vprobs, hprobs1, hstates1 = self._gibbs_sampling_step(nn_input)
                    nn_input = vprobs

            with tf.name_scope('compute_negative_association'):
                negative = tf.matmul(tf.transpose(vprobs), hprobs1)

            self.encode = hprobs1

            with tf.name_scope('calculate_delta_weights'):
                d_w = self._tf_delta_w.assign(self._tf_lr * (positive - negative)
                                              + tf.scalar_mul(self._tf_momentum, self._tf_delta_w)
                                              - tf.nn.tanh(tf.scalar_mul(self._tf_weight_decay, self._tf_w)))
                d_bh = self._tf_delta_bh.assign(self._tf_lr * tf.reduce_mean(hprobs0 - hprobs1, 0))
                d_bv = self._tf_delta_bv.assign(self._tf_lr * tf.reduce_mean(self._tf_input_data - vprobs, 0))

            with tf.name_scope('update_weights'):
                self._tf_w_upd8 = self._tf_w.assign_add(d_w)
                self._tf_bh_upd8 = self._tf_bh.assign_add(d_bh)
                self._tf_bv_upd8 = self._tf_bv.assign_add(d_bv)

        with tf.name_scope('accuracy'):
            result = self._sample_hidden_from_visible(self._tf_input_data)

            correct_prediction = tf.equal(tf.argmax(result[0], 1), tf.argmax(self._tf_desired_output, 1))
            self._tf_accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

        with tf.name_scope('loss'):
            self._tf_loss_function = tf.sqrt(tf.reduce_mean(tf.square(self._tf_input_data - vprobs)))

    def _create_placeholders(self):

        '''
            Creates all needed Placeholders in one call. The placeholders can be feed at runtime.

        :return: self.
        '''
        self._tf_input_data = tf.placeholder(tf.float32, [None, self._input_size], name="x-input")
        self._tf_hrand = tf.placeholder(tf.float32, [None, self._hidden_size], name="hrand")
        self._tf_vrand = tf.placeholder(tf.float32, [None, self._input_size], name="vrand")
        self._tf_desired_output = tf.placeholder(tf.float32, [None, self._hidden_size], name="Desired_Output")
        self._tf_lr = tf.placeholder(tf.float32, [], name="learning_rate")
        self._tf_momentum = tf.placeholder(tf.float32, [], name="momentum_term")
        self._tf_weight_decay = tf.placeholder(tf.float32, [], name="alpha")
        self._tf_gibbs_steps = tf.placeholder(tf.float32, [], name="gibbs_sampling_steps")

    def _create_variables(self):

        '''
            Creates all Variables in one function call. The variables contain the weights and weight changes.
            Tensor specified as tf.Variable are save by the tf.train.Saver after every succefull call of fit.
            If restore_previous_model is true the pre-initialised values we'll be overwritten by the previously save one.

        :return: self.
        '''
        self._tf_w = tf.Variable(
            tf.random_normal((self._input_size, self._hidden_size), mean=0.0, stddev=self._standard_derivation),
            name="weights_" + self._model_name)
        self._tf_bh = tf.Variable(tf.zeros([self._hidden_size]), name="hidden_bias_" + self._model_name)
        self._tf_bv = tf.Variable(tf.zeros([self._input_size]), name="visible_bias_" + self._model_name)

        self._tf_delta_w = tf.Variable(tf.zeros([self._input_size, self._hidden_size]), name="delta_w_" + self._model_name)
        self._tf_delta_bh = tf.Variable(tf.zeros([self._hidden_size]), name="delta_bh_" + self._model_name)
        self._tf_delta_bv = tf.Variable(tf.zeros([self._input_size]), name="delta_bv_" + self._model_name)

    def _gibbs_sampling_step(self, visible):

        '''
            Performes a complete Gibbs-Sapmling-Step.

        :param visible: input vector
        :return: probality of the hidden layer, visible layer and a new hidden layer after one down up pass. Also the
                    binary staes of the hidden and the new hidden layer.
        '''
        hprobs, hstates = self._sample_hidden_from_visible(visible)
        vprobs = self._sample_visible_from_hidden(hprobs)
        hprobs1, hstates1 = self._sample_hidden_from_visible(vprobs)

        return hprobs, hstates, vprobs, hprobs1, hstates1

    def _sample_hidden_from_visible(self, visible):

        '''
            Helper function. Performs an up-pass from the visible into the hidden layer.

        :param visible: Input vector with size of the visible layer.
        :return: The probability and the states of the hidden layer after the up-pass.
        '''
        hprobs = tf.nn.sigmoid(tf.matmul(visible, self._tf_w) + self._tf_bh)
        hstates = self._sample_prob(hprobs, self._tf_hrand)

        return hprobs, hstates

    def _sample_visible_from_hidden(self, hidden):

        '''
        Helper function. Performs an down-pass from the hidden into the visible layer.

        :param hidden: Input vector with size of hidden layer.
        :return: self
        '''
        return tf.nn.sigmoid(tf.matmul(hidden, tf.transpose(self._tf_w)) + self._tf_bv)

    def _compute_positive_association(self, visible, hidden_states):

        '''
        Helper function to compute the positive association in one function call.

        :param visible: A tensor which represents the states of the visible units.
        :param hidden_states: A tensor which represents the states of the hidden units.
        :return: self
        '''

        return tf.matmul(tf.transpose(visible), hidden_states)

    def _sample_prob(self, probs, rand):

        '''
        Helperfunction to randomly convert a value into 1 or 0.

        :param probs: The probability vector of the up or down pass.
        :param rand: A Vector with the same size as probs which contains random values.
        :return: self.
        '''

        return tf.nn.relu(tf.sign(probs - rand))

    def _create_summary_nodes(self):

        '''
        This function adds additional nodes to the computation graph. These are sidenodes which get only evaluated when
                self._tf_merged_summarys runs. The summary nodes create Data which is visualized with Tensorboard.
        :return: self.
        '''

        with tf.name_scope('summaries'):
            with tf.name_scope('learning_progress'):
                tf.summary.scalar("loss", self._tf_loss_function)
            with tf.name_scope('weight_development'):
                tf.summary.scalar("max_Weight", tf.reduce_max(self._tf_w))
                tf.summary.scalar("min_weight", tf.reduce_min(self._tf_w))
                # tf.summary.histogram("delta_weights", self._tf_delta_w)
                # tf.summary.histogram("weights", self._tf_w)
            with tf.name_scope('hyperparameter'):
                tf.summary.scalar("_learning_rate", self._learning_rate)
                tf.summary.scalar("Gibbs_sampling_steps", self._gibbs_sampling_steps)

    def _create_data_directories(self):

        '''
        A Helper function to create the directories to save the logs and the model weights.

        :return: Returns the path as string for the model, data and summary directory.
        '''

        self._main_dir = self._main_dir + '/' if self._main_dir[-1] != '/' else self._main_dir

        models_dir = "models/" + self._main_dir
        data_dir = "data/" + self._main_dir
        summary_dir = "logs/" + self._main_dir

        for d in [models_dir, data_dir, summary_dir]:
            if not os.path.isdir(d):
                os.makedirs(d)

        return models_dir, data_dir, summary_dir

    def get_weights(self):

        '''
        Initializes a Session and loads the model data. The tensorflow eval function returns the variable values as
                a numpy array. The numpy array is used to initialize the DBN weights.

        :return: Return the weights and biases as numpy array
        '''

        self._create_variables()
        with tf.Session() as self._tf_session:
            self._tf_saver = tf.train.Saver()
            self._tf_saver.restore(self._tf_session, self._model_path)

            w = self._tf_w.eval()
            bh = self._tf_bh.eval()

        tf.reset_default_graph()

        return w, bh

    def get_weights_as_images(self, width, height, outdir='img/', n_images=10, img_type='grey'):

        """
        Create and save the weights of the hidden units with respect to the
                visible units as images.

        :param width: width of the image
        :param height: height of the image
        :param outdir: directory where the created images should be saved
        :param n_images: How many images should be createds
        :param img_type: Should the image be color or grey. If it should be colored for every pixel should exist 3 neurons.
        :return: self.
        """

        outdir = self._data_dir + outdir
        self._build_model()
        with tf.Session() as self._tf_session:
            self._tf_saver = tf.train.Saver()

            self._tf_saver.restore(self._tf_session, self._model_path)

            weights = self._tf_w.eval()

            perm = np.random.permutation(self._hidden_size)[:n_images]

            for p in perm:
                w = np.array([i[p] for i in weights])
                image_path = outdir + self._model_name + '_{}.png'.format(p)

                utilities.gen_image(w, width, height, image_path, img_type)
