import warnings

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm, tqdm_notebook
from itertools import chain
import qucumber.cplx as cplx


__all__ = [
    "RBM_Module",
    "BinomialRBM"
]


class RBM_Module(nn.Module):

    def __init__(self, num_visible, num_hidden, zero_weights=False,
                 gpu=True, seed=1234):
        super(RBM_Module, self).__init__()
        self.num_visible = int(num_visible)
        self.num_hidden = int(num_hidden)

        if gpu and not torch.cuda.is_available():
            warnings.warn("Could not find GPU: will continue with CPU.",
                          ResourceWarning)

        self.gpu = gpu and torch.cuda.is_available()

        if seed:
            if self.gpu:
                torch.cuda.manual_seed(seed)
            else:
                torch.manual_seed(seed)

        self.device = torch.device('cuda') if self.gpu else torch.device('cpu')

        if zero_weights:
            self.weights = nn.Parameter((torch.zeros(self.num_hidden,
                                                     self.num_visible,
                                                     device=self.device,
                                                     dtype=torch.double)),
                                        requires_grad=True)
        else:

            self.weights = nn.Parameter(
                (torch.randn(self.num_hidden, self.num_visible,
                             device=self.device, dtype=torch.double)
                 / np.sqrt(self.num_visible)),
                requires_grad=True)

        self.visible_bias = nn.Parameter(torch.zeros(self.num_visible,
                                                     device=self.device,
                                                     dtype=torch.double),
                                         requires_grad=True)
        self.hidden_bias = nn.Parameter(torch.zeros(self.num_hidden,
                                                    device=self.device,
                                                    dtype=torch.double),
                                        requires_grad=True)

    def __repr__(self):
        return ("RBM_Module(num_visible={}, num_hidden={}, gpu={})"
                .format(self.num_visible, self.num_hidden, self.gpu))

    def effective_energy(self, v):
        r"""The effective energies of the given visible states.

        .. math::

            \mathcal{E}(\bm{v}) &= \sum_{j}b_j v_j
                        + \sum_{i}\log
                            \left\lbrack 1 +
                                  \exp\left(c_{i} + \sum_{j} W_{ij} v_j\right)
                            \right\rbrack

        :param v: The visible states.
        :type v: torch.doubleTensor

        :returns: The effective energies of the given visible states.
        :rtype: torch.doubleTensor
        """
        if len(v.shape) < 2:
            v = v.view(1, -1)
        visible_bias_term = torch.mv(v, self.visible_bias)
        hidden_bias_term = F.softplus(
            F.linear(v, self.weights, self.hidden_bias)
        ).sum(1)

        return visible_bias_term + hidden_bias_term

    def prob_v_given_h(self, h):
        """Given a hidden unit configuration, compute the probability
        vector of the visible units being on.

        :param h: The hidden unit
        :type h: torch.doubleTensor

        :returns: The probability of visible units being active given the
                  hidden state.
        :rtype: torch.doubleTensor
        """
        p = F.sigmoid(F.linear(h, self.weights.t(), self.visible_bias))
        return p

    def prob_h_given_v(self, v):
        """Given a visible unit configuration, compute the probability
        vector of the hidden units being on.

        :param h: The hidden unit.
        :type h: torch.doubleTensor

        :returns: The probability of hidden units being active given the
                  visible state.
        :rtype: torch.doubleTensor
        """
        p = F.sigmoid(F.linear(v, self.weights, self.hidden_bias))
        return p

    def sample_v_given_h(self, h):
        """Sample/generate a visible state given a hidden state.

        :param h: The hidden state.
        :type h: torch.doubleTensor

        :returns: Tuple containing prob_v_given_h(h) and the sampled visible
                  state.
        :rtype: tuple(torch.doubleTensor, torch.doubleTensor)
        """
        p = self.prob_v_given_h(h)
        v = p.bernoulli()
        return p, v

    def sample_h_given_v(self, v):
        """Sample/generate a hidden state given a visible state.

        :param h: The visible state.
        :type h: torch.doubleTensor

        :returns: Tuple containing prob_h_given_v(v) and the sampled hidden
                  state.
        :rtype: tuple(torch.doubleTensor, torch.doubleTensor)
        """
        p = self.prob_h_given_v(v)
        h = p.bernoulli()
        return p, h

    def gibbs_sampling(self, k, v0):
        """Performs k steps of Block Gibbs sampling given an initial visible
        state v0.

        :param k: Number of Block Gibbs steps.
        :type k: int
        :param v0: The initial visible state.
        :type v0: torch.doubleTensor

        :returns: Tuple containing the initial visible state, v0,
                  the hidden state sampled from v0,
                  the visible state sampled after k steps,
                  the hidden state sampled after k steps and its corresponding

                  probability vector.
        :rtype: tuple(torch.doubleTensor, torch.doubleTensor,
                      torch.doubleTensor, torch.doubleTensor,
                      torch.doubleTensor)
        """
        ph, h0 = self.sample_h_given_v(v0)
        v, h = v0, h0
        for _ in range(k):
            pv, v = self.sample_v_given_h(h)
            ph, h = self.sample_h_given_v(v)
        return v0, h0, v, h, ph

    def sample(self, num_samples, k=10):
        """Samples from the RBM using k steps of Block Gibbs sampling.
        :param num_samples: The number of samples to be generated
        :type num_samples: int
        :param k: Number of Block Gibbs steps.
        :type k: int
        :returns: Samples drawn from the RBM
        :rtype: torch.doubleTensor
        """
        dist = torch.distributions.bernoulli.Bernoulli(probs=0.5)
        v0 = (dist.sample(torch.Size([num_samples, self.num_visible]))
                  .to(device=self.device, dtype=torch.double))
        _, _, v, _, _ = self.gibbs_sampling(k, v0)
        return v

    def unnormalized_probability(self, v):
        r"""The unnormalized probabilities of the given visible states.

        .. math:: p(\bm{v}) = e^{\mathcal{E}(\bm{v})}

        :param v: The visible states.
        :type v: torch.doubleTensor

        :returns: The unnormalized probability of the given visible state(s).
        :rtype: torch.doubleTensor
        """
        return self.effective_energy(v).exp()

    def generate_visible_space(self):
        """Generates all possible visible states.

        :returns: A tensor of all possible spin configurations.
        :rtype: torch.doubleTensor
        """
        space = torch.zeros((1 << self.num_visible, self.num_visible),
                            device=self.device, dtype=torch.double)
        for i in range(1 << self.num_visible):
            d = i
            for j in range(self.num_visible):
                d, r = divmod(d, 2)
                space[i, self.num_visible - j - 1] = int(r)

        return space

    def log_partition(self, visible_space):
        """The natural logarithm of the partition function of the RBM.

        :param visible_space: A rank 2 tensor of the entire visible space.
        :type visible_space: torch.doubleTensor

        :returns: The natural log of the partition function.
        :rtype: torch.doubleTensor
        """
        free_energies = self.effective_energy(visible_space)
        max_free_energy = free_energies.max()

        f_reduced = free_energies - max_free_energy
        logZ = max_free_energy + f_reduced.exp().sum().log()

        return logZ

    def partition(self, visible_space):
        """The partition function of the RBM.

        :param visible_space: A rank 2 tensor of the entire visible space.
        :type visible_space: torch.doubleTensor

        :returns: The partition function.
        :rtype: torch.doubleTensor
        """
        return self.log_partition(visible_space).exp()

    def probability(self, v, Z):
        """Evaluates the probability of the given vector(s) of visible
        units; NOT RECOMMENDED FOR RBMS WITH A LARGE # OF VISIBLE UNITS

        :param v: The visible states.
        :type v: torch.doubleTensor
        :param Z: The partition function.
        :type Z: float

        :returns: The probability of the given vector(s) of visible units.
        :rtype: torch.doubleTensor
        """
        return self.unnormalized_probability(v) / Z


class BinomialRBM(nn.Module):
    def __init__(self, num_visible, num_hidden=None, gpu=True, seed=1234):
        super(BinomialRBM, self).__init__()
        self.num_visible = int(num_visible)
        self.num_hidden = int(num_hidden) if num_hidden is not None else self.num_visible
        self.rbm_module = RBM_Module(self.num_visible, self.num_hidden,
                                     gpu=gpu, seed=seed)
        self.stop_training = False

    def save(self, location, metadata={}):
        """Saves the RBM parameters to the given location along with
        any given metadata.

        :param location: The location to save the RBM parameters + metadata
        :type location: str or file
        :param metadata: Any extra metadata to store alongside the RBM
                         parameters
        :type metadata: dict
        """
        # add extra metadata to dictionary before saving it to disk
        data = {**self.state_dict(), **metadata}
        torch.save(data, location)

    def load(self, location):
        """Loads the RBM parameters from the given location ignoring any
        metadata stored in the file. Overwrites the RBM's parameters.

        .. note::
            The RBM object on which this function is called must
            have the same shape as the one who's parameters are being
            loaded.

        :param location: The location to load the RBM parameters from
        :type location: str or file
        """
        self.load_state_dict(torch.load(location), strict=False)

    def compute_batch_gradients(self, k, pos_batch, neg_batch):
        """This function will compute the gradients of a batch of the training
        data (data_file) given the basis measurements (chars_file).

        :param k: Number of contrastive divergence steps in training.
        :type k: int
        :param pos_batch: Batch of the input data for the positive phase.
        :type pos_batch: |DoubleTensor|
        :param neg_batch: Batch of the input data for the negative phase.
        :type neg_batch: |DoubleTensor|

        :returns: Dictionary containing all the gradients of the parameters.
        :rtype: dict
        """

        v0, _, _, _, _ = self.rbm_module.gibbs_sampling(k, pos_batch)
        _, _, vk, hk, phk = self.rbm_module.gibbs_sampling(k, neg_batch)

        prob = F.sigmoid(F.linear(v0, self.rbm_module.weights,
                                  self.rbm_module.hidden_bias))

        pos_batch_size = float(len(pos_batch))
        neg_batch_size = float(len(neg_batch))

        w_grad = torch.einsum("ij,ik->jk", (prob, v0))/pos_batch_size
        vb_grad = torch.einsum("ij->j", (v0,))/pos_batch_size
        hb_grad = torch.einsum("ij->j", (prob,))/pos_batch_size

        w_grad -= torch.einsum("ij,ik->jk", (phk, vk))/neg_batch_size
        vb_grad -= torch.einsum("ij->j", (vk,))/neg_batch_size
        hb_grad -= torch.einsum("ij->j", (phk,))/neg_batch_size

        # Return negative gradients to match up nicely with the usual
        # parameter update rules, which *subtract* the gradient from
        # the parameters. This is in contrast with the RBM update
        # rules which ADD the gradients (scaled by the learning rate)
        # to the parameters.

        return {"rbm_module": {"weights": -w_grad,
                               "visible_bias": -vb_grad,
                               "hidden_bias": -hb_grad}}

    def fit(self, data, epochs=100, pos_batch_size=100, neg_batch_size=200,
            k=1, lr=1e-2, progbar=False, callbacks=[]):
        """Execute the training of the RBM.

        :param data: The actual training data
        :type data: list(float)
        :param epochs: The number of parameter (i.e. weights and biases)
                       updates
        :type epochs: int
        :param pos_batch_size: The size of batches for the positive phase
                               taken from the data.
        :type pos_batch_size: int
        :param neg_batch_size: The size of batches for the negative phase
                               taken from the data
        :type neg_batch_size: int
        :param k: The number of contrastive divergence steps
        :type k: int
        :param lr: Learning rate
        :type lr: float
        :param progbar: Whether or not to display a progress bar. If "notebook"
                        is passed, will use a Jupyter notebook compatible
                        progress bar.
        :type progbar: bool or str
        :param callbacks: Callbacks to run while training.
        :type callbacks: list(qucumber.callbacks.Callback)
        """

        disable_progbar = (progbar is False)
        progress_bar = tqdm_notebook if progbar == "notebook" else tqdm

        data = torch.tensor(data, device=self.rbm_module.device,
                            dtype=torch.double)
        optimizer = torch.optim.SGD([self.rbm_module.weights,
                                     self.rbm_module.visible_bias,
                                     self.rbm_module.hidden_bias],
                                    lr=lr)

        for cb in callbacks:
            cb.on_train_start(self)

        for ep in progress_bar(range(epochs), desc="Epochs ",
                               disable=disable_progbar):
            pos_batches = DataLoader(data, batch_size=pos_batch_size,
                                     shuffle=True)

            multiplier = int((neg_batch_size / pos_batch_size) + 0.5)
            neg_batches = [DataLoader(data, batch_size=neg_batch_size,
                                      shuffle=True)
                           for i in range(multiplier)]
            neg_batches = chain(*neg_batches)

            for cb in callbacks:
                cb.on_epoch_start(self, ep)

            if self.stop_training:  # check for stop_training signal
                break

            for batch_num, (pos_batch, neg_batch) in enumerate(zip(pos_batches,
                                                               neg_batches)):
                for cb in callbacks:
                    cb.on_batch_start(self, ep, batch_num)

                all_grads = self.compute_batch_gradients(k, pos_batch,
                                                         neg_batch)
                optimizer.zero_grad()  # clear any cached gradients

                # assign all available gradients to the corresponding parameter
                for name, grads in all_grads.items():
                    selected_RBM = getattr(self, name)
                    for param in grads.keys():
                        getattr(selected_RBM, param).grad = grads[param]

                optimizer.step()  # tell the optimizer to apply the gradients

                for cb in callbacks:
                    cb.on_batch_end(self, ep, batch_num)

            for cb in callbacks:
                cb.on_epoch_end(self, ep)

        for cb in callbacks:
            cb.on_train_end(self)

    def sample(self, num_samples, k):
        """Samples from the RBM using k steps of Block Gibbs sampling.
        :param num_samples: The number of samples to be generated
        :type num_samples: int
        :param k: Number of Block Gibbs steps.
        :type k: int
        :returns: Samples drawn from the RBM.
        :rtype: torch.doubleTensor
        """
        return self.rbm_module.sample(num_samples, k)


class ComplexRBM:
    # NOTE: In development. Might be unstable.
    # NOTE: The 'full_unitaries' argument is not needed for training/sampling.
    # This is only here for debugging the gradients. Delete this argument when
    # gradients have been debugged.
    def __init__(self, full_unitaries, psi_dictionary, num_visible,
                 num_hidden_amp, num_hidden_phase, test_grads, gpu=True,
                 seed=1234):

        self.num_visible = int(num_visible)
        self.num_hidden_amp = int(num_hidden_amp)
        self.num_hidden_phase = int(num_hidden_phase)
        self.full_unitaries = full_unitaries
        self.psi_dictionary = psi_dictionary
        self.rbm_amp = RBM_Module(num_visible, num_hidden_amp, gpu=gpu,
                                  seed=seed)
        self.rbm_phase = RBM_Module(num_visible, num_hidden_phase,
                                    zero_weights=True, gpu=gpu, seed=None)
        self.device = self.rbm_amp.device
        self.test_grads = test_grads

    def basis_state_generator(self, s):
        """Only works for binary visible units at the moment. Generates a vector
        given a spin value (0 or 1).

        :param s: A spin's value (either 0 or 1).
        :type s: float

        :returns: If s = 0, this is the (1,0) state in the basis of the
                  measurement. If s = 1, this is the (0,1) state in the basis
                  of the measurement.
        :rtype: torch.doubleTensor
        """
        if s == 0.:
            return torch.tensor([[1., 0.], [0., 0.]], dtype=torch.double)
        if s == 1.:
            return torch.tensor([[0., 1.], [0., 0.]], dtype=torch.double)

    def state_generator(self, num_non_trivial_unitaries):
        """A function that returns all possible configurations of
        'num_non_trivial_unitaries' spins. Similar to generate_visible_space.

        :param num_non_trivial_unitaries: The number of sites measured in the
                                          non-computational basis.
        :type num_non_trivial_unitaries: int

        :returns: An array of all possible spin configurations of
                  'num_non_trivial_unitaries' spins.
        :rtype: torch.doubleTensor
        """
        states = torch.zeros((2**num_non_trivial_unitaries,
                              num_non_trivial_unitaries),
                             device=self.device,
                             dtype=torch.double)
        for i in range(2**num_non_trivial_unitaries):
            temp = i
            for j in range(num_non_trivial_unitaries):
                temp, remainder = divmod(temp, 2)
                states[i][num_non_trivial_unitaries - j - 1] = remainder
        return states

    def unnormalized_probability_amp(self, v):
        r"""The effective energy of the phase RBM.

        :param v: Visible unit(s).
        :type v: torch.doubleTensor

        :returns:
            :math:`p_{\lambda}(\bm{v}) = e^{\mathcal{E}_{\lambda}(\bm{v})}`
        :rtype: torch.doubleTensor
        """
        return self.rbm_amp.unnormalized_probability(v)

    def unnormalized_probability_phase(self, v):
        r"""The effective energy of the phase RBM.

        :param v: Visible unit(s).
        :type v: torch.doubleTensor

        :returns: :math:`p_{\mu}(\bm{v}) = e^{\mathcal{E}_{\mu}(\bm{v})}`
        :rtype: torch.doubleTensor
        """
        return self.rbm_phase.unnormalized_probability(v)

    def normalized_wavefunction(self, v):
        r"""The RBM wavefunction.

        :param v: Visible unit(s).
        :type v: torch.doubleTensor

        :returns:

        .. math:: \psi_{\lambda\mu} =
                    \sqrt{\frac{p_{\lambda}}{Z_{\lambda}}}
                    \exp\left(\frac{i\log(p_{\mu})}{2}\right)

        :rtype: torch.doubleTensor
        """
        v_prime = v.view(-1, self.num_visible)
        temp1 = (self.unnormalized_probability_amp(v_prime)).sqrt()
        temp2 = ((self.unnormalized_probability_phase(v_prime)).log())*0.5

        cos_angle = temp2.cos()
        sin_angle = temp2.sin()

        psi = torch.zeros(2, v_prime.size()[0], dtype=torch.double)
        psi[0] = temp1*cos_angle
        psi[1] = temp1*sin_angle

        sqrt_Z = (self.rbm_amp.partition(
            self.rbm_amp.generate_visible_space())).sqrt()

        return psi / sqrt_Z

    def unnormalized_wavefunction(self, v):
        r"""The unnormalized RBM wavefunction.

        :param v: Visible unit(s).
        :type v: torch.doubleTensor

        :returns:

        .. math:: \tilde{\psi}_{\lambda\mu} =
                        \sqrt{p_{\lambda}}
                        \exp\left(\frac{i\log(p_{\mu})}{2}\right)

        :rtype: torch.doubleTensor
        """
        v_prime = v.view(-1, self.num_visible)
        temp1 = (self.unnormalized_probability_amp(v_prime)).sqrt()
        temp2 = ((self.unnormalized_probability_phase(v_prime)).log())*0.5
        cos_angle = temp2.cos()
        sin_angle = temp2.sin()

        psi = torch.zeros(2, v_prime.size()[0], dtype=torch.double)

        psi[0] = temp1*cos_angle
        psi[1] = temp1*sin_angle

        return psi

    def compute_batch_gradients(self, unitary_dict, k, batch, chars_batch):
        """This function will compute the gradients of a batch of the training
        data (data_file) given the basis measurements (chars_file).

        :param k: Number of contrastive divergence steps in amplitude training.
        :type k: int
        :param batch: Batch of the input data.
        :type batch: torch.doubleTensor
        :param chars_batch: Batch of bases that correspondingly indicates the
                            basis each site in the batch was measured in.
        :type chars_batch: list(str)

        :returns: Dictionary containing all the gradients (negative): Gradient
                  of weights, visible bias and hidden bias for the amplitude,
                  Gradients of weights, visible bias and hidden bias for the
                  phase.
        :rtype: dict
        """

        batch_size = len(batch)

        g_weights_amp = torch.zeros_like(self.rbm_amp.weights)
        g_vb_amp = torch.zeros_like(self.rbm_amp.visible_bias)
        g_hb_amp = torch.zeros_like(self.rbm_amp.hidden_bias)

        g_weights_phase = torch.zeros_like(self.rbm_phase.weights)
        g_vb_phase = torch.zeros_like(self.rbm_phase.visible_bias)
        g_hb_phase = torch.zeros_like(self.rbm_phase.hidden_bias)

        [batch, h0_amp_batch, vk_amp_batch, hk_amp_batch, phk_amp_batch] = \
            self.rbm_amp.gibbs_sampling(k, batch)
        # Iterate through every data point in the batch.
        for row_count, v0 in enumerate(batch):

            # A counter for the number of non-trivial unitaries
            # (non-computational basis) in the data point.
            num_non_trivial_unitaries = 0

            # tau_indices will contain the index numbers of spins not in the
            # computational basis (Z). z_indices will contain the index numbers
            # of spins in the computational basis.
            tau_indices = []
            z_indices = []

            for j in range(self.num_visible):
                # Go through the unitaries (chars_batch[row_count]) of each
                # site in the data point, v0, and save inidices of non-trivial.
                if chars_batch[row_count][j] != 'Z':
                    num_non_trivial_unitaries += 1
                    tau_indices.append(j)
                else:
                    z_indices.append(j)

            if num_non_trivial_unitaries == 0:
                # If there are no non-trivial unitaries for the data point v0,
                # calculate the positive phase of regular (i.e. non-complex
                # RBM) gradient. Use the actual data point, v0.
                prob_amp = F.sigmoid(F.linear(v0, self.rbm_amp.weights,
                                              self.rbm_amp.hidden_bias))
                g_weights_amp -= (torch.einsum("i,j->ij", (prob_amp, v0))
                                  / batch_size)
                g_vb_amp -= v0 / batch_size
                g_hb_amp -= prob_amp / batch_size

            else:
                # Compute the rotated gradients.
                [L_weights_amp, L_vb_amp, L_hb_amp,
                 L_weights_phase, L_vb_phase, L_hb_phase] = \
                    self.compute_rotated_grads(unitary_dict, k, v0,
                                               chars_batch[row_count],
                                               num_non_trivial_unitaries,
                                               z_indices, tau_indices)

                # Gradents of amplitude parameters take the real part of the
                # rotated gradients.
                g_weights_amp -= L_weights_amp[0] / batch_size
                g_vb_amp -= L_vb_amp[0] / batch_size
                g_hb_amp -= L_hb_amp[0] / batch_size

                # Gradents of phase parameters take the imaginary part of the
                # rotated gradients.
                g_weights_phase += L_weights_phase[1] / batch_size
                g_vb_phase += L_vb_phase[1] / batch_size
                g_hb_phase += L_hb_phase[1] / batch_size

        # Block gibbs sampling for negative phase.
        g_weights_amp += (torch.einsum("ij,ik->jk",
                                       (phk_amp_batch, vk_amp_batch))
                          / batch_size)
        g_vb_amp += torch.einsum("ij->j", (vk_amp_batch,)) / batch_size
        g_hb_amp += torch.einsum("ij->j", (phk_amp_batch,)) / batch_size

        """Return negative gradients to match up nicely with the usual
        parameter update rules, which *subtract* the gradient from
        the parameters. This is in contrast with the RBM update
        rules which ADD the gradients (scaled by the learning rate)
        to the parameters."""

        return {
            "rbm_amp": {
                "weights": g_weights_amp,
                "visible_bias": g_vb_amp,
                "hidden_bias": g_hb_amp
            },
            "rbm_phase": {
                "weights": g_weights_phase,
                "visible_bias": g_vb_phase,
                "hidden_bias": g_hb_phase
            }
        }

    def compute_rotated_grads(self, unitary_dict, k, v0, characters,
                              num_non_trivial_unitaries,
                              z_indices, tau_indices):
        """Computes the rotated gradients.

        :param v0: A visible unit.
        :type v0: torch.doubleTensor
        :param characters: A string of characters corresponding to the basis
                           that each site in v0 was measured in.
        :type characters: str
        :param num_non_trivial_unitaries: The number of sites in v0 that aren't
                                          measured in the computational basis.
        :type num_non_trivial_unitaries: int
        :param z_indices: A list of indices that correspond to sites of v0 that
                          are measured in the computational basis.
        :type z_indices: list(int)
        :param tau_indices: A list of indices that correspond to sites of v0
                            that are not measured in the computational basis.
        :type tau_indices: list(int)

        :returns: Dictionary of the rotated gradients: L_weights_amp, L_vb_amp,
                  L_hb_amp, L_weights_phase, L_vb_phase, L_hb_phase
        :rtype: dict
        """
        """Initialize the 'A' parameters (see alg 4.2)."""
        A_weights_amp = torch.zeros(2, self.rbm_amp.weights.size()[0],
                                    self.rbm_amp.weights.size()[1],
                                    device=self.device, dtype=torch.double)
        A_vb_amp = torch.zeros(2, self.rbm_amp.visible_bias.size()[0],
                               device=self.device, dtype=torch.double)
        A_hb_amp = torch.zeros(2, self.rbm_amp.hidden_bias.size()[0],
                               device=self.device, dtype=torch.double)

        A_weights_phase = torch.zeros(2, self.rbm_phase.weights.size()[0],
                                      self.rbm_phase.weights.size()[1],
                                      device=self.device, dtype=torch.double)
        A_vb_phase = torch.zeros(2, self.rbm_phase.visible_bias.size()[0],
                                 device=self.device, dtype=torch.double)
        A_hb_phase = torch.zeros(2, self.rbm_phase.hidden_bias.size()[0],
                                 device=self.device, dtype=torch.double)
        # 'B' will contain the coefficients of the rotated unnormalized
        # wavefunction.
        B = torch.zeros(2, device=self.device, dtype=torch.double)

        w_grad_amp = torch.zeros_like(self.rbm_amp.weights)
        vb_grad_amp = torch.zeros_like(self.rbm_amp.visible_bias)
        hb_grad_amp = torch.zeros_like(self.rbm_amp.hidden_bias)

        w_grad_phase = torch.zeros_like(self.rbm_phase.weights)
        vb_grad_phase = torch.zeros_like(self.rbm_phase.visible_bias)
        hb_grad_phase = torch.zeros_like(self.rbm_phase.hidden_bias)

        zeros_for_w_amp = torch.zeros_like(w_grad_amp)
        zeros_for_w_phase = torch.zeros_like(w_grad_phase)
        zeros_for_vb = torch.zeros_like(vb_grad_amp)
        zeros_for_hb_amp = torch.zeros_like(hb_grad_amp)
        zeros_for_hb_phase = torch.zeros_like(hb_grad_phase)

        # Loop over Hilbert space of the non trivial unitaries to build
        # the state.
        for j in range(2**num_non_trivial_unitaries):
            s = self.state_generator(num_non_trivial_unitaries)[j]
            # Creates a matrix where the jth row is the desired state, |S>,
            # a vector.

            # This is the sigma state.
            constructed_state = torch.zeros(
                self.num_visible, dtype=torch.double)

            U = torch.tensor([1., 0.], dtype=torch.double, device=self.device)

            # Populate the |sigma> state (aka constructed_state) accirdingly.
            for index in range(len(z_indices)):
                # These are the sites in the computational basis.
                constructed_state[z_indices[index]] = v0[z_indices[index]]

            for index in range(len(tau_indices)):
                # These are the sites that are NOT in the computational basis.
                constructed_state[tau_indices[index]] = s[index]

                aa = unitary_dict[characters[tau_indices[index]]]
                bb = self.basis_state_generator(v0[tau_indices[index]])
                cc = self.basis_state_generator(s[index])

                temp = cplx.inner_prod(cplx.MV_mult(
                    cplx.compT_matrix(aa), bb), cc)

                U = cplx.scalar_mult(U, temp)

            # Positive phase gradients for phase and amp. Will be added into
            # the 'A' parameters.

            prob_amp = F.sigmoid(F.linear(constructed_state,
                                          self.rbm_amp.weights,
                                          self.rbm_amp.hidden_bias))
            prob_phase = F.sigmoid(F.linear(constructed_state,
                                            self.rbm_phase.weights,
                                            self.rbm_phase.hidden_bias))

            w_grad_amp = torch.einsum("i,j->ij", (prob_amp, constructed_state))
            vb_grad_amp = constructed_state
            hb_grad_amp = prob_amp

            w_grad_phase = torch.einsum("i,j->ij",
                                        (prob_phase, constructed_state))
            vb_grad_phase = constructed_state
            hb_grad_phase = prob_phase

            """
            In order to calculate the 'A' parameters below with the current
            complex library, I need to make the weights and biases complex.
            I fill the complex parts of the parameters with a tensor of zeros.
            """
            temp_w_grad_amp = cplx.make_complex_matrix(w_grad_amp,
                                                       zeros_for_w_amp)
            temp_vb_grad_amp = cplx.make_complex_vector(vb_grad_amp,
                                                        zeros_for_vb)
            temp_hb_grad_amp = cplx.make_complex_vector(hb_grad_amp,
                                                        zeros_for_hb_amp)

            temp_w_grad_phase = cplx.make_complex_matrix(w_grad_phase,
                                                         zeros_for_w_phase)
            temp_vb_grad_phase = cplx.make_complex_vector(vb_grad_phase,
                                                          zeros_for_vb)
            temp_hb_grad_phase = cplx.make_complex_vector(hb_grad_phase,
                                                          zeros_for_hb_phase)

            # Temp = U*psi(sigma)
            temp = cplx.scalar_mult(
                U, self.unnormalized_wavefunction(constructed_state))

            A_weights_amp += cplx.MS_mult(temp, temp_w_grad_amp)
            A_vb_amp += cplx.VS_mult(temp, temp_vb_grad_amp)
            A_hb_amp += cplx.VS_mult(temp, temp_hb_grad_amp)

            A_weights_phase += cplx.MS_mult(temp, temp_w_grad_phase)
            A_vb_phase += cplx.VS_mult(temp, temp_vb_grad_phase)
            A_hb_phase += cplx.VS_mult(temp, temp_hb_grad_phase)

            # Rotated wavefunction.
            B += temp

        L_weights_amp = cplx.MS_divide(A_weights_amp, B)
        L_vb_amp = cplx.VS_divide(A_vb_amp, B)
        L_hb_amp = cplx.VS_divide(A_hb_amp, B)

        L_weights_phase = cplx.MS_divide(A_weights_phase, B)
        L_vb_phase = cplx.VS_divide(A_vb_phase, B)
        L_hb_phase = cplx.VS_divide(A_hb_phase, B)

        return [L_weights_amp, L_vb_amp, L_hb_amp, L_weights_phase, L_vb_phase,
                L_hb_phase]

    def fit(self, data, character_data, unitary_dict, epochs, batch_size,
            k=1, lr=1e-2, log_every=0, progbar=False):
        """Execute the training of the RBM.

        :param data: The actual training data
        :type data: list(float)
        :param character_data: The corresponding bases that each site in the
                               data has been measured in.
        :type character_data: list(str)
        :param epochs: The number of parameter (i.e. weights and biases)
                       updates
        :type epochs: int
        :param batch_size: The size of batches taken from the data
        :type batch_size: int
        :param k: The number of contrastive divergence steps
        :type k: int
        :param lr: Learning rate
        :type lr: float
        :param progbar: Whether or not to display a progress bar. If "notebook"
                        is passed, will use a Jupyter notebook compatible
                        progress bar.
        :type progbar: bool or str
        """
        # Make data file into a torch tensor.
        data = torch.tensor(data, dtype=torch.double).to(device=self.device)

        # Use the Adam optmizer to update the weights and biases.
        optimizer = torch.optim.Adam([self.rbm_amp.weights,
                                      self.rbm_amp.visible_bias,
                                      self.rbm_amp.hidden_bias,
                                      self.rbm_phase.weights,
                                      self.rbm_phase.visible_bias,
                                      self.rbm_phase.hidden_bias],
                                     lr=lr)

        disable_progbar = (progbar is False)
        progress_bar = tqdm_notebook if progbar == "notebook" else tqdm

        vis = self.rbm_amp.generate_visible_space()

        for ep in progress_bar(range(0, epochs + 1),
                               desc="Epochs ", total=epochs,
                               disable=disable_progbar):
            # Shuffle the data to ensure that the batches taken from the data
            # are random data points.
            random_permutation = torch.randperm(data.shape[0])

            shuffled_data = data[random_permutation]
            shuffled_character_data = character_data[random_permutation]

            # List of all the batches.
            batches = [shuffled_data[batch_start:(batch_start + batch_size)]
                       for batch_start in range(0, len(data), batch_size)]

            # List of all the bases.
            char_batches = [shuffled_character_data[batch_start:(batch_start + batch_size)]
                            for batch_start in range(0, len(data), batch_size)]

            # Calculate convergence quantities every "log-every" steps.
            if ep % log_every == 0:
                fidelity_ = self.fidelity(vis, 'Z' 'Z')
                print ('Epoch = ',ep,'\nFidelity = ',fidelity_)

            # Save parameters at the end of training.
            if ep == epochs:
                print('Finished training. Saving results...')
                self.save_params()
                print('Done.')
                break

            # Loop through all of the batches and calculate the batch
            # gradients.
            for index, batch in progress_bar(enumerate(batches),
                                             desc="Batches",
                                             leave=False, disable=True):

                all_grads = self.compute_batch_gradients(unitary_dict, k, batch,
                                                         char_batches[index])

                if self.test_grads:
                    self.test_gradients(unitary_dict, vis, k, batches[index],
                                        char_batches[index], all_grads)

                # Clear any cached gradients.
                optimizer.zero_grad()

                # Assign all available gradients to the
                # corresponding parameter.
                for name, grads in all_grads.items():
                    selected_RBM = getattr(self, name)
                    for param in grads.keys():
                        getattr(selected_RBM, param).grad = grads[param]

                # Tell the optimizer to apply the gradients and update
                # the parameters.
                optimizer.step()

    def save_params(self):
        """A function that saves the weights and biases for the amplitude and
        phase individually."""
        trained_params = [self.rbm_amp.weights.data.numpy(),
                          self.rbm_amp.visible_bias.data.numpy(),
                          self.rbm_amp.hidden_bias.data.numpy(),
                          self.rbm_phase.weights.data.numpy(),
                          self.rbm_phase.visible_bias.data.numpy(),
                          self.rbm_phase.hidden_bias.data.numpy()]

        with open('trained_weights_amp.csv', 'w') as csvfile:
            np.savetxt(csvfile, trained_params[0], fmt='%.5f', delimiter=',')

        with open('trained_visible_bias_amp.csv', 'w') as csvfile:
            np.savetxt(csvfile, trained_params[1], fmt='%.5f', delimiter=',')

        with open('trained_hidden_bias_amp.csv', 'w') as csvfile:
            np.savetxt(csvfile, trained_params[2], fmt='%.5f', delimiter=',')

        with open('trained_weights_phase.csv', 'w') as csvfile:
            np.savetxt(csvfile, trained_params[3], fmt='%.5f', delimiter=',')

        with open('trained_visible_bias_phase.csv', 'w') as csvfile:
            np.savetxt(csvfile, trained_params[4], fmt='%.5f', delimiter=',')

        with open('trained_hidden_bias_phase.csv', 'w') as csvfile:
            np.savetxt(csvfile, trained_params[5], fmt='%.5f', delimiter=',')

    def get_true_psi(self, basis):
        """Picks out the true psi in the correct basis.

        :param basis: E.g. XZZZX.
        :type basis: str

        :returns: The true wavefunction in the basis.
        :rtype: torch.doubleTensor
        """
        key = ''
        for i in range(len(basis)):
            key += basis[i]
        return self.psi_dictionary[key]

    def overlap(self, visible_space, basis):
        """Computes the overlap between the RBM and true wavefunctions.

        :param visible_space: An array of all possible spin configurations.
        :type visible_space: torch.doubleTensor
        :param basis: E.g. XZZZX.
        :type basis: str

        :returns: :math:`O = \\langle{\\psi_{true}}\\vert\\psi_{\\lambda\\mu}\\rangle`.
        :rtype: float
        """
        overlap_ = cplx.inner_prod(self.get_true_psi(basis),
                           self.normalized_wavefunction(visible_space))
        return overlap_

    def fidelity(self, visible_space, basis):
        """Computed the fidelity of the RBM and true wavefunctions.

        :param visible_space: An array of all possible spin configurations.
        :type visible_space: torch.doubleTensor
        :param basis: E.g. XZZZX.
        :type basis: str

        :returns: :math:`F = |O|^2`.
        :rtype: float
        """
        return cplx.norm(self.overlap(visible_space, basis))

    def KL_divergence(self, visible_space):
        '''Computes the total KL divergence.
        '''
        KL = 0.0
        basis_list = ['Z' 'Z', 'X' 'Z', 'Z' 'X', 'Y' 'Z', 'Z' 'Y']

        '''Wavefunctions (RBM and true) in the computational basis.'''
        # psi_ZZ      = self.normalized_wavefunction(visible_space)
        # true_psi_ZZ = self.get_true_psi('ZZ')

        '''Compute the KL divergence for the non computational bases.'''
        for i in range(len(basis_list)):
            rotated_RBM_psi = cplx.MV_mult(
                self.full_unitaries[basis_list[i]],
                self.normalized_wavefunction(visible_space))
            rotated_true_psi = self.get_true_psi(basis_list[i])

            for j in range(len(visible_space)):
                elementof_rotated_RBM_psi = torch.tensor(
                    [rotated_RBM_psi[0][j], rotated_RBM_psi[1][j]]).view(2, 1)
                elementof_rotated_true_psi = (torch.tensor(
                    [rotated_true_psi[0][j], rotated_true_psi[1][j]]
                ).view(2, 1))

                norm_true_psi = cplx.norm(cplx.inner_prod(
                    elementof_rotated_true_psi, elementof_rotated_true_psi))
                norm_RBM_psi = cplx.norm(cplx.inner_prod(
                    elementof_rotated_RBM_psi, elementof_rotated_RBM_psi))

                if norm_true_psi > 0.0:
                    KL += norm_true_psi*torch.log(norm_true_psi)
                if norm_RBM_psi > 0.0:
                    KL -= norm_true_psi*torch.log(norm_RBM_psi)

        '''Compute KL divergence for the computational basis.'''
        '''
        for j in range(len(visible_space)):

            elementof_ZZ_RBM_psi  = torch.tensor([psi_ZZ[0][j], psi_ZZ[1][j]]).view(2,1)
            elementof_ZZ_true_psi = torch.tensor([true_psi_ZZ[0][j], true_psi_ZZ[1][j]]).view(2,1)

            norm_ZZ_true_psi = cplx.norm( cplx.inner_prod(elementof_ZZ_true_psi, elementof_ZZ_true_psi) )
            norm_ZZ_RBM_psi  = cplx.norm( cplx.inner_prod(elementof_ZZ_RBM_psi, elementof_ZZ_RBM_psi) )

            if norm_ZZ_true_psi > 0.0:
                KL += norm_ZZ_true_psi*torch.log(norm_ZZ_true_psi)

            KL -= norm_ZZ_true_psi*torch.log(norm_ZZ_RBM_psi)
        '''
        return KL

    def compute_numerical_gradient(self, visible_space, param, alg_grad):
        eps = 1.e-6
        print("Numerical\t Exact\t\t Abs. Diff.")

        for i in range(len(param)):

            param[i].data += eps
            KL_pos = self.KL_divergence(visible_space)

            param[i].data -= 2*eps
            KL_neg = self.KL_divergence(visible_space)

            param[i].data += eps

            num_grad = (KL_pos - KL_neg) / (2*eps)

            print("{: 10.8f}\t{: 10.8f}\t{: 10.8f}\t"
                  .format(num_grad, alg_grad[i],
                          abs(num_grad - alg_grad[i])))

    def test_gradients(self, unitary_dict, visible_space, k, batch, chars_batch, alg_grads):
        # Must have negative sign because the compute_batch_grads returns the neg of the grads.
        # key_list = ["weights_amp", "visible_bias_amp", "hidden_bias_amp", "weights_phase", "visible_bias_phase", "hidden_bias_phase"]

        flat_weights_amp = self.rbm_amp.weights.data.view(-1)
        flat_weights_phase = self.rbm_phase.weights.data.view(-1)

        flat_grad_weights_amp = alg_grads["rbm_amp"]["weights"].view(-1)
        flat_grad_weights_phase = alg_grads["rbm_phase"]["weights"].view(-1)

        print('-------------------------------------------------------------------------------')

        print('Weights amp gradient')
        self.compute_numerical_gradient(
            visible_space, flat_weights_amp, -flat_grad_weights_amp)
        print ('\n')

        print('Visible bias amp gradient')
        self.compute_numerical_gradient(
            visible_space, self.rbm_amp.visible_bias, -alg_grads["rbm_amp"]["visible_bias"])
        print ('\n')

        print('Hidden bias amp gradient')
        self.compute_numerical_gradient(
            visible_space, self.rbm_amp.hidden_bias, -alg_grads["rbm_amp"]["hidden_bias"])
        print ('\n')

        print('Weights phase gradient')
        self.compute_numerical_gradient(
            visible_space, flat_weights_phase, -flat_grad_weights_phase)
        print ('\n')

        print('Visible bias phase gradient')
        self.compute_numerical_gradient(
            visible_space, self.rbm_phase.visible_bias, -alg_grads["rbm_phase"]["visible_bias"])
        print ('\n')

        print('Hidden bias phase gradient')
        self.compute_numerical_gradient(
            visible_space, self.rbm_phase.hidden_bias, -alg_grads["rbm_phase"]["hidden_bias"])

    def state_to_index(self, state):
        ''' Only for debugging how the unitary is applied to the unnormalized wavefunction - the 'B' term in alg 4.2.'''
        states = torch.zeros(2**self.num_visible, self.num_visible)
        npstates = states.numpy()
        npstate = state.numpy()
        for i in range(2**self.num_visible):
            temp = i

            for j in range(self.num_visible):
                temp, remainder = divmod(temp, 2)
                npstates[i][self.num_visible - j - 1] = remainder

            if np.array_equal(npstates[i], npstate):
                return i

import warnings

import numpy as np
import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm, tqdm_notebook
from itertools import chain
import qucumber.cplx as cplx


__all__ = [
    "RBM_Module",
    "BinomialRBM"
]


class RBM_Module(nn.Module):

    def __init__(self, num_visible, num_hidden, zero_weights=False,
                 gpu=True, seed=1234):
        super(RBM_Module, self).__init__()
        self.num_visible = int(num_visible)
        self.num_hidden = int(num_hidden)

        if gpu and not torch.cuda.is_available():
            warnings.warn("Could not find GPU: will continue with CPU.",
                          ResourceWarning)

        self.gpu = gpu and torch.cuda.is_available()

        if seed:
            if self.gpu:
                torch.cuda.manual_seed(seed)
            else:
                torch.manual_seed(seed)

        self.device = torch.device('cuda') if self.gpu else torch.device('cpu')

        if zero_weights:
            self.weights = nn.Parameter((torch.zeros(self.num_hidden,
                                                     self.num_visible,
                                                     device=self.device,
                                                     dtype=torch.double)),
                                        requires_grad=True)
        else:

            self.weights = nn.Parameter(
                (torch.randn(self.num_hidden, self.num_visible,
                             device=self.device, dtype=torch.double)
                 / np.sqrt(self.num_visible)),
                requires_grad=True)

        self.visible_bias = nn.Parameter(torch.zeros(self.num_visible,
                                                     device=self.device,
                                                     dtype=torch.double),
                                         requires_grad=True)
        self.hidden_bias = nn.Parameter(torch.zeros(self.num_hidden,
                                                    device=self.device,
                                                    dtype=torch.double),
                                        requires_grad=True)

    def __repr__(self):
        return ("RBM_Module(num_visible={}, num_hidden={}, gpu={})"
                .format(self.num_visible, self.num_hidden, self.gpu))

    def effective_energy(self, v):
        r"""The effective energies of the given visible states.

        .. math::

            \mathcal{E}(\bm{v}) &= \sum_{j}b_j v_j
                        + \sum_{i}\log
                            \left\lbrack 1 +
                                  \exp\left(c_{i} + \sum_{j} W_{ij} v_j\right)
                            \right\rbrack

        :param v: The visible states.
        :type v: torch.doubleTensor

        :returns: The effective energies of the given visible states.
        :rtype: torch.doubleTensor
        """
        if len(v.shape) < 2:
            v = v.view(1, -1)
        visible_bias_term = torch.mv(v, self.visible_bias)
        hidden_bias_term = F.softplus(
            F.linear(v, self.weights, self.hidden_bias)
        ).sum(1)

        return visible_bias_term + hidden_bias_term

    def prob_v_given_h(self, h):
        """Given a hidden unit configuration, compute the probability
        vector of the visible units being on.

        :param h: The hidden unit
        :type h: torch.doubleTensor

        :returns: The probability of visible units being active given the
                  hidden state.
        :rtype: torch.doubleTensor
        """
        p = F.sigmoid(F.linear(h, self.weights.t(), self.visible_bias))
        return p

    def prob_h_given_v(self, v):
        """Given a visible unit configuration, compute the probability
        vector of the hidden units being on.

        :param h: The hidden unit.
        :type h: torch.doubleTensor

        :returns: The probability of hidden units being active given the
                  visible state.
        :rtype: torch.doubleTensor
        """
        p = F.sigmoid(F.linear(v, self.weights, self.hidden_bias))
        return p

    def sample_v_given_h(self, h):
        """Sample/generate a visible state given a hidden state.

        :param h: The hidden state.
        :type h: torch.doubleTensor

        :returns: Tuple containing prob_v_given_h(h) and the sampled visible
                  state.
        :rtype: tuple(torch.doubleTensor, torch.doubleTensor)
        """
        p = self.prob_v_given_h(h)
        v = p.bernoulli()
        return p, v

    def sample_h_given_v(self, v):
        """Sample/generate a hidden state given a visible state.

        :param h: The visible state.
        :type h: torch.doubleTensor

        :returns: Tuple containing prob_h_given_v(v) and the sampled hidden
                  state.
        :rtype: tuple(torch.doubleTensor, torch.doubleTensor)
        """
        p = self.prob_h_given_v(v)
        h = p.bernoulli()
        return p, h

    def gibbs_sampling(self, k, v0):
        """Performs k steps of Block Gibbs sampling given an initial visible
        state v0.

        :param k: Number of Block Gibbs steps.
        :type k: int
        :param v0: The initial visible state.
        :type v0: torch.doubleTensor

        :returns: Tuple containing the initial visible state, v0,
                  the hidden state sampled from v0,
                  the visible state sampled after k steps,
                  the hidden state sampled after k steps and its corresponding

                  probability vector.
        :rtype: tuple(torch.doubleTensor, torch.doubleTensor,
                      torch.doubleTensor, torch.doubleTensor,
                      torch.doubleTensor)
        """
        ph, h0 = self.sample_h_given_v(v0)
        v, h = v0, h0
        for _ in range(k):
            pv, v = self.sample_v_given_h(h)
            ph, h = self.sample_h_given_v(v)
        return v0, h0, v, h, ph

    def sample(self, num_samples, k=10):
        """Samples from the RBM using k steps of Block Gibbs sampling.
        :param num_samples: The number of samples to be generated
        :type num_samples: int
        :param k: Number of Block Gibbs steps.
        :type k: int
        :returns: Samples drawn from the RBM
        :rtype: torch.doubleTensor
        """
        dist = torch.distributions.bernoulli.Bernoulli(probs=0.5)
        v0 = (dist.sample(torch.Size([num_samples, self.num_visible]))
                  .to(device=self.device, dtype=torch.double))
        _, _, v, _, _ = self.gibbs_sampling(k, v0)
        return v

    def unnormalized_probability(self, v):
        r"""The unnormalized probabilities of the given visible states.

        .. math:: p(\bm{v}) = e^{\mathcal{E}(\bm{v})}

        :param v: The visible states.
        :type v: torch.doubleTensor

        :returns: The unnormalized probability of the given visible state(s).
        :rtype: torch.doubleTensor
        """
        return self.effective_energy(v).exp()

    def generate_visible_space(self):
        """Generates all possible visible states.

        :returns: A tensor of all possible spin configurations.
        :rtype: torch.doubleTensor
        """
        space = torch.zeros((1 << self.num_visible, self.num_visible),
                            device=self.device, dtype=torch.double)
        for i in range(1 << self.num_visible):
            d = i
            for j in range(self.num_visible):
                d, r = divmod(d, 2)
                space[i, self.num_visible - j - 1] = int(r)

        return space

    def log_partition(self, visible_space):
        """The natural logarithm of the partition function of the RBM.

        :param visible_space: A rank 2 tensor of the entire visible space.
        :type visible_space: torch.doubleTensor

        :returns: The natural log of the partition function.
        :rtype: torch.doubleTensor
        """
        free_energies = self.effective_energy(visible_space)
        max_free_energy = free_energies.max()

        f_reduced = free_energies - max_free_energy
        logZ = max_free_energy + f_reduced.exp().sum().log()

        return logZ

    def partition(self, visible_space):
        """The partition function of the RBM.

        :param visible_space: A rank 2 tensor of the entire visible space.
        :type visible_space: torch.doubleTensor

        :returns: The partition function.
        :rtype: torch.doubleTensor
        """
        return self.log_partition(visible_space).exp()

    def probability(self, v, Z):
        """Evaluates the probability of the given vector(s) of visible
        units; NOT RECOMMENDED FOR RBMS WITH A LARGE # OF VISIBLE UNITS

        :param v: The visible states.
        :type v: torch.doubleTensor
        :param Z: The partition function.
        :type Z: float

        :returns: The probability of the given vector(s) of visible units.
        :rtype: torch.doubleTensor
        """
        return self.unnormalized_probability(v) / Z


class BinomialRBM(nn.Module):
    def __init__(self, num_visible, num_hidden=None, gpu=True, seed=1234):
        super(BinomialRBM, self).__init__()
        self.num_visible = int(num_visible)
        self.num_hidden = int(num_hidden) if num_hidden is not None else self.num_visible
        self.rbm_module = RBM_Module(self.num_visible, self.num_hidden,
                                     gpu=gpu, seed=seed)
        self.stop_training = False

    def save(self, location, metadata={}):
        """Saves the RBM parameters to the given location along with
        any given metadata.

        :param location: The location to save the RBM parameters + metadata
        :type location: str or file
        :param metadata: Any extra metadata to store alongside the RBM
                         parameters
        :type metadata: dict
        """
        # add extra metadata to dictionary before saving it to disk
        data = {**self.state_dict(), **metadata}
        torch.save(data, location)

    def load(self, location):
        """Loads the RBM parameters from the given location ignoring any
        metadata stored in the file. Overwrites the RBM's parameters.

        .. note::
            The RBM object on which this function is called must
            have the same shape as the one who's parameters are being
            loaded.

        :param location: The location to load the RBM parameters from
        :type location: str or file
        """
        self.load_state_dict(torch.load(location), strict=False)

    def compute_batch_gradients(self, k, pos_batch, neg_batch):
        """This function will compute the gradients of a batch of the training
        data (data_file) given the basis measurements (chars_file).

        :param k: Number of contrastive divergence steps in training.
        :type k: int
        :param pos_batch: Batch of the input data for the positive phase.
        :type pos_batch: |DoubleTensor|
        :param neg_batch: Batch of the input data for the negative phase.
        :type neg_batch: |DoubleTensor|

        :returns: Dictionary containing all the gradients of the parameters.
        :rtype: dict
        """

        v0, _, _, _, _ = self.rbm_module.gibbs_sampling(k, pos_batch)
        _, _, vk, hk, phk = self.rbm_module.gibbs_sampling(k, neg_batch)

        prob = F.sigmoid(F.linear(v0, self.rbm_module.weights,
                                  self.rbm_module.hidden_bias))

        pos_batch_size = float(len(pos_batch))
        neg_batch_size = float(len(neg_batch))

        w_grad = torch.einsum("ij,ik->jk", (prob, v0))/pos_batch_size
        vb_grad = torch.einsum("ij->j", (v0,))/pos_batch_size
        hb_grad = torch.einsum("ij->j", (prob,))/pos_batch_size

        w_grad -= torch.einsum("ij,ik->jk", (phk, vk))/neg_batch_size
        vb_grad -= torch.einsum("ij->j", (vk,))/neg_batch_size
        hb_grad -= torch.einsum("ij->j", (phk,))/neg_batch_size

        # Return negative gradients to match up nicely with the usual
        # parameter update rules, which *subtract* the gradient from
        # the parameters. This is in contrast with the RBM update
        # rules which ADD the gradients (scaled by the learning rate)
        # to the parameters.

        return {"rbm_module": {"weights": -w_grad,
                               "visible_bias": -vb_grad,
                               "hidden_bias": -hb_grad}}

    def fit(self, data, epochs=100, pos_batch_size=100, neg_batch_size=200,
            k=1, lr=1e-2, progbar=False, callbacks=[]):
        """Execute the training of the RBM.

        :param data: The actual training data
        :type data: list(float)
        :param epochs: The number of parameter (i.e. weights and biases)
                       updates
        :type epochs: int
        :param pos_batch_size: The size of batches for the positive phase
                               taken from the data.
        :type pos_batch_size: int
        :param neg_batch_size: The size of batches for the negative phase
                               taken from the data
        :type neg_batch_size: int
        :param k: The number of contrastive divergence steps
        :type k: int
        :param lr: Learning rate
        :type lr: float
        :param progbar: Whether or not to display a progress bar. If "notebook"
                        is passed, will use a Jupyter notebook compatible
                        progress bar.
        :type progbar: bool or str
        :param callbacks: Callbacks to run while training.
        :type callbacks: list(qucumber.callbacks.Callback)
        """

        disable_progbar = (progbar is False)
        progress_bar = tqdm_notebook if progbar == "notebook" else tqdm

        data = torch.tensor(data, device=self.rbm_module.device,
                            dtype=torch.double)
        optimizer = torch.optim.SGD([self.rbm_module.weights,
                                     self.rbm_module.visible_bias,
                                     self.rbm_module.hidden_bias],
                                    lr=lr)

        for cb in callbacks:
            cb.on_train_start(self)

        for ep in progress_bar(range(epochs), desc="Epochs ",
                               disable=disable_progbar):
            pos_batches = DataLoader(data, batch_size=pos_batch_size,
                                     shuffle=True)

            multiplier = int((neg_batch_size / pos_batch_size) + 0.5)
            neg_batches = [DataLoader(data, batch_size=neg_batch_size,
                                      shuffle=True)
                           for i in range(multiplier)]
            neg_batches = chain(*neg_batches)

            for cb in callbacks:
                cb.on_epoch_start(self, ep)

            if self.stop_training:  # check for stop_training signal
                break

            for batch_num, (pos_batch, neg_batch) in enumerate(zip(pos_batches,
                                                               neg_batches)):
                for cb in callbacks:
                    cb.on_batch_start(self, ep, batch_num)

                all_grads = self.compute_batch_gradients(k, pos_batch,
                                                         neg_batch)
                optimizer.zero_grad()  # clear any cached gradients

                # assign all available gradients to the corresponding parameter
                for name, grads in all_grads.items():
                    selected_RBM = getattr(self, name)
                    for param in grads.keys():
                        getattr(selected_RBM, param).grad = grads[param]

                optimizer.step()  # tell the optimizer to apply the gradients

                for cb in callbacks:
                    cb.on_batch_end(self, ep, batch_num)

            for cb in callbacks:
                cb.on_epoch_end(self, ep)

        for cb in callbacks:
            cb.on_train_end(self)

    def sample(self, num_samples, k):
        """Samples from the RBM using k steps of Block Gibbs sampling.
        :param num_samples: The number of samples to be generated
        :type num_samples: int
        :param k: Number of Block Gibbs steps.
        :type k: int
        :returns: Samples drawn from the RBM.
        :rtype: torch.doubleTensor
        """
        return self.rbm_module.sample(num_samples, k)


class ComplexRBM:
    # NOTE: In development. Might be unstable.
    # NOTE: The 'full_unitaries' argument is not needed for training/sampling.
    # This is only here for debugging the gradients. Delete this argument when
    # gradients have been debugged.
    def __init__(self, full_unitaries, psi_dictionary, num_visible,
                 num_hidden_amp, num_hidden_phase, test_grads, gpu=True,
                 seed=1234):

        self.num_visible = int(num_visible)
        self.num_hidden_amp = int(num_hidden_amp)
        self.num_hidden_phase = int(num_hidden_phase)
        self.full_unitaries = full_unitaries
        self.psi_dictionary = psi_dictionary
        self.rbm_amp = RBM_Module(num_visible, num_hidden_amp, gpu=gpu,
                                  seed=seed)
        self.rbm_phase = RBM_Module(num_visible, num_hidden_phase,
                                    zero_weights=True, gpu=gpu, seed=None)
        self.device = self.rbm_amp.device
        self.test_grads = test_grads

    def basis_state_generator(self, s):
        """Only works for binary visible units at the moment. Generates a vector
        given a spin value (0 or 1).

        :param s: A spin's value (either 0 or 1).
        :type s: float

        :returns: If s = 0, this is the (1,0) state in the basis of the
                  measurement. If s = 1, this is the (0,1) state in the basis
                  of the measurement.
        :rtype: torch.doubleTensor
        """
        if s == 0.:
            return torch.tensor([[1., 0.], [0., 0.]], dtype=torch.double)
        if s == 1.:
            return torch.tensor([[0., 1.], [0., 0.]], dtype=torch.double)

    def state_generator(self, num_non_trivial_unitaries):
        """A function that returns all possible configurations of
        'num_non_trivial_unitaries' spins. Similar to generate_visible_space.

        :param num_non_trivial_unitaries: The number of sites measured in the
                                          non-computational basis.
        :type num_non_trivial_unitaries: int

        :returns: An array of all possible spin configurations of
                  'num_non_trivial_unitaries' spins.
        :rtype: torch.doubleTensor
        """
        states = torch.zeros((2**num_non_trivial_unitaries,
                              num_non_trivial_unitaries),
                             device=self.device,
                             dtype=torch.double)
        for i in range(2**num_non_trivial_unitaries):
            temp = i
            for j in range(num_non_trivial_unitaries):
                temp, remainder = divmod(temp, 2)
                states[i][num_non_trivial_unitaries - j - 1] = remainder
        return states

    def unnormalized_probability_amp(self, v):
        r"""The effective energy of the phase RBM.

        :param v: Visible unit(s).
        :type v: torch.doubleTensor

        :returns:
            :math:`p_{\lambda}(\bm{v}) = e^{\mathcal{E}_{\lambda}(\bm{v})}`
        :rtype: torch.doubleTensor
        """
        return self.rbm_amp.unnormalized_probability(v)

    def unnormalized_probability_phase(self, v):
        r"""The effective energy of the phase RBM.

        :param v: Visible unit(s).
        :type v: torch.doubleTensor

        :returns: :math:`p_{\mu}(\bm{v}) = e^{\mathcal{E}_{\mu}(\bm{v})}`
        :rtype: torch.doubleTensor
        """
        return self.rbm_phase.unnormalized_probability(v)

    def normalized_wavefunction(self, v):
        r"""The RBM wavefunction.

        :param v: Visible unit(s).
        :type v: torch.doubleTensor

        :returns:

        .. math:: \psi_{\lambda\mu} =
                    \sqrt{\frac{p_{\lambda}}{Z_{\lambda}}}
                    \exp\left(\frac{i\log(p_{\mu})}{2}\right)

        :rtype: torch.doubleTensor
        """
        v_prime = v.view(-1, self.num_visible)
        temp1 = (self.unnormalized_probability_amp(v_prime)).sqrt()
        temp2 = ((self.unnormalized_probability_phase(v_prime)).log())*0.5

        cos_angle = temp2.cos()
        sin_angle = temp2.sin()

        psi = torch.zeros(2, v_prime.size()[0], dtype=torch.double)
        psi[0] = temp1*cos_angle
        psi[1] = temp1*sin_angle

        sqrt_Z = (self.rbm_amp.partition(
            self.rbm_amp.generate_visible_space())).sqrt()

        return psi / sqrt_Z

    def unnormalized_wavefunction(self, v):
        r"""The unnormalized RBM wavefunction.

        :param v: Visible unit(s).
        :type v: torch.doubleTensor

        :returns:

        .. math:: \tilde{\psi}_{\lambda\mu} =
                        \sqrt{p_{\lambda}}
                        \exp\left(\frac{i\log(p_{\mu})}{2}\right)

        :rtype: torch.doubleTensor
        """
        v_prime = v.view(-1, self.num_visible)
        temp1 = (self.unnormalized_probability_amp(v_prime)).sqrt()
        temp2 = ((self.unnormalized_probability_phase(v_prime)).log())*0.5
        cos_angle = temp2.cos()
        sin_angle = temp2.sin()

        psi = torch.zeros(2, v_prime.size()[0], dtype=torch.double)

        psi[0] = temp1*cos_angle
        psi[1] = temp1*sin_angle

        return psi

    def compute_batch_gradients(self, unitary_dict, k, batch, chars_batch):
        """This function will compute the gradients of a batch of the training
        data (data_file) given the basis measurements (chars_file).

        :param k: Number of contrastive divergence steps in amplitude training.
        :type k: int
        :param batch: Batch of the input data.
        :type batch: torch.doubleTensor
        :param chars_batch: Batch of bases that correspondingly indicates the
                            basis each site in the batch was measured in.
        :type chars_batch: list(str)

        :returns: Dictionary containing all the gradients (negative): Gradient
                  of weights, visible bias and hidden bias for the amplitude,
                  Gradients of weights, visible bias and hidden bias for the
                  phase.
        :rtype: dict
        """

        batch_size = len(batch)

        g_weights_amp = torch.zeros_like(self.rbm_amp.weights)
        g_vb_amp = torch.zeros_like(self.rbm_amp.visible_bias)
        g_hb_amp = torch.zeros_like(self.rbm_amp.hidden_bias)

        g_weights_phase = torch.zeros_like(self.rbm_phase.weights)
        g_vb_phase = torch.zeros_like(self.rbm_phase.visible_bias)
        g_hb_phase = torch.zeros_like(self.rbm_phase.hidden_bias)

        [batch, h0_amp_batch, vk_amp_batch, hk_amp_batch, phk_amp_batch] = \
            self.rbm_amp.gibbs_sampling(k, batch)
        # Iterate through every data point in the batch.
        for row_count, v0 in enumerate(batch):

            # A counter for the number of non-trivial unitaries
            # (non-computational basis) in the data point.
            num_non_trivial_unitaries = 0

            # tau_indices will contain the index numbers of spins not in the
            # computational basis (Z). z_indices will contain the index numbers
            # of spins in the computational basis.
            tau_indices = []
            z_indices = []

            for j in range(self.num_visible):
                # Go through the unitaries (chars_batch[row_count]) of each
                # site in the data point, v0, and save inidices of non-trivial.
                if chars_batch[row_count][j] != 'Z':
                    num_non_trivial_unitaries += 1
                    tau_indices.append(j)
                else:
                    z_indices.append(j)

            if num_non_trivial_unitaries == 0:
                # If there are no non-trivial unitaries for the data point v0,
                # calculate the positive phase of regular (i.e. non-complex
                # RBM) gradient. Use the actual data point, v0.
                prob_amp = F.sigmoid(F.linear(v0, self.rbm_amp.weights,
                                              self.rbm_amp.hidden_bias))
                g_weights_amp -= (torch.einsum("i,j->ij", (prob_amp, v0))
                                  / batch_size)
                g_vb_amp -= v0 / batch_size
                g_hb_amp -= prob_amp / batch_size

            else:
                # Compute the rotated gradients.
                [L_weights_amp, L_vb_amp, L_hb_amp,
                 L_weights_phase, L_vb_phase, L_hb_phase] = \
                    self.compute_rotated_grads(unitary_dict, k, v0,
                                               chars_batch[row_count],
                                               num_non_trivial_unitaries,
                                               z_indices, tau_indices)

                # Gradents of amplitude parameters take the real part of the
                # rotated gradients.
                g_weights_amp -= L_weights_amp[0] / batch_size
                g_vb_amp -= L_vb_amp[0] / batch_size
                g_hb_amp -= L_hb_amp[0] / batch_size

                # Gradents of phase parameters take the imaginary part of the
                # rotated gradients.
                g_weights_phase += L_weights_phase[1] / batch_size
                g_vb_phase += L_vb_phase[1] / batch_size
                g_hb_phase += L_hb_phase[1] / batch_size

        # Block gibbs sampling for negative phase.
        g_weights_amp += (torch.einsum("ij,ik->jk",
                                       (phk_amp_batch, vk_amp_batch))
                          / batch_size)
        g_vb_amp += torch.einsum("ij->j", (vk_amp_batch,)) / batch_size
        g_hb_amp += torch.einsum("ij->j", (phk_amp_batch,)) / batch_size

        """Return negative gradients to match up nicely with the usual
        parameter update rules, which *subtract* the gradient from
        the parameters. This is in contrast with the RBM update
        rules which ADD the gradients (scaled by the learning rate)
        to the parameters."""

        return {
            "rbm_amp": {
                "weights": g_weights_amp,
                "visible_bias": g_vb_amp,
                "hidden_bias": g_hb_amp
            },
            "rbm_phase": {
                "weights": g_weights_phase,
                "visible_bias": g_vb_phase,
                "hidden_bias": g_hb_phase
            }
        }

    def compute_rotated_grads(self, unitary_dict, k, v0, characters,
                              num_non_trivial_unitaries,
                              z_indices, tau_indices):
        """Computes the rotated gradients.

        :param v0: A visible unit.
        :type v0: torch.doubleTensor
        :param characters: A string of characters corresponding to the basis
                           that each site in v0 was measured in.
        :type characters: str
        :param num_non_trivial_unitaries: The number of sites in v0 that aren't
                                          measured in the computational basis.
        :type num_non_trivial_unitaries: int
        :param z_indices: A list of indices that correspond to sites of v0 that
                          are measured in the computational basis.
        :type z_indices: list(int)
        :param tau_indices: A list of indices that correspond to sites of v0
                            that are not measured in the computational basis.
        :type tau_indices: list(int)

        :returns: Dictionary of the rotated gradients: L_weights_amp, L_vb_amp,
                  L_hb_amp, L_weights_phase, L_vb_phase, L_hb_phase
        :rtype: dict
        """
        """Initialize the 'A' parameters (see alg 4.2)."""
        A_weights_amp = torch.zeros(2, self.rbm_amp.weights.size()[0],
                                    self.rbm_amp.weights.size()[1],
                                    device=self.device, dtype=torch.double)
        A_vb_amp = torch.zeros(2, self.rbm_amp.visible_bias.size()[0],
                               device=self.device, dtype=torch.double)
        A_hb_amp = torch.zeros(2, self.rbm_amp.hidden_bias.size()[0],
                               device=self.device, dtype=torch.double)

        A_weights_phase = torch.zeros(2, self.rbm_phase.weights.size()[0],
                                      self.rbm_phase.weights.size()[1],
                                      device=self.device, dtype=torch.double)
        A_vb_phase = torch.zeros(2, self.rbm_phase.visible_bias.size()[0],
                                 device=self.device, dtype=torch.double)
        A_hb_phase = torch.zeros(2, self.rbm_phase.hidden_bias.size()[0],
                                 device=self.device, dtype=torch.double)
        # 'B' will contain the coefficients of the rotated unnormalized
        # wavefunction.
        B = torch.zeros(2, device=self.device, dtype=torch.double)

        w_grad_amp = torch.zeros_like(self.rbm_amp.weights)
        vb_grad_amp = torch.zeros_like(self.rbm_amp.visible_bias)
        hb_grad_amp = torch.zeros_like(self.rbm_amp.hidden_bias)

        w_grad_phase = torch.zeros_like(self.rbm_phase.weights)
        vb_grad_phase = torch.zeros_like(self.rbm_phase.visible_bias)
        hb_grad_phase = torch.zeros_like(self.rbm_phase.hidden_bias)

        zeros_for_w_amp = torch.zeros_like(w_grad_amp)
        zeros_for_w_phase = torch.zeros_like(w_grad_phase)
        zeros_for_vb = torch.zeros_like(vb_grad_amp)
        zeros_for_hb_amp = torch.zeros_like(hb_grad_amp)
        zeros_for_hb_phase = torch.zeros_like(hb_grad_phase)

        # Loop over Hilbert space of the non trivial unitaries to build
        # the state.
        for j in range(2**num_non_trivial_unitaries):
            s = self.state_generator(num_non_trivial_unitaries)[j]
            # Creates a matrix where the jth row is the desired state, |S>,
            # a vector.

            # This is the sigma state.
            constructed_state = torch.zeros(
                self.num_visible, dtype=torch.double)

            U = torch.tensor([1., 0.], dtype=torch.double, device=self.device)

            # Populate the |sigma> state (aka constructed_state) accirdingly.
            for index in range(len(z_indices)):
                # These are the sites in the computational basis.
                constructed_state[z_indices[index]] = v0[z_indices[index]]

            for index in range(len(tau_indices)):
                # These are the sites that are NOT in the computational basis.
                constructed_state[tau_indices[index]] = s[index]

                aa = unitary_dict[characters[tau_indices[index]]]
                bb = self.basis_state_generator(v0[tau_indices[index]])
                cc = self.basis_state_generator(s[index])

                temp = cplx.inner_prod(cplx.MV_mult(
                    cplx.compT_matrix(aa), bb), cc)

                U = cplx.scalar_mult(U, temp)

            # Positive phase gradients for phase and amp. Will be added into
            # the 'A' parameters.

            prob_amp = F.sigmoid(F.linear(constructed_state,
                                          self.rbm_amp.weights,
                                          self.rbm_amp.hidden_bias))
            prob_phase = F.sigmoid(F.linear(constructed_state,
                                            self.rbm_phase.weights,
                                            self.rbm_phase.hidden_bias))

            w_grad_amp = torch.einsum("i,j->ij", (prob_amp, constructed_state))
            vb_grad_amp = constructed_state
            hb_grad_amp = prob_amp

            w_grad_phase = torch.einsum("i,j->ij",
                                        (prob_phase, constructed_state))
            vb_grad_phase = constructed_state
            hb_grad_phase = prob_phase

            """
            In order to calculate the 'A' parameters below with the current
            complex library, I need to make the weights and biases complex.
            I fill the complex parts of the parameters with a tensor of zeros.
            """
            temp_w_grad_amp = cplx.make_complex_matrix(w_grad_amp,
                                                       zeros_for_w_amp)
            temp_vb_grad_amp = cplx.make_complex_vector(vb_grad_amp,
                                                        zeros_for_vb)
            temp_hb_grad_amp = cplx.make_complex_vector(hb_grad_amp,
                                                        zeros_for_hb_amp)

            temp_w_grad_phase = cplx.make_complex_matrix(w_grad_phase,
                                                         zeros_for_w_phase)
            temp_vb_grad_phase = cplx.make_complex_vector(vb_grad_phase,
                                                          zeros_for_vb)
            temp_hb_grad_phase = cplx.make_complex_vector(hb_grad_phase,
                                                          zeros_for_hb_phase)

            # Temp = U*psi(sigma)
            temp = cplx.scalar_mult(
                U, self.unnormalized_wavefunction(constructed_state))

            A_weights_amp += cplx.MS_mult(temp, temp_w_grad_amp)
            A_vb_amp += cplx.VS_mult(temp, temp_vb_grad_amp)
            A_hb_amp += cplx.VS_mult(temp, temp_hb_grad_amp)

            A_weights_phase += cplx.MS_mult(temp, temp_w_grad_phase)
            A_vb_phase += cplx.VS_mult(temp, temp_vb_grad_phase)
            A_hb_phase += cplx.VS_mult(temp, temp_hb_grad_phase)

            # Rotated wavefunction.
            B += temp

        L_weights_amp = cplx.MS_divide(A_weights_amp, B)
        L_vb_amp = cplx.VS_divide(A_vb_amp, B)
        L_hb_amp = cplx.VS_divide(A_hb_amp, B)

        L_weights_phase = cplx.MS_divide(A_weights_phase, B)
        L_vb_phase = cplx.VS_divide(A_vb_phase, B)
        L_hb_phase = cplx.VS_divide(A_hb_phase, B)

        return [L_weights_amp, L_vb_amp, L_hb_amp, L_weights_phase, L_vb_phase,
                L_hb_phase]

    def fit(self, data, character_data, unitary_dict, epochs, batch_size,
            k=1, lr=1e-2, log_every=0, progbar=False):
        """Execute the training of the RBM.

        :param data: The actual training data
        :type data: list(float)
        :param character_data: The corresponding bases that each site in the
                               data has been measured in.
        :type character_data: list(str)
        :param epochs: The number of parameter (i.e. weights and biases)
                       updates
        :type epochs: int
        :param batch_size: The size of batches taken from the data
        :type batch_size: int
        :param k: The number of contrastive divergence steps
        :type k: int
        :param lr: Learning rate
        :type lr: float
        :param progbar: Whether or not to display a progress bar. If "notebook"
                        is passed, will use a Jupyter notebook compatible
                        progress bar.
        :type progbar: bool or str
        """
        # Make data file into a torch tensor.
        data = torch.tensor(data, dtype=torch.double).to(device=self.device)

        # Use the Adam optmizer to update the weights and biases.
        optimizer = torch.optim.Adam([self.rbm_amp.weights,
                                      self.rbm_amp.visible_bias,
                                      self.rbm_amp.hidden_bias,
                                      self.rbm_phase.weights,
                                      self.rbm_phase.visible_bias,
                                      self.rbm_phase.hidden_bias],
                                     lr=lr)

        disable_progbar = (progbar is False)
        progress_bar = tqdm_notebook if progbar == "notebook" else tqdm

        vis = self.rbm_amp.generate_visible_space()

        for ep in progress_bar(range(0, epochs + 1),
                               desc="Epochs ", total=epochs,
                               disable=disable_progbar):
            # Shuffle the data to ensure that the batches taken from the data
            # are random data points.
            random_permutation = torch.randperm(data.shape[0])

            shuffled_data = data[random_permutation]
            shuffled_character_data = character_data[random_permutation]

            # List of all the batches.
            batches = [shuffled_data[batch_start:(batch_start + batch_size)]
                       for batch_start in range(0, len(data), batch_size)]

            # List of all the bases.
            char_batches = [shuffled_character_data[batch_start:(batch_start + batch_size)]
                            for batch_start in range(0, len(data), batch_size)]

            # Calculate convergence quantities every "log-every" steps.
            if ep % log_every == 0:
                fidelity_ = self.fidelity(vis, 'Z' 'Z')
                print ('Epoch = ',ep,'\nFidelity = ',fidelity_)

            # Save parameters at the end of training.
            if ep == epochs:
                print('Finished training. Saving results...')
                self.save_params()
                print('Done.')
                break

            # Loop through all of the batches and calculate the batch
            # gradients.
            for index, batch in progress_bar(enumerate(batches),
                                             desc="Batches",
                                             leave=False, disable=True):

                all_grads = self.compute_batch_gradients(unitary_dict, k, batch,
                                                         char_batches[index])

                if self.test_grads:
                    self.test_gradients(unitary_dict, vis, k, batches[index],
                                        char_batches[index], all_grads)

                # Clear any cached gradients.
                optimizer.zero_grad()

                # Assign all available gradients to the
                # corresponding parameter.
                for name, grads in all_grads.items():
                    selected_RBM = getattr(self, name)
                    for param in grads.keys():
                        getattr(selected_RBM, param).grad = grads[param]

                # Tell the optimizer to apply the gradients and update
                # the parameters.
                optimizer.step()

    def save_params(self):
        """A function that saves the weights and biases for the amplitude and
        phase individually."""
        trained_params = [self.rbm_amp.weights.data.numpy(),
                          self.rbm_amp.visible_bias.data.numpy(),
                          self.rbm_amp.hidden_bias.data.numpy(),
                          self.rbm_phase.weights.data.numpy(),
                          self.rbm_phase.visible_bias.data.numpy(),
                          self.rbm_phase.hidden_bias.data.numpy()]

        with open('trained_weights_amp.csv', 'w') as csvfile:
            np.savetxt(csvfile, trained_params[0], fmt='%.5f', delimiter=',')

        with open('trained_visible_bias_amp.csv', 'w') as csvfile:
            np.savetxt(csvfile, trained_params[1], fmt='%.5f', delimiter=',')

        with open('trained_hidden_bias_amp.csv', 'w') as csvfile:
            np.savetxt(csvfile, trained_params[2], fmt='%.5f', delimiter=',')

        with open('trained_weights_phase.csv', 'w') as csvfile:
            np.savetxt(csvfile, trained_params[3], fmt='%.5f', delimiter=',')

        with open('trained_visible_bias_phase.csv', 'w') as csvfile:
            np.savetxt(csvfile, trained_params[4], fmt='%.5f', delimiter=',')

        with open('trained_hidden_bias_phase.csv', 'w') as csvfile:
            np.savetxt(csvfile, trained_params[5], fmt='%.5f', delimiter=',')

    def get_true_psi(self, basis):
        """Picks out the true psi in the correct basis.

        :param basis: E.g. XZZZX.
        :type basis: str

        :returns: The true wavefunction in the basis.
        :rtype: torch.doubleTensor
        """
        key = ''
        for i in range(len(basis)):
            key += basis[i]
        return self.psi_dictionary[key]

    def overlap(self, visible_space, basis):
        """Computes the overlap between the RBM and true wavefunctions.

        :param visible_space: An array of all possible spin configurations.
        :type visible_space: torch.doubleTensor
        :param basis: E.g. XZZZX.
        :type basis: str

        :returns: :math:`O = \\langle{\\psi_{true}}\\vert\\psi_{\\lambda\\mu}\\rangle`.
        :rtype: float
        """
        overlap_ = cplx.inner_prod(self.get_true_psi(basis),
                           self.normalized_wavefunction(visible_space))
        return overlap_

    def fidelity(self, visible_space, basis):
        """Computed the fidelity of the RBM and true wavefunctions.

        :param visible_space: An array of all possible spin configurations.
        :type visible_space: torch.doubleTensor
        :param basis: E.g. XZZZX.
        :type basis: str

        :returns: :math:`F = |O|^2`.
        :rtype: float
        """
        return cplx.norm(self.overlap(visible_space, basis))

    def KL_divergence(self, visible_space):
        '''Computes the total KL divergence.
        '''
        KL = 0.0
        basis_list = ['Z' 'Z', 'X' 'Z', 'Z' 'X', 'Y' 'Z', 'Z' 'Y']

        '''Wavefunctions (RBM and true) in the computational basis.'''
        # psi_ZZ      = self.normalized_wavefunction(visible_space)
        # true_psi_ZZ = self.get_true_psi('ZZ')

        '''Compute the KL divergence for the non computational bases.'''
        for i in range(len(basis_list)):
            rotated_RBM_psi = cplx.MV_mult(
                self.full_unitaries[basis_list[i]],
                self.normalized_wavefunction(visible_space))
            rotated_true_psi = self.get_true_psi(basis_list[i])

            for j in range(len(visible_space)):
                elementof_rotated_RBM_psi = torch.tensor(
                    [rotated_RBM_psi[0][j], rotated_RBM_psi[1][j]]).view(2, 1)
                elementof_rotated_true_psi = (torch.tensor(
                    [rotated_true_psi[0][j], rotated_true_psi[1][j]]
                ).view(2, 1))

                norm_true_psi = cplx.norm(cplx.inner_prod(
                    elementof_rotated_true_psi, elementof_rotated_true_psi))
                norm_RBM_psi = cplx.norm(cplx.inner_prod(
                    elementof_rotated_RBM_psi, elementof_rotated_RBM_psi))

                if norm_true_psi > 0.0:
                    KL += norm_true_psi*torch.log(norm_true_psi)
                if norm_RBM_psi > 0.0:
                    KL -= norm_true_psi*torch.log(norm_RBM_psi)

        '''Compute KL divergence for the computational basis.'''
        '''
        for j in range(len(visible_space)):

            elementof_ZZ_RBM_psi  = torch.tensor([psi_ZZ[0][j], psi_ZZ[1][j]]).view(2,1)
            elementof_ZZ_true_psi = torch.tensor([true_psi_ZZ[0][j], true_psi_ZZ[1][j]]).view(2,1)

            norm_ZZ_true_psi = cplx.norm( cplx.inner_prod(elementof_ZZ_true_psi, elementof_ZZ_true_psi) )
            norm_ZZ_RBM_psi  = cplx.norm( cplx.inner_prod(elementof_ZZ_RBM_psi, elementof_ZZ_RBM_psi) )

            if norm_ZZ_true_psi > 0.0:
                KL += norm_ZZ_true_psi*torch.log(norm_ZZ_true_psi)

            KL -= norm_ZZ_true_psi*torch.log(norm_ZZ_RBM_psi)
        '''
        return KL

    def compute_numerical_gradient(self, visible_space, param, alg_grad):
        eps = 1.e-6
        print("Numerical\t Exact\t\t Abs. Diff.")

        for i in range(len(param)):

            param[i].data += eps
            KL_pos = self.KL_divergence(visible_space)

            param[i].data -= 2*eps
            KL_neg = self.KL_divergence(visible_space)

            param[i].data += eps

            num_grad = (KL_pos - KL_neg) / (2*eps)

            print("{: 10.8f}\t{: 10.8f}\t{: 10.8f}\t"
                  .format(num_grad, alg_grad[i],
                          abs(num_grad - alg_grad[i])))

    def test_gradients(self, unitary_dict, visible_space, k, batch, chars_batch, alg_grads):
        # Must have negative sign because the compute_batch_grads returns the neg of the grads.
        # key_list = ["weights_amp", "visible_bias_amp", "hidden_bias_amp", "weights_phase", "visible_bias_phase", "hidden_bias_phase"]

        flat_weights_amp = self.rbm_amp.weights.data.view(-1)
        flat_weights_phase = self.rbm_phase.weights.data.view(-1)

        flat_grad_weights_amp = alg_grads["rbm_amp"]["weights"].view(-1)
        flat_grad_weights_phase = alg_grads["rbm_phase"]["weights"].view(-1)

        print('-------------------------------------------------------------------------------')

        print('Weights amp gradient')
        self.compute_numerical_gradient(
            visible_space, flat_weights_amp, -flat_grad_weights_amp)
        print ('\n')

        print('Visible bias amp gradient')
        self.compute_numerical_gradient(
            visible_space, self.rbm_amp.visible_bias, -alg_grads["rbm_amp"]["visible_bias"])
        print ('\n')

        print('Hidden bias amp gradient')
        self.compute_numerical_gradient(
            visible_space, self.rbm_amp.hidden_bias, -alg_grads["rbm_amp"]["hidden_bias"])
        print ('\n')

        print('Weights phase gradient')
        self.compute_numerical_gradient(
            visible_space, flat_weights_phase, -flat_grad_weights_phase)
        print ('\n')

        print('Visible bias phase gradient')
        self.compute_numerical_gradient(
            visible_space, self.rbm_phase.visible_bias, -alg_grads["rbm_phase"]["visible_bias"])
        print ('\n')

        print('Hidden bias phase gradient')
        self.compute_numerical_gradient(
            visible_space, self.rbm_phase.hidden_bias, -alg_grads["rbm_phase"]["hidden_bias"])

    def state_to_index(self, state):
        ''' Only for debugging how the unitary is applied to the unnormalized wavefunction - the 'B' term in alg 4.2.'''
        states = torch.zeros(2**self.num_visible, self.num_visible)
        npstates = states.numpy()
        npstate = state.numpy()
        for i in range(2**self.num_visible):
            temp = i

            for j in range(self.num_visible):
                temp, remainder = divmod(temp, 2)
                npstates[i][self.num_visible - j - 1] = remainder

            if np.array_equal(npstates[i], npstate):
                return i

The effective energies of the given visible states.          .. math::              \mathcal{E}(\bm{v}) &= \sum_{j}b_j v_j                         + \sum_{i}\log                             \left\lbrack 1 +                                   \exp\left(c_{i} + \sum_{j} W_{ij} v_j\right)                             \right\rbrack          :param v: The visible states.         :type v: torch.doubleTensor          :returns: The effective energies of the given visible states.         :rtype: torch.doubleTensor Given a hidden unit configuration, compute the probability         vector of the visible units being on.          :param h: The hidden unit         :type h: torch.doubleTensor          :returns: The probability of visible units being active given the                   hidden state.         :rtype: torch.doubleTensor Given a visible unit configuration, compute the probability         vector of the hidden units being on.          :param h: The hidden unit.         :type h: torch.doubleTensor          :returns: The probability of hidden units being active given the                   visible state.         :rtype: torch.doubleTensor Sample/generate a visible state given a hidden state.          :param h: The hidden state.         :type h: torch.doubleTensor          :returns: Tuple containing prob_v_given_h(h) and the sampled visible                   state.         :rtype: tuple(torch.doubleTensor, torch.doubleTensor) Sample/generate a hidden state given a visible state.          :param h: The visible state.         :type h: torch.doubleTensor          :returns: Tuple containing prob_h_given_v(v) and the sampled hidden                   state.         :rtype: tuple(torch.doubleTensor, torch.doubleTensor) Performs k steps of Block Gibbs sampling given an initial visible         state v0.          :param k: Number of Block Gibbs steps.         :type k: int         :param v0: The initial visible state.         :type v0: torch.doubleTensor          :returns: Tuple containing the initial visible state, v0,                   the hidden state sampled from v0,                   the visible state sampled after k steps,                   the hidden state sampled after k steps and its corresponding                    probability vector.         :rtype: tuple(torch.doubleTensor, torch.doubleTensor,                       torch.doubleTensor, torch.doubleTensor,                       torch.doubleTensor) Samples from the RBM using k steps of Block Gibbs sampling.         :param num_samples: The number of samples to be generated         :type num_samples: int         :param k: Number of Block Gibbs steps.         :type k: int         :returns: Samples drawn from the RBM         :rtype: torch.doubleTensor The unnormalized probabilities of the given visible states.          .. math:: p(\bm{v}) = e^{\mathcal{E}(\bm{v})}          :param v: The visible states.         :type v: torch.doubleTensor          :returns: The unnormalized probability of the given visible state(s).         :rtype: torch.doubleTensor Generates all possible visible states.          :returns: A tensor of all possible spin configurations.         :rtype: torch.doubleTensor The natural logarithm of the partition function of the RBM.          :param visible_space: A rank 2 tensor of the entire visible space.         :type visible_space: torch.doubleTensor          :returns: The natural log of the partition function.         :rtype: torch.doubleTensor The partition function of the RBM.          :param visible_space: A rank 2 tensor of the entire visible space.         :type visible_space: torch.doubleTensor          :returns: The partition function.         :rtype: torch.doubleTensor Evaluates the probability of the given vector(s) of visible         units; NOT RECOMMENDED FOR RBMS WITH A LARGE # OF VISIBLE UNITS          :param v: The visible states.         :type v: torch.doubleTensor         :param Z: The partition function.         :type Z: float          :returns: The probability of the given vector(s) of visible units.         :rtype: torch.doubleTensor Saves the RBM parameters to the given location along with         any given metadata.          :param location: The location to save the RBM parameters + metadata         :type location: str or file         :param metadata: Any extra metadata to store alongside the RBM                          parameters         :type metadata: dict # add extra metadata to dictionary before saving it to disk Loads the RBM parameters from the given location ignoring any         metadata stored in the file. Overwrites the RBM's parameters.          .. note::             The RBM object on which this function is called must             have the same shape as the one who's parameters are being             loaded.          :param location: The location to load the RBM parameters from         :type location: str or file This function will compute the gradients of a batch of the training         data (data_file) given the basis measurements (chars_file).          :param k: Number of contrastive divergence steps in training.         :type k: int         :param pos_batch: Batch of the input data for the positive phase.         :type pos_batch: |DoubleTensor|         :param neg_batch: Batch of the input data for the negative phase.         :type neg_batch: |DoubleTensor|          :returns: Dictionary containing all the gradients of the parameters.         :rtype: dict # Return negative gradients to match up nicely with the usual # parameter update rules, which *subtract* the gradient from # the parameters. This is in contrast with the RBM update # rules which ADD the gradients (scaled by the learning rate) # to the parameters. Execute the training of the RBM.          :param data: The actual training data         :type data: list(float)         :param epochs: The number of parameter (i.e. weights and biases)                        updates         :type epochs: int         :param pos_batch_size: The size of batches for the positive phase                                taken from the data.         :type pos_batch_size: int         :param neg_batch_size: The size of batches for the negative phase                                taken from the data         :type neg_batch_size: int         :param k: The number of contrastive divergence steps         :type k: int         :param lr: Learning rate         :type lr: float         :param progbar: Whether or not to display a progress bar. If "notebook"                         is passed, will use a Jupyter notebook compatible                         progress bar.         :type progbar: bool or str         :param callbacks: Callbacks to run while training.         :type callbacks: list(qucumber.callbacks.Callback) # check for stop_training signal # clear any cached gradients # assign all available gradients to the corresponding parameter # tell the optimizer to apply the gradients Samples from the RBM using k steps of Block Gibbs sampling.         :param num_samples: The number of samples to be generated         :type num_samples: int         :param k: Number of Block Gibbs steps.         :type k: int         :returns: Samples drawn from the RBM.         :rtype: torch.doubleTensor # NOTE: In development. Might be unstable. # NOTE: The 'full_unitaries' argument is not needed for training/sampling. # This is only here for debugging the gradients. Delete this argument when # gradients have been debugged. Only works for binary visible units at the moment. Generates a vector         given a spin value (0 or 1).          :param s: A spin's value (either 0 or 1).         :type s: float          :returns: If s = 0, this is the (1,0) state in the basis of the                   measurement. If s = 1, this is the (0,1) state in the basis                   of the measurement.         :rtype: torch.doubleTensor A function that returns all possible configurations of         'num_non_trivial_unitaries' spins. Similar to generate_visible_space.          :param num_non_trivial_unitaries: The number of sites measured in the                                           non-computational basis.         :type num_non_trivial_unitaries: int          :returns: An array of all possible spin configurations of                   'num_non_trivial_unitaries' spins.         :rtype: torch.doubleTensor The effective energy of the phase RBM.          :param v: Visible unit(s).         :type v: torch.doubleTensor          :returns:             :math:`p_{\lambda}(\bm{v}) = e^{\mathcal{E}_{\lambda}(\bm{v})}`         :rtype: torch.doubleTensor The effective energy of the phase RBM.          :param v: Visible unit(s).         :type v: torch.doubleTensor          :returns: :math:`p_{\mu}(\bm{v}) = e^{\mathcal{E}_{\mu}(\bm{v})}`         :rtype: torch.doubleTensor The RBM wavefunction.          :param v: Visible unit(s).         :type v: torch.doubleTensor          :returns:          .. math:: \psi_{\lambda\mu} =                     \sqrt{\frac{p_{\lambda}}{Z_{\lambda}}}                     \exp\left(\frac{i\log(p_{\mu})}{2}\right)          :rtype: torch.doubleTensor The unnormalized RBM wavefunction.          :param v: Visible unit(s).         :type v: torch.doubleTensor          :returns:          .. math:: \tilde{\psi}_{\lambda\mu} =                         \sqrt{p_{\lambda}}                         \exp\left(\frac{i\log(p_{\mu})}{2}\right)          :rtype: torch.doubleTensor This function will compute the gradients of a batch of the training         data (data_file) given the basis measurements (chars_file).          :param k: Number of contrastive divergence steps in amplitude training.         :type k: int         :param batch: Batch of the input data.         :type batch: torch.doubleTensor         :param chars_batch: Batch of bases that correspondingly indicates the                             basis each site in the batch was measured in.         :type chars_batch: list(str)          :returns: Dictionary containing all the gradients (negative): Gradient                   of weights, visible bias and hidden bias for the amplitude,                   Gradients of weights, visible bias and hidden bias for the                   phase.         :rtype: dict # Iterate through every data point in the batch. # A counter for the number of non-trivial unitaries # (non-computational basis) in the data point. # tau_indices will contain the index numbers of spins not in the # computational basis (Z). z_indices will contain the index numbers # of spins in the computational basis. # Go through the unitaries (chars_batch[row_count]) of each # site in the data point, v0, and save inidices of non-trivial. # If there are no non-trivial unitaries for the data point v0, # calculate the positive phase of regular (i.e. non-complex # RBM) gradient. Use the actual data point, v0. # Compute the rotated gradients. # Gradents of amplitude parameters take the real part of the # rotated gradients. # Gradents of phase parameters take the imaginary part of the # rotated gradients. # Block gibbs sampling for negative phase. Return negative gradients to match up nicely with the usual         parameter update rules, which *subtract* the gradient from         the parameters. This is in contrast with the RBM update         rules which ADD the gradients (scaled by the learning rate)         to the parameters. Computes the rotated gradients.          :param v0: A visible unit.         :type v0: torch.doubleTensor         :param characters: A string of characters corresponding to the basis                            that each site in v0 was measured in.         :type characters: str         :param num_non_trivial_unitaries: The number of sites in v0 that aren't                                           measured in the computational basis.         :type num_non_trivial_unitaries: int         :param z_indices: A list of indices that correspond to sites of v0 that                           are measured in the computational basis.         :type z_indices: list(int)         :param tau_indices: A list of indices that correspond to sites of v0                             that are not measured in the computational basis.         :type tau_indices: list(int)          :returns: Dictionary of the rotated gradients: L_weights_amp, L_vb_amp,                   L_hb_amp, L_weights_phase, L_vb_phase, L_hb_phase         :rtype: dict Initialize the 'A' parameters (see alg 4.2). # 'B' will contain the coefficients of the rotated unnormalized # wavefunction. # Loop over Hilbert space of the non trivial unitaries to build # the state. # Creates a matrix where the jth row is the desired state, |S>, # a vector. # This is the sigma state. # Populate the |sigma> state (aka constructed_state) accirdingly. # These are the sites in the computational basis. # These are the sites that are NOT in the computational basis. # Positive phase gradients for phase and amp. Will be added into # the 'A' parameters. In order to calculate the 'A' parameters below with the current             complex library, I need to make the weights and biases complex.             I fill the complex parts of the parameters with a tensor of zeros. # Temp = U*psi(sigma) # Rotated wavefunction. Execute the training of the RBM.          :param data: The actual training data         :type data: list(float)         :param character_data: The corresponding bases that each site in the                                data has been measured in.         :type character_data: list(str)         :param epochs: The number of parameter (i.e. weights and biases)                        updates         :type epochs: int         :param batch_size: The size of batches taken from the data         :type batch_size: int         :param k: The number of contrastive divergence steps         :type k: int         :param lr: Learning rate         :type lr: float         :param progbar: Whether or not to display a progress bar. If "notebook"                         is passed, will use a Jupyter notebook compatible                         progress bar.         :type progbar: bool or str # Make data file into a torch tensor. # Use the Adam optmizer to update the weights and biases. # Shuffle the data to ensure that the batches taken from the data # are random data points. # List of all the batches. # List of all the bases. # Calculate convergence quantities every "log-every" steps. # Save parameters at the end of training. # Loop through all of the batches and calculate the batch # gradients. # Clear any cached gradients. # Assign all available gradients to the # corresponding parameter. # Tell the optimizer to apply the gradients and update # the parameters. A function that saves the weights and biases for the amplitude and         phase individually. Picks out the true psi in the correct basis.          :param basis: E.g. XZZZX.         :type basis: str          :returns: The true wavefunction in the basis.         :rtype: torch.doubleTensor Computes the overlap between the RBM and true wavefunctions.          :param visible_space: An array of all possible spin configurations.         :type visible_space: torch.doubleTensor         :param basis: E.g. XZZZX.         :type basis: str          :returns: :math:`O = \\langle{\\psi_{true}}\\vert\\psi_{\\lambda\\mu}\\rangle`.         :rtype: float Computed the fidelity of the RBM and true wavefunctions.          :param visible_space: An array of all possible spin configurations.         :type visible_space: torch.doubleTensor         :param basis: E.g. XZZZX.         :type basis: str          :returns: :math:`F = |O|^2`.         :rtype: float Computes the total KL divergence. Wavefunctions (RBM and true) in the computational basis. # psi_ZZ      = self.normalized_wavefunction(visible_space) # true_psi_ZZ = self.get_true_psi('ZZ') Compute the KL divergence for the non computational bases. Compute KL divergence for the computational basis. for j in range(len(visible_space)):              elementof_ZZ_RBM_psi  = torch.tensor([psi_ZZ[0][j], psi_ZZ[1][j]]).view(2,1)             elementof_ZZ_true_psi = torch.tensor([true_psi_ZZ[0][j], true_psi_ZZ[1][j]]).view(2,1)              norm_ZZ_true_psi = cplx.norm( cplx.inner_prod(elementof_ZZ_true_psi, elementof_ZZ_true_psi) )             norm_ZZ_RBM_psi  = cplx.norm( cplx.inner_prod(elementof_ZZ_RBM_psi, elementof_ZZ_RBM_psi) )              if norm_ZZ_true_psi > 0.0:                 KL += norm_ZZ_true_psi*torch.log(norm_ZZ_true_psi)              KL -= norm_ZZ_true_psi*torch.log(norm_ZZ_RBM_psi) # Must have negative sign because the compute_batch_grads returns the neg of the grads. # key_list = ["weights_amp", "visible_bias_amp", "hidden_bias_amp", "weights_phase", "visible_bias_phase", "hidden_bias_phase"] Only for debugging how the unitary is applied to the unnormalized wavefunction - the 'B' term in alg 4.2.