from util import *


class RestrictedBoltzmannMachine:
    """
    For more details: A Practical Guide to Training Restricted Boltzmann Machines.
    https://www.cs.toronto.edu/~hinton/absps/guideTR.pdf
    """

    def __init__(self, ndim_visible, ndim_hidden, is_bottom=False, image_size=(28, 28), is_top=False, n_labels=10,
                 batch_size=10):
        """
        Args:
            ndim_visible: Number of units in visible layer.
            ndim_hidden: Number of units in hidden layer.
            is_bottom: True only if this rbm is at the bottom of the stack in a deep belief net. Used to interpret
                visible layer as image data with dimensions "image_size".
            image_size: Image dimension for visible layer.
            is_top: True only if this rbm is at the top of stack in deep belief net. Used to interpret visible layer as
                concatenated with "n_label" unit of label data at the end.
            n_labels: Number of label categories.
            batch_size: Size of mini-batch.
        """
        # architecture
        self.ndim_visible = ndim_visible
        self.ndim_hidden = ndim_hidden

        # position in stack (for DBN)
        self.is_bottom = is_bottom
        if is_bottom:
            self.image_size = image_size
        self.is_top = is_top
        if is_top:
            self.n_labels = n_labels

        # init parameters
        self.bias_v = np.random.normal(loc=0.0, scale=0.01, size=(self.ndim_visible,))
        self.weight_vh = np.random.normal(loc=0.0, scale=0.01, size=(self.ndim_visible, self.ndim_hidden))
        self.bias_h = np.random.normal(loc=0.0, scale=0.01, size=(self.ndim_hidden,))
        self.weight_v_to_h = None
        self.weight_h_to_v = None

        # update vectors (for momentum)
        self.delta_bias_v = 0
        self.delta_weight_vh = 0
        self.delta_bias_h = 0
        self.delta_weight_v_to_h = 0
        self.delta_weight_h_to_v = 0

        # hyper-parameters for learning
        self.batch_size = batch_size
        self.learning_rate = 0.01
        self.momentum = 0.7
        self.print_period = 500

        # receptive-fields (only applicable when visible layer is input data)
        self.rf = {
            "period": 5000,                                         # iteration period to visualize
            "grid": [5, 5],                                         # size of the grid
            "ids": np.random.randint(0, self.ndim_hidden, 25)       # pick some random hidden units
        }

        # stats
        self.reconstruction_loss = []
        self.measured_iterations = []

    def cd1(self, visible_trainset, n_iterations=10000):
        """Contrastive Divergence with k=1 full alternating Gibbs sampling

        Args:
          visible_trainset: training data for this rbm, shape is (size of training set, size of visible layer)
          n_iterations: number of iterations of learning (each iteration learns a mini-batch)
        """
        print("learning CD1")

        n_samples = visible_trainset.shape[0]
        n_batches = int(n_samples / self.batch_size)

        # shuffle
        visible_trainset = visible_trainset.copy()
        np.random.shuffle(visible_trainset)

        for it in range(n_iterations):
            # [TODO TASK 4.1] run k=1 alternating Gibbs sampling : v_0 -> h_0 ->  v_1 -> h_1.
            # select batch
            batch = it % n_batches
            batch = visible_trainset[batch*self.batch_size:(batch+1)*self.batch_size]

            # positive phase
            v_0 = batch                             # clamp inputs to visible nodes
            p_h_0, h_0 = self.get_h_given_v(v_0)

            # negative phase
            p_v_1, v_1 = self.get_v_given_h(h_0)    # binary samples for h (strong regularizer)
            p_h_1, h_1 = self.get_h_given_v(p_v_1)  # probabilities for v

            # [TODO TASK 4.1] update the parameters using function 'update_params'
            self.update_params(v_0, h_0, p_v_1, p_h_1)      # binary samples only for h_0

            # visualize once in a while when visible layer is input images
            if it % self.rf["period"] == 0 and self.is_bottom:
                viz_rf(weights=self.weight_vh[:, self.rf["ids"]].reshape((self.image_size[0], self.image_size[1], -1)),
                       it=it, grid=self.rf["grid"])

            # print progress
            if it % self.print_period == 0:
                self.reconstruction_loss.append(self.compute_reconstruction_loss(visible_trainset))
                self.measured_iterations.append(it)
                print("iteration=%07d recon_loss=%4.4f" % (it, self.reconstruction_loss[-1]))

    def update_params(self, v_0, h_0, v_k, h_k):
        """Update the weight and bias parameters.

        You could also add weight decay and momentum for weight updates.

        Args:
           v_0: activities or probabilities of visible layer (data to the rbm)
           h_0: activities or probabilities of hidden layer
           v_k: activities or probabilities of visible layer
           h_k: activities or probabilities of hidden layer
           all args have shape (size of mini-batch, size of respective layer)
        """
        n_samples = v_0.shape[0]

        # [TODO TASK 4.1] get the gradients from the arguments (replace the 0s below) and update the weight and bias parameters
        self.delta_weight_vh = self.learning_rate / n_samples * (v_0.T @ h_0 - v_k.T @ h_k) + self.momentum * self.delta_weight_vh
        self.delta_bias_v = self.learning_rate / n_samples * (v_0.sum(axis=0) - v_k.sum(axis=0)) + self.momentum * self.delta_bias_v
        self.delta_bias_h = self.learning_rate / n_samples * (h_0.sum(axis=0) - h_k.sum(axis=0)) + self.momentum * self.delta_bias_h

        self.bias_v += self.delta_bias_v
        self.weight_vh += self.delta_weight_vh
        self.bias_h += self.delta_bias_h

    def get_h_given_v(self, visible_minibatch):
        """Compute probabilities p(h|v) and activations h ~ p(h|v) 

        Uses undirected weight "weight_vh" and bias "bias_h"
        
        Args: 
           visible_minibatch: shape is (size of mini-batch, size of visible layer)
        Returns:        
           tuple ( p(h|v) , h) 
           both are shaped (size of mini-batch, size of hidden layer)
        """
        assert self.weight_vh is not None

        # [TODO TASK 4.1] compute probabilities and activations (samples from probabilities) of hidden layer (replace the zeros below)
        p_h = sigmoid(self.bias_h + visible_minibatch @ self.weight_vh)
        h = sample_binary(p_h)
        return p_h, h

    def get_v_given_h(self, hidden_minibatch):
        """Compute probabilities p(v|h) and activations v ~ p(v|h)

        Uses undirected weight "weight_vh" and bias "bias_v"
        
        Args: 
           hidden_minibatch: shape is (size of mini-batch, size of hidden layer)
        Returns:        
           tuple ( p(v|h) , v) 
           both are shaped (size of mini-batch, size of visible layer)
        """
        assert self.weight_vh is not None

        if self.is_top:
            """Here visible layer has both data and labels. Compute total input for each unit (identical for both 
            cases), and split into two parts, something like support[:, :-self.n_labels] and support[:, 
            -self.n_labels:]. Then, for both parts, use the appropriate activation function to get probabilities and 
            a sampling method to get activities. The probabilities as well as activities can then be concatenated 
            back into a normal visible layer. """

            # [TODO TASK 4.1] compute probabilities and activations (samples from probabilities) of visible layer (replace the pass below).
            support = self.bias_v + hidden_minibatch @ self.weight_vh.T

            # labels are in the last units (on the right)!
            p_v_labels = softmax(support[:, -self.n_labels:])
            v_labels = sample_categorical(p_v_labels)

            # other units are representations of data
            p_v_data = sigmoid(support[:, :-self.n_labels])
            v_data = sample_binary(p_v_data)

            p_v = np.concatenate((p_v_data, p_v_labels), axis=1)
            v = np.concatenate((v_data, v_labels), axis=1)
        else:
            # [TODO TASK 4.1] compute probabilities and activations (samples from probabilities) of visible layer (replace the pass and zeros below)
            p_v = sigmoid(self.bias_v + hidden_minibatch @ self.weight_vh.T)
            v = sample_binary(p_v)

        return p_v, v

    """ RBM as a belief layer: the functions below do not have to be changed until running a DBN """

    def untwine_weights(self):
        self.weight_v_to_h = np.copy(self.weight_vh)
        self.weight_h_to_v = np.copy(np.transpose(self.weight_vh))
        self.weight_vh = None

    def get_h_given_v_dir(self, visible_minibatch):
        """Compute probabilities p(h|v) and activations h ~ p(h|v)

        Uses directed weight "weight_v_to_h" and bias "bias_h"
        
        Args: 
           visible_minibatch: shape is (size of mini-batch, size of visible layer)
        Returns:        
           tuple ( p(h|v) , h) 
           both are shaped (size of mini-batch, size of hidden layer)
        """
        assert self.weight_v_to_h is not None

        # [TODO TASK 4.2] perform same computation as the function 'get_h_given_v' but with directed connections (replace the zeros below)
        p_h = sigmoid(self.bias_h + visible_minibatch @ self.weight_v_to_h)
        h = sample_binary(p_h)
        return p_h, h

    def get_v_given_h_dir(self, hidden_minibatch):
        """Compute probabilities p(v|h) and activations v ~ p(v|h)

        Uses directed weight "weight_h_to_v" and bias "bias_v"
        
        Args: 
           hidden_minibatch: shape is (size of mini-batch, size of hidden layer)
        Returns:        
           tuple ( p(v|h) , v) 
           both are shaped (size of mini-batch, size of visible layer)
        """
        assert self.weight_h_to_v is not None

        if self.is_top:
            """Here visible layer has both data and labels. Compute total input for each unit (identical for both 
            cases), and split into two parts, something like support[:, :-self.n_labels] and support[:, 
            -self.n_labels:]. Then, for both parts, use the appropriate activation function to get probabilities and 
            a sampling method to get activities. The probabilities as well as activities can then be concatenated 
            back into a normal visible layer. """

            # [TODO TASK 4.2] Note that even though this function performs same computation as 'get_v_given_h' but
            #  with directed connections, this case should never be executed : when the RBM is a part of a DBN and is
            #  at the top, it will not have directed connections. Appropriate code here is to raise an error (replace
            #  pass below)
            raise RuntimeError('Directed connections used in the top RBM.')
        else:
            # [TODO TASK 4.2] performs same computaton as the function 'get_v_given_h' but with directed connections (replace the pass and zeros below)
            p_v = sigmoid(self.bias_v + hidden_minibatch @ self.weight_h_to_v)
            v = sample_binary(p_v)

        return p_v, v

    def update_generate_params(self, inps, trgs, preds):
        """Update generative weight "weight_h_to_v" and bias "bias_v"
        
        Args:
           inps: activities or probabilities of input unit
           trgs: activities or probabilities of output unit (target)
           preds: activities or probabilities of output unit (prediction)
           all args have shape (size of mini-batch, size of respective layer)
        """
        n_samples = inps.shape[0]

        # [TODO TASK 4.3] find the gradients from the arguments (replace the 0s below) and update the weight and bias parameters.
        self.delta_weight_h_to_v = self.learning_rate / n_samples * inps.T @ (trgs - preds) + self.momentum * self.delta_weight_h_to_v
        self.delta_bias_v = self.learning_rate / n_samples * (trgs - preds).sum(axis=0) + self.momentum * self.delta_bias_v

        self.weight_h_to_v += self.delta_weight_h_to_v
        self.bias_v += self.delta_bias_v

    def update_recognize_params(self, inps, trgs, preds):
        """Update recognition weight "weight_v_to_h" and bias "bias_h"
        
        Args:
           inps: activities or probabilities of input unit
           trgs: activities or probabilities of output unit (target)
           preds: activities or probabilities of output unit (prediction)
           all args have shape (size of mini-batch, size of respective layer)
        """
        n_samples = inps.shape[0]

        # [TODO TASK 4.3] find the gradients from the arguments (replace the 0s below) and update the weight and bias parameters.
        self.delta_weight_v_to_h = self.learning_rate / n_samples * inps.T @ (trgs - preds) + self.momentum * self.delta_weight_v_to_h
        self.delta_bias_h += self.learning_rate / n_samples * (trgs - preds).sum(axis=0) + self.momentum * self.delta_bias_h

        self.weight_v_to_h += self.delta_weight_v_to_h
        self.bias_h += self.delta_bias_h

    def reconstruct(self, visible_data, sample_v1=False):
        """Reconstruct data. v_0 -> h_0 -> v_1."""
        if sample_v1:
            sample_v1 = 1
        else:
            sample_v1 = 0

        h = self.get_h_given_v(visible_data)[1]     # binary samples (bottleneck)
        v = self.get_v_given_h(h)[sample_v1]
        return v

    def compute_reconstruction_loss(self, visible_data):
        """Compute the average reconstruction error."""
        reconstructions = self.reconstruct(visible_data)
        n_samples = len(visible_data)
        return 1./n_samples * np.linalg.norm(visible_data - reconstructions)
