import numpy as np
from qsim.codes import qubit, rydberg
from qsim.codes.quantum_state import State
from qsim import tools
from scipy.linalg import expm
import scipy.sparse as sparse
from scipy.sparse.linalg import expm_multiply
from qsim.graph_algorithms.graph import Graph, enumerate_independent_sets


class HamiltonianDriver(object):
    def __init__(self, transition: tuple = (0, 1), energies: tuple = (1,), pauli='X', code=qubit, graph=None,
                 IS_subspace=False):
        """Default is that the first element in transition is the higher energy s."""
        self.transition = transition
        self.energies = energies
        assert pauli in ['X', 'Y', 'Z']
        self.pauli = pauli
        self.code = code
        self.graph = graph
        if self.pauli == 'X' and not self.code.logical_code:
            self._operator = np.zeros((self.code.d, self.code.d))
            self._operator[self.transition[1], self.transition[0]] = 1
            self._operator[self.transition[0], self.transition[1]] = 1
        elif self.pauli == 'Y' and not self.code.logical_code:
            self._operator = np.zeros((self.code.d, self.code.d))
            self._operator[self.transition[1], self.transition[0]] = 1j
            self._operator[self.transition[0], self.transition[1]] = -1j
        elif self.pauli == 'Z' and not self.code.logical_code:
            self._operator = np.zeros((self.code.d, self.code.d))
            self._operator[self.transition[0], self.transition[0]] = 1
            self._operator[self.transition[1], self.transition[1]] = -1
        # If a logical code, we should use the normal qubit operators because we assume the code is a qubit
        elif self.pauli == 'X' and self.code.logical_code:
            self._operator = self.code.X
        elif self.pauli == 'Y' and self.code.logical_code:
            self._operator = self.code.Y
        elif self.pauli == 'Z' and self.code.logical_code:
            self._operator = self.code.Z
        self.IS_subspace = IS_subspace

        if self.IS_subspace:
            # Generate sparse mixing Hamiltonian
            assert isinstance(graph, Graph)
            if code is not qubit:
                IS, num_IS = graph.generate_independent_sets_qudit(self.code)
                if self.pauli == 'Z':
                    self._diagonal_hamiltonian = np.zeros((num_IS, 1))
                    for k in range(num_IS):
                        self._diagonal_hamiltonian[k, 0] = np.sum(IS[k, ...] == self.transition[0]) - np.sum(
                            IS[k, ...] == self.transition[1])

                    self._csr_hamiltonian = sparse.csr_matrix((self._diagonal_hamiltonian.T[0], (np.arange(num_IS),
                                                                                                 np.arange(num_IS))))

                    self._hamiltonian = self._csr_hamiltonian

                else:
                    # For each IS, look at spin flips generated by the laser
                    # Over-allocate space
                    rows = np.zeros(graph.n * num_IS, dtype=int)
                    columns = np.zeros(graph.n * num_IS, dtype=int)
                    entries = np.zeros(graph.n * num_IS, dtype=int)
                    num_terms = 0
                    for i in range(num_IS):
                        for j in range(graph.n):
                            if IS[i, j] == self.transition[0]:
                                # Flip spin at this location
                                # Get binary representation
                                temp = IS[i, ...].copy()
                                temp[j] = self.transition[1]
                                where_matched = (np.argwhere(np.sum(np.abs(IS - temp), axis=1) == 0).flatten())
                                if len(where_matched) > 0:
                                    # This is a valid spin flip by removing a node
                                    rows[num_terms] = where_matched[0]
                                    columns[num_terms] = i
                                    if self.pauli == 'X':
                                        entries[num_terms] = 1
                                    elif self.pauli == 'Y':
                                        # entries[num_terms] = -1j
                                        entries[num_terms] = 1j
                                    num_terms += 1
                    """else:
                        for (i, key) in enumerate(self.graph.independent_sets_dict):
                            num_neighbors = len(self.graph.independent_sets_dict[key][1])
                            rows[num_terms:num_terms + num_neighbors] = self.graph.independent_sets_dict[key][1]
                            columns[num_terms:num_terms + num_neighbors] = np.ones(num_neighbors) * \
                                                                           self.graph.independent_sets_dict[key][0]
                            if self.pauli == 'X':
                                entries[num_terms:num_terms + num_neighbors] = 1
                            elif self.pauli == 'Y':
                                entries[num_terms:num_terms + num_neighbors] = -1j
                            num_terms += num_neighbors"""

                    # Cut off the excess in the arrays
                    columns = columns[:2 * num_terms]
                    rows = rows[:2 * num_terms]
                    entries = entries[:2 * num_terms]
                    # Populate the second half of the entries according to self.pauli
                    if self.pauli == 'X':
                        columns[num_terms:2 * num_terms] = rows[:num_terms]
                        rows[num_terms:2 * num_terms] = columns[:num_terms]
                        entries[num_terms:2 * num_terms] = entries[:num_terms]
                    elif self.pauli == 'Y':
                        columns[num_terms:2 * num_terms] = rows[:num_terms]
                        rows[num_terms:2 * num_terms] = columns[:num_terms]
                        entries[num_terms:2 * num_terms] = -1 * entries[:num_terms]
                    # Now, construct the Hamiltonian
                    self._csr_hamiltonian = sparse.csr_matrix((entries, (rows, columns)), shape=(num_IS, num_IS))
                    self._hamiltonian = self._csr_hamiltonian
            else:
                # Use graph generator functions
                if self.pauli == 'Z':
                    independent_sets = enumerate_independent_sets(graph.graph)
                    # Generate a list of integers corresponding to the independent sets in binary
                    # All ones
                    k = self.graph.num_independent_sets - 2
                    self.mis_size = 0
                    hamiltonian = np.zeros(self.graph.num_independent_sets, dtype=float)
                    hamiltonian[-1] = -1 * self.graph.n
                    for i in independent_sets:
                        hamiltonian[k] = len(i) - (self.graph.n - len(i))
                        k -= 1

                    self._hamiltonian = sparse.csr_matrix(
                        (hamiltonian,
                         (np.arange(self.graph.num_independent_sets), np.arange(self.graph.num_independent_sets))),
                        shape=(self.graph.num_independent_sets, self.graph.num_independent_sets))

                else:
                    independent_sets = enumerate_independent_sets(graph.graph)
                    # Generate a list of integers corresponding to the independent sets in binary
                    previous_size = 0
                    self.mis_size = 0
                    independent_sets_dict = {(): self.graph.num_independent_sets - 1}
                    rows = []
                    columns = []
                    entries = []
                    k = self.graph.num_independent_sets - 2
                    for i in independent_sets:
                        current_size = len(i)
                        if current_size - previous_size > 1:
                            previous_size = current_size - 1
                            # Clear out the dictionary with terms we can't connect to
                            for key in list(independent_sets_dict):
                                if len(key) != previous_size:
                                    independent_sets_dict.pop(key)
                        independent_sets_dict[tuple(i)] = k
                        for (j, node) in enumerate(i):
                            i_removed = i.copy()
                            i_removed.pop(j)
                            index = independent_sets_dict[tuple(i_removed)]
                            # Index is the current independent set with a single node removed
                            rows.append(k)
                            columns.append(index)
                            rows.append(index)
                            columns.append(k)
                            if self.pauli == 'Y':
                                entries.append(-1j)
                                entries.append(1j)
                            else:
                                entries.append(1)
                                entries.append(1)
                        k -= 1
                    # Now, construct the Hamiltonian
                    self._csr_hamiltonian = sparse.csr_matrix((entries, (rows, columns)),
                                                              shape=(self.graph.num_independent_sets,
                                                                     self.graph.num_independent_sets))
                    self._hamiltonian = self._csr_hamiltonian
        else:
            self._hamiltonian = None
        self._left_acting_hamiltonian = None
        self._right_acting_hamiltonian = None

    @property
    def hamiltonian(self):
        if self._hamiltonian is None:
            assert not self.IS_subspace
            try:
                assert self.graph is not None
            except AssertionError:
                print('self.graph must be not None to generate the Hamiltonian property.')
            self._hamiltonian = sparse.csr_matrix(((self.code.d * self.code.n) ** self.graph.n,
                                                   (self.code.d * self.code.n) ** self.graph.n))
            for i in range(self.graph.n):
                self._hamiltonian = self._hamiltonian + tools.tensor_product(
                    [sparse.identity((self.code.d * self.code.n) ** i),
                     self._operator,
                     sparse.identity((self.code.d * self.code.n) ** (self.graph.n - i - 1))],
                    sparse=True)
        return self.energies[0] * self._hamiltonian

    @property
    def evolution_operator(self, vector_space='hilbert'):
        if vector_space != 'hilbert' and vector_space != 'liouville':
            raise Exception('Attribute vector_space must be hilbert or liouville')
        if vector_space == 'liouville':
            if self._left_acting_hamiltonian is None:
                self._left_acting_hamiltonian = sparse.kron(sparse.identity(self._hamiltonian.shape[0]),
                                                            self._hamiltonian)
                self._right_acting_hamiltonian = sparse.kron(self._hamiltonian.T,
                                                             sparse.identity(self._hamiltonian.shape[0]))

            return -1j * self.energies[0] * self._left_acting_hamiltonian + 1j * self.energies[0] * \
                   self._right_acting_hamiltonian
        else:
            return -1j * self.hamiltonian

    def left_multiply(self, state: State):
        if not self.IS_subspace:
            temp = np.zeros_like(state)
            # For each logical qubit
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    if self.pauli == 'X':
                        temp = temp + self.code.left_multiply(state, [i], ['X'])
                    elif self.pauli == 'Y':
                        temp = temp + self.code.left_multiply(state, [i], ['Y'])
                    elif self.pauli == 'Z':
                        temp = temp + self.code.left_multiply(state, [i], ['Z'])
                elif not self.code.logical_code:
                    ind = self.code.d ** i
                    out = np.zeros_like(state, dtype=np.complex128)
                    if state.is_ket:
                        state = state.reshape((-1, self.code.d, ind), order='F')
                        # Note index start from the right (sN,...,s3,s2,s1)
                        out = out.reshape((-1, self.code.d, ind), order='F')
                        if self.pauli == 'X':  # Sigma_X
                            # We want to exchange two indices
                            out[:, [self.transition[0], self.transition[1]], :] = \
                                state[:, [self.transition[1], self.transition[0]], :]
                        elif self.pauli == 'Y':  # Sigma_Y
                            out[:, [self.transition[0], self.transition[1]], :] = \
                                state[:, [self.transition[1], self.transition[0]], :]
                            out[:, self.transition[0], :] = -1j * out[:, self.transition[0], :]
                            out[:, self.transition[1], :] = 1j * out[:, self.transition[1], :]
                        elif self.pauli == 'Z':  # Sigma_Z
                            out[:, [self.transition[0], self.transition[1]], :] = \
                                state[:, [self.transition[0], self.transition[1]], :]
                            out[:, self.transition[1], :] = -1 * out[:, self.transition[1], :]
                        state = state.reshape(state_shape, order='F')
                        out = out.reshape(state_shape, order='F')
                    else:
                        out = out.reshape((-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1),
                                           self.code.d, ind), order='F')
                        state = state.reshape((-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1),
                                               self.code.d, ind), order='F')

                        if self.pauli == 'X':  # Sigma_X
                            out[:, [self.transition[0], self.transition[1]], :, :, :] = \
                                state[:, [self.transition[1], self.transition[0]], :, :, :]
                        elif self.pauli == 'Y':  # Sigma_Y
                            out[:, [self.transition[0], self.transition[1]], :, :, :] = \
                                state[:, [self.transition[1], self.transition[0]], :, :, :]
                            out[:, self.transition[0], :, :, :] = -1j * out[:, self.transition[0], :, :, :]
                            out[:, self.transition[1], :, :, :] = 1j * out[:, self.transition[1], :, :, :]
                        elif self.pauli == 'Z':  # Sigma_Z
                            out[:, [self.transition[0], self.transition[1]], :, :, :] = \
                                state[:, [self.transition[0], self.transition[1]], :, :, :]
                            out[:, self.transition[1], :, :, :] = -1 * out[:, self.transition[1], :, :, :]
                        state = state.reshape(state_shape, order='F')
                        out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                         graph=self.graph)
        else:
            # Handle dimensions
            if self.pauli == 'Z' and not self.code.logical_code:  # In this case, the Hamiltonian is diagonal
                return State(self.energies[0] * self._diagonal_hamiltonian * state, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                return State(self.energies[0] * self._csr_hamiltonian @ state, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            print('Warning: right multiply functionality currently applies the operator and daggers the state.')
            return self.left_multiply(state).conj().T
        if not self.IS_subspace:
            temp = np.zeros_like(state)
            # For each physical qubit
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    if self.pauli == 'X':
                        temp = temp + self.code.right_multiply(state, [i], ['X'])
                    elif self.pauli == 'Y':
                        temp = temp + self.code.right_multiply(state, [i], ['Y'])
                    elif self.pauli == 'Z':
                        temp = temp + self.code.right_multiply(state, [i], ['Z'])
                elif not self.code.logical_code:
                    ind = self.code.d ** i
                    out = np.zeros_like(state)
                    out = out.reshape(
                        (-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1), self.code.d, ind),
                        order='F')
                    state = state.reshape(
                        (-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1), self.code.d, ind),
                        order='F')
                    if self.pauli == 'X' and not self.code.logical_code:  # Sigma_X
                        out[:, :, :, [self.transition[0], self.transition[1]], :] = state[:, :, :, [self.transition[1],
                                                                                                    self.transition[0]],
                                                                                    :]
                    elif self.pauli == 'Y' and not self.code.logical_code:  # Sigma_Y
                        out[:, :, :, [self.transition[0], self.transition[1]], :] = state[:, :, :, [self.transition[1],
                                                                                                    self.transition[0]],
                                                                                    :]
                        out[:, :, :, self.transition[0], :] = -1j * out[:, :, :, self.transition[0], :]
                        out[:, :, :, self.transition[1], :] = 1j * out[:, :, :, self.transition[1], :]
                    elif self.pauli == 'Z' and not self.code.logical_code:  # Sigma_Z
                        out[:, :, :, [self.transition[0], self.transition[1]], :] = state[:, :, :, [self.transition[0],
                                                                                                    self.transition[1]],
                                                                                    :]
                        out[:, :, :, self.transition[1], :] = -1 * state[:, :, :, self.transition[1], :]
                        state = state.reshape(state_shape, order='F')
                    state = state.reshape(state_shape, order='F')
                    out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                         graph=self.graph)
        else:
            if self.pauli == 'Z' and not self.code.logical_code:  # In this case, the Hamiltonian is diagonal
                return State(state * self.hamiltonian.T, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                             code=state.code, graph=self.graph)
            else:
                return State(state @ self.hamiltonian.T.conj(), is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                             code=state.code, graph=self.graph)

    def evolve(self, state: State, time):
        r"""
        Use reshape to efficiently implement evolution under :math:`H_B=\\sum_i X_i`
        """
        if not self.IS_subspace:
            # We don't want to modify the original s
            out = state.copy()
            for i in range(state.number_logical_qudits):
                # Note that self._operator is not necessarily involutary
                if self.pauli == 'X':
                    out = self.code.rotation(out, [i], self.energies[0] * time, self._operator)
                elif self.pauli == 'Y':
                    out = self.code.rotation(out, [i], self.energies[0] * time, self._operator)
                elif self.pauli == 'Z':
                    out = self.code.rotation(out, [i], self.energies[0] * time, self._operator)
            return out
        else:
            if state.is_ket:
                # Handle dimensions
                if self.hamiltonian.shape[1] == 1:
                    return State(np.exp(-1j * time * self.hamiltonian) * state, is_ket=state.is_ket,
                                 IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
                else:
                    return State(expm_multiply(-1j * time * self.hamiltonian, state),
                                 is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                if self.hamiltonian.shape[1] == 1:
                    exp_hamiltonian = np.exp(-1j * time * self.hamiltonian)
                    return State(exp_hamiltonian * state * exp_hamiltonian.conj().T,
                                 is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
                else:
                    exp_hamiltonian = expm(-1j * time * self.hamiltonian)
                    return State(exp_hamiltonian @ state @ exp_hamiltonian.conj().T,
                                 is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)


class HamiltonianMaxCut(object):
    def __init__(self, graph: Graph, code=qubit, energies=(1,), cost_function=True, use_Z2_symmetry=False):
        # If MIS is true, create an MIS Hamiltonian. Otherwise, make a MaxCut Hamiltonian
        r"""
        Generate a vector corresponding to the diagonal of the MaxCut Hamiltonian.
        """
        self.code = code
        self.energies = energies
        # Make sure all edges have weight attribute; default to 1

        self.graph = graph
        self.optimization = 'max'
        self.n = self.graph.n
        if use_Z2_symmetry:
            c = np.zeros([self.code.d ** (self.code.n * (self.n - 1)), 1])
        else:
            c = np.zeros([self.code.d ** (self.code.n * self.n), 1])
        if tools.is_diagonal(self.code.Z):
            self._is_diagonal = True
            z = np.expand_dims(np.diagonal(self.code.Z), axis=0).T

            def my_eye(n):
                return np.ones((np.asarray(self.code.d ** self.code.n) ** n, 1))
        else:
            self._is_diagonal = False
            # Compute the optimum first. We don't care that this takes extra time, since it only needs to run once
            z = np.expand_dims(np.diagonal(qubit.Z), axis=0).T

            def my_eye(n):
                return np.ones((np.asarray(qubit.d) ** n, 1))

            for a, b in self.graph.edges:
                if b < a:
                    a, b = b, a
                if use_Z2_symmetry:
                    if cost_function:
                        if a == min(self.graph.nodes):
                            c = c - 1 / 2 * graph.graph[a][b]['weight'] * (tools.tensor_product(
                                [my_eye(b - 1), z, my_eye(self.n - b - 1)]) - my_eye(self.n - 1))
                        else:
                            c = c - 1 / 2 * graph.graph[a][b]['weight'] * (tools.tensor_product(
                                [my_eye(a - 1), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)]) - my_eye(self.n - 1))
                    else:
                        if a == min(self.graph.nodes):
                            c = c + graph.graph[a][b]['weight'] * (tools.tensor_product(
                                [my_eye(b - 1), z, my_eye(self.n - b - 1)]))
                        else:
                            c = c + graph.graph[a][b]['weight'] * (tools.tensor_product(
                                [my_eye(a - 1), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)]))
                else:
                    if cost_function:
                        c = c - 1 / 2 * graph.graph[a][b]['weight'] * (tools.tensor_product(
                            [my_eye(a), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)]) - my_eye(self.n))
                    else:
                        c = c + graph.graph[a][b]['weight'] * tools.tensor_product(
                            [my_eye(a), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)])
            self._optimum = np.max(c).real
            c = sparse.csr_matrix((self.code.d ** (self.code.n * self.n), self.code.d ** (self.code.n * self.n)))

            z = sparse.csr_matrix(self.code.Z)

            def my_eye(n):
                return sparse.csr_matrix(np.ones(np.asarray(z.shape[0]) ** n),
                                         (np.asarray(z.shape[0]) ** n, np.asarray(z.shape[0]) ** n))

        for a, b in self.graph.edges:
            if b < a:
                a, b = b, a

            if cost_function:
                if use_Z2_symmetry:
                    if a == min(self.graph.nodes):
                        c = c - 1 / 2 * graph.graph[a][b]['weight'] * (tools.tensor_product(
                            [my_eye(b - 1), z, my_eye(self.n - b - 1)]) - my_eye(self.n - 1))
                    else:
                        c = c - 1 / 2 * graph.graph[a][b]['weight'] * (tools.tensor_product(
                            [my_eye(a - 1), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)]) - my_eye(self.n - 1))
                else:
                    c = c - 1 / 2 * graph.graph[a][b]['weight'] * (tools.tensor_product(
                        [my_eye(a), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)],
                        sparse=(not self._is_diagonal)) - my_eye(
                        self.n))
            else:
                if use_Z2_symmetry:
                    if a == min(self.graph.nodes):
                        c = c + graph.graph[a][b]['weight'] * (tools.tensor_product(
                            [my_eye(b - 1), z, my_eye(self.n - b - 1)]))
                    else:
                        c = c + graph.graph[a][b]['weight'] * (tools.tensor_product(
                            [my_eye(a - 1), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)]))

                else:
                    c = c + graph.graph[a][b]['weight'] * (tools.tensor_product(
                        [my_eye(a), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)],
                        sparse=(not self._is_diagonal)))
        if self._is_diagonal:
            self._diagonal_hamiltonian = c
            self._optimum = np.max(c).real
            if use_Z2_symmetry:
                c = sparse.csr_matrix((c.flatten(), (np.arange(self.code.d ** (self.code.n * (self.n - 1))),
                                                     np.arange(self.code.d ** (self.code.n * (self.n - 1))))),
                                      shape=(self.code.d ** (self.code.n * (self.n - 1)),
                                             self.code.d ** (self.code.n * (self.n - 1))))
            else:
                c = sparse.csr_matrix((c.flatten(), (np.arange(self.code.d ** (self.code.n * self.n)),
                                                     np.arange(self.code.d ** (self.code.n * self.n)))),
                                      shape=(self.code.d ** (self.code.n * self.n),
                                             self.code.d ** (self.code.n * self.n)))
        else:
            # c is already the right shape, just convert it to a csr matrix
            c = sparse.csr_matrix(c)
        self._hamiltonian = c
        self._left_acting_hamiltonian = None
        self._right_acting_hamiltonian = None

    @property
    def hamiltonian(self):
        return self.energies[0] * self._hamiltonian

    @property
    def evolution_operator(self, vector_space='hilbert'):
        if vector_space != 'hilbert' and vector_space != 'liouville':
            raise Exception('Attribute vector_space must be hilbert or liouville')
        if vector_space == 'liouville':
            if self._left_acting_hamiltonian is None:
                self._left_acting_hamiltonian = sparse.kron(sparse.identity(self._hamiltonian.shape[0]),
                                                            self._hamiltonian)
                self._right_acting_hamiltonian = sparse.kron(self._hamiltonian.T,
                                                             sparse.identity(self._hamiltonian.shape[0]))

            return -1j * self.energies[0] * self._left_acting_hamiltonian + 1j * self.energies[0] * \
                   self._right_acting_hamiltonian
        else:
            return -1j * self.hamiltonian

    @property
    def optimum(self):
        # Optimum for non-diagonal Hamiltonians can be found by computing the optimum in the standard basis,
        # which is done in self.__init__()
        return self.energies[0] * self._optimum

    def evolve(self, state: State, time):
        if state.is_ket:
            if self._is_diagonal:
                # It's quicker to exponentiate a diagonal array than use expm_multiply
                return State(np.exp(-1j * time * self._diagonal_hamiltonian) * state, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                return State(expm_multiply(-1j * time * self.hamiltonian, state), is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
        else:
            if self._is_diagonal:
                return State(np.exp(-1j * time * self._diagonal_hamiltonian) * state * np.exp(
                    1j * time * self._diagonal_hamiltonian).T, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                             code=state.code, graph=self.graph)
            else:
                temp = expm(-1j * time * self.hamiltonian)
                return State(temp @ state @ temp.conj().T, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                             code=state.code, graph=self.graph)

    def left_multiply(self, state: State):
        if self._is_diagonal:
            return State(self._diagonal_hamiltonian * state, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                         code=state.code, graph=self.graph)
        else:
            return State(self.hamiltonian @ state, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                         code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        # Already real, so you don't need to conjugate
        if state.is_ket:
            if self._is_diagonal:
                return State(state.conj().T * self._diagonal_hamiltonian.T, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                return State((state @ self.hamiltonian.T).conj().T, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
        else:
            if self._is_diagonal:
                # Density matrices are already Hermitian, so you don't need to dagger
                return State(state * self._diagonal_hamiltonian.T, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                return State(state @ self.hamiltonian.T, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

    def cost_function(self, state: State):
        # Need to project into the IS subspace
        # Returns <s|C|s>
        if state.is_ket:
            if self._is_diagonal:
                return np.real(np.vdot(state, self._diagonal_hamiltonian * state))
            else:
                return np.real(np.vdot(state, self.hamiltonian @ state))
        else:
            # Density matrix
            if self._is_diagonal:
                return np.real(np.squeeze(tools.trace(self._diagonal_hamiltonian * state)))
            else:
                return np.real(np.squeeze(tools.trace(self.hamiltonian @ state)))

    def optimum_overlap(self, state: State):
        # Returns \sum_i <s|opt_i><opt_i|s>
        if self._is_diagonal:
            optimum_indices = np.argwhere(self._diagonal_hamiltonian == self.optimum).T[0]
            # Construct an operator that is zero everywhere except at the optimum
            optimum = np.zeros(self._diagonal_hamiltonian.shape)
            optimum[optimum_indices] = 1
        else:
            # The plan for this is to basically use the code for _is_diagonal to identify the logical qubit subspaces
            # which encode the optimum. Then, make an operator that's the identity in those subspaces
            raise NotImplementedError('Optimum overlap not implemented for non-diagonal Hamiltonians')
        if state.is_ket:
            return np.real(np.vdot(state, optimum * state))
        else:
            # Density matrix
            return np.real(np.squeeze(tools.trace(optimum * state)))

    def approximation_ratio(self, state: State):
        # Returns <s|C|s>/optimum
        return self.cost_function(state) / self.optimum


class HamiltonianMIS(object):
    def __init__(self, graph: Graph, energies=(1, 1), code=qubit, IS_subspace=False):
        r"""
        Generate a vector corresponding to the diagonal of the MIS Hamiltonian.
        """
        if energies == (1, 1) and IS_subspace:
            energies = (1,)
        self.code = code
        self.graph = graph
        self.n = self.graph.n
        self.energies = energies
        self.IS_subspace = IS_subspace
        self.optimization = 'max'
        if not self.IS_subspace:
            # Store node and edge terms separately so the Hamiltonian can be dynamically updated when energies
            # are changed

            if tools.is_diagonal(self.code.Q):
                self._hamiltonian_edge_terms = np.zeros([1, (self.code.d ** self.code.n) ** self.n])
                self._hamiltonian_node_terms = np.zeros([1, (self.code.d ** self.code.n) ** self.n])
                self._is_diagonal = True

                Q = np.expand_dims(np.diagonal(self.code.Q), axis=0)

                def my_eye(n):
                    return np.ones(np.asarray(self.code.d ** self.code.n) ** n)
            else:
                # TODO: generate a sparse matrix instead
                self._hamiltonian_edge_terms = np.zeros([(self.code.d ** self.code.n) ** self.n,
                                                         (self.code.d ** self.code.n) ** self.n])
                self._hamiltonian_node_terms = np.zeros([(self.code.d ** self.code.n) ** self.n,
                                                         (self.code.d ** self.code.n) ** self.n])

                Q = np.expand_dims(np.diagonal(qubit.Q), axis=0)

                def my_eye(n):
                    return np.ones(np.asarray(qubit.d ** qubit.n) ** n)

                self._optimum_edge_terms = np.zeros([(qubit.d ** qubit.n) ** self.n,
                                                     (qubit.d ** qubit.n) ** self.n])
                self._optimum_node_terms = np.zeros([(qubit.d ** qubit.n) ** self.n,
                                                     (qubit.d ** qubit.n) ** self.n])

                for i, j in graph.graph.edges:
                    if j < i:
                        i, j = j, i
                    self._optimum_edge_terms = self._optimum_edge_terms + graph.graph.edges[(i, j)]['weight'] * \
                                               tools.tensor_product(
                                                   [my_eye(i), Q, my_eye(j - i - 1), Q, my_eye(self.n - j - 1)])

                for i in graph.graph.nodes:
                    self._optimum_node_terms = self._optimum_node_terms + graph.graph.nodes[i]['weight'] * \
                                               tools.tensor_product([my_eye(i), Q, my_eye(self.n - i - 1)])

                self._is_diagonal = False
                Q = self.code.Q

                def my_eye(n):
                    return np.identity(np.asarray(self.code.d ** self.code.n) ** n)
            for i, j in graph.graph.edges:
                if j < i:
                    i, j = j, i
                self._hamiltonian_edge_terms = self._hamiltonian_edge_terms + graph.graph.edges[(i, j)]['weight'] * \
                                               tools.tensor_product(
                                                   [my_eye(i), Q, my_eye(j - i - 1), Q, my_eye(self.n - j - 1)])
            for i in graph.graph.nodes:
                self._hamiltonian_node_terms = self._hamiltonian_node_terms + graph.graph.nodes[i]['weight'] * \
                                               tools.tensor_product([my_eye(i), Q, my_eye(self.n - i - 1)])
            self._hamiltonian_node_terms = self._hamiltonian_node_terms.T
            self._hamiltonian_edge_terms = self._hamiltonian_edge_terms.T
            if self._is_diagonal:
                self._optimum_edge_terms = self._hamiltonian_edge_terms
                self._optimum_node_terms = self._hamiltonian_node_terms
                self._diagonal_hamiltonian_edge_terms = self._hamiltonian_edge_terms.copy()
                self._diagonal_hamiltonian_node_terms = self._hamiltonian_node_terms.copy()
                self._hamiltonian_node_terms = sparse.csr_matrix(
                    (self._hamiltonian_node_terms.flatten(), (np.arange(self.code.d ** (self.code.n * self.n)),
                                                              np.arange(self.code.d ** (self.code.n * self.n)))),
                    shape=(self.code.d ** (self.code.n * self.n),
                           self.code.d ** (self.code.n * self.n)))
                self._hamiltonian_edge_terms = sparse.csr_matrix(
                    (self._hamiltonian_edge_terms.flatten(), (np.arange(self.code.d ** (self.code.n * self.n)),
                                                              np.arange(self.code.d ** (self.code.n * self.n)))),
                    shape=(self.code.d ** (self.code.n * self.n),
                           self.code.d ** (self.code.n * self.n)))
                # TODO: what happens to _optimum_node_terms if not _is_diagonal
            else:
                self._diagonal_hamiltonian_edge_terms = self._hamiltonian_edge_terms.copy()
                self._diagonal_hamiltonian_node_terms = self._hamiltonian_node_terms.copy()
                self._hamiltonian_edge_terms = sparse.csr_matrix(self._hamiltonian_edge_terms)
                self._hamiltonian_node_terms = sparse.csr_matrix(self._hamiltonian_node_terms)
            self._left_acting_hamiltonian_edge_terms = None
            self._right_acting_hamiltonian_edge_terms = None

        else:
            self._is_diagonal = True
            if not (self.code == qubit or self.code == rydberg):
                raise NotImplementedError("IS subspace only implemented for qubit and Rydberg codes.")
            # Don't generate anything that depends on the entire Hilbert space as to save space

            # These are your independent sets of the original graphs, ordered by node and size
            if self.code == qubit:
                node_weights = np.asarray([self.graph.graph.nodes[i]['weight'] for i in range(self.graph.n)])
                independent_sets = enumerate_independent_sets(self.graph.graph)
                # Generate a list of integers corresponding to the independent sets in binary
                # All ones
                k = self.graph.num_independent_sets - 2
                self.mis_size = 0
                C = np.zeros(self.graph.num_independent_sets, dtype=float)
                C[-1] = 0
                for i in independent_sets:
                    C[k] = np.sum([node_weights[j] for j in i])
                    k -= 1
                self._hamiltonian = sparse.csr_matrix((C, (np.arange(self.graph.num_independent_sets),
                                                           np.arange(self.graph.num_independent_sets))))

                C = np.expand_dims(C, axis=0).T

            # Otherwise, we need to include the possibility that we are in one of many ground space states
            elif self.code == rydberg:
                # TODO: fix this to reflect the new way of notating independent sets!
                # Count the number of elements in the ground space and map to their representation in ternary
                # Determine how large to make the array
                independent_sets, num_IS = self.graph.generate_independent_sets_qudit(self.code)
                # Generate Hamiltonian from independent sets
                node_weights = np.asarray([self.graph.graph.nodes[i]['weight'] for i in range(self.graph.n)])
                C = np.zeros((num_IS, 1), dtype=np.complex128)
                for k in independent_sets:
                    C[k, 0] = np.sum((independent_sets[k][2] == 0) * node_weights)
            self._diagonal_hamiltonian_node_terms = C
            C = C.flatten()

            self._hamiltonian_node_terms = sparse.csr_matrix((
                C, (np.arange(len(C)), np.arange(len(C)))), shape=(len(C), len(C)))

        self._left_acting_hamiltonian_node_terms = None
        self._right_acting_hamiltonian_node_terms = None

    @property
    def hamiltonian(self):
        if not self.IS_subspace:
            return self.energies[0] * self._hamiltonian_node_terms - self.energies[1] * self._hamiltonian_edge_terms
        else:
            return self.energies[0] * self._hamiltonian_node_terms

    @property
    def evolution_operator(self, vector_space='hilbert'):
        if vector_space != 'hilbert' and vector_space != 'liouville':
            raise Exception('Attribute vector_space must be hilbert or liouville')
        if vector_space == 'liouville':
            if self._left_acting_hamiltonian_node_terms is None and not self.IS_subspace:
                self._left_acting_hamiltonian_edge_terms = sparse.kron(sparse.identity(
                    self._hamiltonian_node_terms.shape[0]), self._hamiltonian_edge_terms)
                self._right_acting_hamiltonian_edge_terms = sparse.kron(
                    self._hamiltonian_edge_terms.T, sparse.identity(self._hamiltonian_edge_terms.shape[0]))
                self._left_acting_hamiltonian_node_terms = sparse.kron(sparse.identity(
                    self._hamiltonian_node_terms.shape[0]), self._hamiltonian_node_terms)
                self._right_acting_hamiltonian_node_terms = sparse.kron(
                    self._hamiltonian_node_terms.T, sparse.identity(self._hamiltonian_node_terms.shape[0]))
            elif self._left_acting_hamiltonian_node_terms is None and self.IS_subspace:
                self._left_acting_hamiltonian_node_terms = sparse.kron(sparse.identity(
                    self._hamiltonian_node_terms.shape[0]), self._hamiltonian_node_terms)
                self._right_acting_hamiltonian_node_terms = sparse.kron(
                    self._hamiltonian_node_terms.T, sparse.identity(self._hamiltonian_node_terms.shape[0]))
            if not self.IS_subspace:
                return -1j * (self.energies[0] * self._left_acting_hamiltonian_node_terms - self.energies[1] *
                              self._left_acting_hamiltonian_edge_terms) + 1j * \
                       (self.energies[0] * self._right_acting_hamiltonian_node_terms - self.energies[1] *
                        self._right_acting_hamiltonian_edge_terms)
            elif self.IS_subspace:
                return -1j * (self.energies[0] * self._left_acting_hamiltonian_node_terms) + 1j * \
                       (self.energies[0] * self._right_acting_hamiltonian_node_terms)
        else:
            return -1j * self.hamiltonian

    @property
    def _diagonal_hamiltonian(self):
        if not self.IS_subspace:
            return self.energies[0] * self._diagonal_hamiltonian_node_terms - self.energies[
                1] * self._diagonal_hamiltonian_edge_terms
        else:
            return self.energies[0] * self._diagonal_hamiltonian_node_terms

    @property
    def optimum(self):
        # This needs to be recomputed because the optimum depends on the energies
        # TODO: figure out what to compute if not _is_diagonal
        if self._is_diagonal:
            return np.max(self._diagonal_hamiltonian).real
        else:
            raise NotImplementedError('Optimum unknown for non-diagonal Hamiltonians')

    def evolve(self, state: State, time):
        if state.is_ket:
            if self._is_diagonal:
                return State(np.exp(-1j * time * self._diagonal_hamiltonian) * state, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                return State(expm_multiply(-1j * time * self.hamiltonian, state), is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
        else:
            if self._is_diagonal:
                return State(np.exp(-1j * time * self._diagonal_hamiltonian) * state * np.exp(
                    1j * time * self._diagonal_hamiltonian).T, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                             code=state.code, graph=self.graph)
            else:
                temp = expm(-1j * time * self.hamiltonian)
                return State(temp @ state @ temp.conj().T, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                             code=state.code, graph=self.graph)

    def left_multiply(self, state: State):
        if self._is_diagonal:
            return State(self._diagonal_hamiltonian * state, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                         code=state.code, graph=self.graph)
        else:
            return State(self.hamiltonian @ state, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                         code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        # Already real, so you don't need to conjugate
        if state.is_ket:
            if self._is_diagonal:
                return State(state.conj().T * self._diagonal_hamiltonian.T, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                return State((state @ self.hamiltonian.T).conj().T, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
        else:
            if self._is_diagonal:
                # Density matrices are already Hermitian, so you don't need to dagger
                return State(state * self._diagonal_hamiltonian.T, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                return State(state @ self.hamiltonian.T, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

    def cost_function(self, state: State):
        # Returns <s|C|s>
        if state.is_ket:
            if self._is_diagonal:
                return np.real(np.vdot(state, self._diagonal_hamiltonian * state))
            else:
                return np.real(np.vdot(state, self.hamiltonian @ state))
        else:
            # Density matrix
            if self._is_diagonal:
                return np.real(np.squeeze(tools.trace(self._diagonal_hamiltonian * state)))
            else:
                return np.real(np.squeeze(tools.trace(self.hamiltonian @ state)))

    def optimum_overlap(self, state: State):
        # Returns \sum_i <s|opt_i><opt_i|s>
        if self._is_diagonal:
            optimum_indices = np.argwhere(self._diagonal_hamiltonian == self.optimum).T[0]
            # Construct an operator that is zero everywhere except at the optimum
            optimum = np.zeros(self._diagonal_hamiltonian.shape)
            optimum[optimum_indices] = 1
        else:
            raise NotImplementedError('Optimum overlap not implemented for non-diagonal Hamiltonians')
        if state.is_ket:
            if self._is_diagonal:
                return np.real(np.vdot(state, optimum * state))
            else:
                return np.real(np.vdot(state, self.hamiltonian @ state))
        else:
            # Density matrix
            if self._is_diagonal:
                return np.real(np.squeeze(tools.trace(optimum * state)))
            else:
                return np.real(np.squeeze(tools.trace(self.hamiltonian @ state)))

    def approximation_ratio(self, state: State):
        # Returns <s|C|s>/optimum
        return self.cost_function(state) / self.optimum


class HamiltonianGlobalPauli(object):
    def __init__(self, pauli: str = 'X', code=qubit):
        self.code = code
        self.pauli = pauli
        if self.pauli == 'X':
            self._operator = self.code.X
        elif self.pauli == 'Y':
            self._operator = self.code.Y
        elif self.pauli == 'Z':
            self._operator = self.code.Z
        self.hamiltonian = None

    def evolve(self, state: State, alpha):
        if self.hamiltonian is None:
            """Initialize the Hamiltonian only once, as it is costly."""
            # TODO: make this a sparse matrix!
            self.hamiltonian = tools.tensor_product([self._operator] * state.number_logical_qudits)
        return self.code.multiply(np.cos(alpha) * np.identity(state.dimension) - 1j * np.sin(alpha) * self.hamiltonian)

    def left_multiply(self, state: State):
        all_qubits = list(range(state.number_logical_qudits))
        return self.code.left_multiply(state, all_qubits, [self.pauli] * state.number_logical_qudits)

    def right_multiply(self, state: State):
        all_qubits = list(range(state.number_logical_qudits))
        return self.code.right_multiply(state, all_qubits, [self.pauli] * state.number_logical_qudits)


class HamiltonianBookatzPenalty(object):
    def __init__(self, code=qubit, energies=(1,)):
        self.code = code
        self.projector = np.identity(self.code.d ** self.code.n) - self.code.code_space_projector
        self.energies = energies

    def evolve(self, state: State, time):
        # Term for a single qubit
        for i in range(state.number_logical_qudits):
            state = self.code.rotation(state, [i], self.energies[0] * time, self.projector, is_idempotent=True)
        return state

    def left_multiply(self, state: State):
        out = np.zeros_like(state, dtype=np.complex128)
        for i in range(state.number_logical_qudits):
            out = out + self.code.left_multiply(state, [i], self.projector)
        return State(self.energies[0] * out, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                     graph=self.graph)

    def right_multiply(self, state: State):
        out = np.zeros_like(state, dtype=np.complex128)
        for i in range(state.number_logical_qudits):
            out = out + self.code.right_multiply(state, [i], self.projector)
        return State(self.energies[0] * out, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                     graph=self.graph)


class HamiltonianMarvianPenalty(object):
    def __init__(self, Nx, Ny):
        super().__init__()
        self.Nx = Nx
        self.Ny = Ny
        self.n = 3 * Nx * Ny
        # Generate Hamiltonian
        # Two by two geometry (can be generalized in the future)
        hp = np.zeros([2 ** self.n, 2 ** self.n])
        for i in range(int(self.Nx * self.Ny)):
            # Add gauge interactions within a single logical qubit
            hp = hp + tools.tensor_product(
                [tools.identity(i * 3), tools.Z(), tools.Z(),
                 tools.identity(self.n - i * 3 - 2)]) + tools.tensor_product(
                [tools.identity(i * 3 + 1), tools.X(), tools.X(), tools.identity(self.n - i * 3 - 3)])
        # Between rows
        for j in range(self.Ny):
            # j is the number of rows
            for k in range(self.Nx):
                # k is the number of columns
                # Need to deal with edge effects
                # Add gauge interactions within a single logical qubit
                if k != self.Nx - 1:
                    # Along the same row
                    hp = hp + tools.tensor_product(
                        [tools.identity(j * self.Nx * 3 + k * 3), tools.X(), tools.identity(2), tools.X(),
                         tools.identity(self.n - (j * self.Nx * 3 + k * 3) - 4)]) + \
                         tools.tensor_product(
                             [tools.identity(j * self.Nx * 3 + k * 3 + 2), tools.Z(), tools.identity(2), tools.Z(),
                              tools.identity(self.n - (j * self.Nx * 3 + k * 3 + 2) - 4)])
                    # Along the same column
                if j != self.Ny - 1:
                    hp = hp + tools.tensor_product(
                        [tools.identity(j * self.Nx * 3 + k * 3), tools.X(), tools.identity(3 * self.Nx - 1),
                         tools.X(),
                         tools.identity(self.n - (j * self.Nx * 3 + k * 3) - 3 * self.Nx - 1)]) + \
                         tools.tensor_product(
                             [tools.identity(j * self.Nx * 3 + k * 3 + 2), tools.Z(),
                              tools.identity(3 * self.Nx - 1),
                              tools.Z(), tools.identity(self.n - (j * self.Nx * 3 + k * 3 + 2) - 3 * self.Nx - 1)])
        self.hamiltonian = -1 * hp


class HamiltonianHeisenberg(object):
    def __init__(self, graph: Graph, energies=(1, 1), subspace='all', code=qubit, IS_subspace=False):
        self.code = code
        self.graph = graph
        self.n = self.graph.n
        self.energies = energies
        self._is_diagonal = True
        assert code is qubit
        if IS_subspace:
            raise NotImplementedError
        self.IS_subspace = IS_subspace
        self._hamiltonian = None
        self.subspace = subspace
        if self.subspace == 'all':
            """Initialize the Hamiltonian."""
            self._hamiltonian_zz = None
            self._hamiltonian_xy = None
        else:
            # For each IS, look at spin flips generated by the laser
            # Over-allocate space
            from scipy.special import comb
            dim = comb(graph.n, (graph.n + self.subspace) / 2)
            if dim % 1 != 0:
                raise Exception('Invalid subspace for number of spins')
            dim = int(dim)
            hamiltonian_zz = np.zeros([dim, 1])
            rows = np.zeros(graph.n * dim, dtype=int)
            columns = np.zeros(graph.n * dim, dtype=int)
            entries = np.zeros(graph.n * dim, dtype=int)
            num_terms = 0
            states = np.zeros((dim, graph.n))
            for i in range(self.code.d ** (self.code.n * self.n)):
                nary = tools.int_to_nary(i, size=graph.n)
                if graph.n - np.sum(nary) == (self.subspace + graph.n) / 2:
                    states[num_terms, ...] = nary
                    num_terms += 1
            num_terms = 0
            for i in range(dim):
                zz = 0
                for a, b in self.graph.edges:
                    if b < a:
                        a, b = b, a
                    if np.abs(states[i, a] - states[i, b]) == 1:
                        zz -= 1
                    else:
                        zz += 1
                    if states[i, a] == 0 and states[i, b] == 1:
                        # Flip spin at this location
                        # Get binary representation
                        temp = states[i, ...].copy()
                        temp[a] = 1
                        temp[b] = 0
                        where_matched_temp = (np.argwhere(np.sum(np.abs(states - temp), axis=1) == 0).flatten())
                        entries[num_terms] = 1
                        rows[num_terms] = where_matched_temp
                        columns[num_terms] = i
                        entries[num_terms + 1] = 1
                        rows[num_terms + 1] = i
                        columns[num_terms + 1] = where_matched_temp
                        num_terms += 2
                hamiltonian_zz[i, 0] = zz
            # Now, construct the Hamiltonian
            self._hamiltonian_zz = sparse.csr_matrix(
                (hamiltonian_zz.flatten(), (np.arange(dim),
                                            np.arange(dim))), shape=(dim, dim))
            self._hamiltonian_xy = sparse.csr_matrix((entries, (rows, columns)), shape=(dim, dim))
            self.states = states

    @property
    def hamiltonian(self):
        if self._hamiltonian_zz is None or self._hamiltonian_xy is None and self.subspace == 'all':
            z = np.expand_dims(np.diagonal(self.code.Z), axis=0).T

            def my_eye(n):
                return np.ones((np.asarray(self.code.d ** self.code.n) ** n, 1))

            hamiltonian_zz = np.zeros([self.code.d ** (self.code.n * self.n), 1])
            for a, b in self.graph.edges:
                if b < a:
                    a, b = b, a
                hamiltonian_zz = hamiltonian_zz + self.graph.graph[a][b]['weight'] * (tools.tensor_product(
                    [my_eye(a), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)],
                    sparse=(not self._is_diagonal)))
            self._hamiltonian_zz = sparse.csr_matrix(
                (hamiltonian_zz.flatten(), (np.arange(self.code.d ** (self.code.n * self.n)),
                                            np.arange(self.code.d ** (self.code.n * self.n)))),
                shape=(self.code.d ** (self.code.n * self.n),
                       self.code.d ** (self.code.n * self.n)))

            # For each IS, look at spin flips generated by the laser
            # Over-allocate space
            rows = np.zeros(self.graph.n * self.code.d ** (self.code.n * self.n), dtype=int)
            columns = np.zeros(self.graph.n * self.code.d ** (self.code.n * self.n), dtype=int)
            entries = np.zeros(self.graph.n * self.code.d ** (self.code.n * self.n), dtype=int)
            num_terms = 0
            for i in range(self.code.d ** (self.code.n * self.n)):
                nary = tools.int_to_nary(i, size=self.graph.n)
                for a, b in self.graph.edges:
                    if b < a:
                        a, b = b, a
                    if nary[a] == 0 and nary[b] == 1:
                        # Flip spin at this location
                        # Get binary representation
                        temp = nary.copy()
                        temp[a] = 1
                        temp[b] = 0
                        temp = tools.nary_to_int(temp)
                        entries[num_terms] = 1
                        rows[num_terms] = temp
                        columns[num_terms] = i
                        entries[num_terms + 1] = 1
                        rows[num_terms + 1] = i
                        columns[num_terms + 1] = temp
                        num_terms += 2
            # Now, construct the Hamiltonian
            self._hamiltonian_xy = sparse.csr_matrix((entries, (rows, columns)),
                                                     shape=(self.code.d ** (self.code.n * self.n),
                                                            self.code.d ** (self.code.n * self.n)))
        return self.energies[0] * self._hamiltonian_zz + self.energies[1] * self._hamiltonian_xy

    def left_multiply(self, state: State):
        temp = np.zeros(state.shape, dtype=np.complex128)
        for edge in self.graph.edges:
            if self.energies[0] != 0:
                term = self.code.left_multiply(state, [edge[0], edge[1]], ['X', 'X'])
                temp = temp + self.energies[0] * term
                term = self.code.left_multiply(state, [edge[0], edge[1]], ['Y', 'Y'])
                temp = temp + self.energies[0] * term
            if self.energies[1] != 0:
                term = self.code.left_multiply(state, [edge[0], edge[1]], ['Z', 'Z'])
                temp = temp + self.energies[1] * term
        return State(temp, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        temp = np.zeros(state.shape, dtype=np.complex128)
        for edge in self.graph.edges:
            if self.energies[0] != 0:
                term = self.code.right_multiply(state, [edge[0], edge[1]], ['X', 'X'])
                temp = temp + self.energies[0] * term
                term = self.code.right_multiply(state, [edge[0], edge[1]], ['Y', 'Y'])
                temp = temp + self.energies[0] * term
            if self.energies[1] != 0:
                term = self.code.right_multiply(state, [edge[0], edge[1]], ['Z', 'Z'])
                temp = temp + self.energies[1] * term
        return State(temp, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

    def evolve(self, state: State, time):
        if not state.is_ket:
            exp_hamiltonian = expm(-1j * time * self.hamiltonian)
            return State(exp_hamiltonian @ state @ exp_hamiltonian.conj().T,
                         is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
        if state.is_ket:
            return State(expm_multiply(-1j * time * self.hamiltonian, state), is_ket=state.is_ket,
                         IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

    def cost_function(self, state: State):
        # Need to project into the IS subspace
        # Returns <s|C|s>
        if state.is_ket:
            return np.real(np.vdot(state, self.hamiltonian * state))
        else:
            # Density matrix
            return np.real(np.squeeze(tools.trace(self.hamiltonian * state)))


class HamiltonianEnergyShift(object):
    def __init__(self, index: int = 0, energies=(1,), code=qubit, IS_subspace=False, graph=None):
        """Default is that the first element in transition is the higher energy s."""
        self.index = index
        self.graph = graph
        self.energies = energies
        self.code = code
        if not self.code.logical_code:
            if not 0 <= self.index < self.code.d:
                raise Exception('Index exceeds qudit dimension.')
            self._operator = np.zeros((self.code.d, self.code.d))
            self._operator[self.index, self.index] = 1
        else:
            if self.index != 0 and self.index != 1:
                raise Exception('Logical codes are qubits, so index must be 0 or 1.')
            if self.index == 0:
                self._operator = self.code.Q
            elif self.index == 1:
                self._operator = self.code.P
        self.IS_subspace = IS_subspace
        if self.IS_subspace:
            # Generate sparse mixing Hamiltonian
            assert graph is not None
            assert isinstance(graph, Graph)
            if code is not qubit:
                IS, num_IS = graph.independent_sets_qudit(self.code)
                self._diagonal_hamiltonian = np.zeros((num_IS, 1), dtype=float)
                for k in range(num_IS):
                    self._diagonal_hamiltonian[k, 0] = np.sum(IS[k, ...] == self.index)
                self._hamiltonian = sparse.csr_matrix(
                    (self._diagonal_hamiltonian.T[0], (np.arange(num_IS), np.arange(num_IS))),
                    shape=(num_IS, num_IS))
            else:
                # We have already solved for this information
                independent_sets = enumerate_independent_sets(graph.graph)
                # Generate a list of integers corresponding to the independent sets in binary
                # All ones
                k = self.graph.num_independent_sets - 2
                self.mis_size = 0
                hamiltonian = np.zeros(self.graph.num_independent_sets, dtype=float)
                hamiltonian[-1] = 0
                for i in independent_sets:
                    hamiltonian[k] = len(i)
                    k -= 1

                self._hamiltonian = sparse.csr_matrix(
                    (hamiltonian,
                     (np.arange(self.graph.num_independent_sets), np.arange(self.graph.num_independent_sets))),
                    shape=(self.graph.num_independent_sets, self.graph.num_independent_sets))
        else:
            # Use full Hilbert space
            self._hamiltonian = None

    @property
    def hamiltonian(self):
        if self._hamiltonian is None:
            assert not self.IS_subspace
            try:
                assert self.graph is not None
            except AssertionError:
                print('self.graph must be not None to generate the Hamiltonian property.')
            self._hamiltonian = sparse.csr_matrix(((self.code.d * self.code.n) ** self.graph.n,
                                                   (self.code.d * self.code.n) ** self.graph.n))
            for i in range(self.graph.n):
                self._hamiltonian = self._hamiltonian + tools.tensor_product(
                    [sparse.identity((self.code.d * self.code.n) ** i),
                     self._operator,
                     sparse.identity((self.code.d * self.code.n) ** (self.graph.n - i - 1))],
                    sparse=True)
        return self.energies[0] * self._hamiltonian

    def left_multiply(self, state: State):
        if not self.IS_subspace:
            temp = np.zeros_like(state)
            # For each logical qubit
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.left_multiply(state, [i], self._operator)
                elif not self.code.logical_code:
                    ind = self.code.d ** i
                    out = np.zeros_like(state, dtype=np.complex128)
                    if state.is_ket:
                        state = state.reshape((-1, self.code.d, ind), order='F')
                        # Note index start from the right (sN,...,s3,s2,s1)
                        out = out.reshape((-1, self.code.d, ind), order='F')
                        out[:, self.index, :] = state[:, self.index, :]
                        state = state.reshape(state_shape, order='F')
                        out = out.reshape(state_shape, order='F')
                    else:
                        out = out.reshape((-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1),
                                           self.code.d, ind), order='F')
                        state = state.reshape((-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1),
                                               self.code.d, ind), order='F')
                        out[:, self.index, :, :, :] = state[:, self.index, :, :, :]
                        state = state.reshape(state_shape, order='F')
                        out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                         graph=self.graph)
        else:
            # Handle dimensions
            return State(self.energies[0] * self._diagonal_hamiltonian * state, is_ket=state.is_ket,
                         IS_subspace=state.IS_subspace,
                         code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            print('Warning: right multiply functionality currently applies the operator and daggers the state.')
            return self.left_multiply(state).conj().T
        elif not self.IS_subspace:
            temp = np.zeros_like(state)
            # For each physical qubit
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.right_multiply(state, [i], self._operator)
                else:
                    ind = self.code.d ** i
                    out = np.zeros_like(state)
                    out = out.reshape(
                        (-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1), self.code.d, ind),
                        order='F')
                    state = state.reshape(
                        (-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1), self.code.d, ind),
                        order='F')
                    out[:, :, :, self.index, :] = state[:, :, :, self.index, :]
                    state = state.reshape(state_shape, order='F')
                    out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                         graph=self.graph)
        else:
            return State(self.energies[0] * state * self._diagonal_hamiltonian.T, is_ket=state.is_ket,
                         IS_subspace=state.IS_subspace,
                         code=state.code, graph=self.graph)

    def evolve(self, state: State, time):
        r"""
        Use reshape to efficiently implement evolution under :math:`H_B=\\sum_i X_i`
        """
        if not self.IS_subspace:
            # We don't want to modify the original s
            out = state.copy()
            for i in range(state.number_logical_qudits):
                # Note that self._operator is not necessarily involutary
                out = self.code.rotation(out, [i], self.energies[0] * time, self._operator)
            return out
        else:
            if state.is_ket:
                # Handle dimensions
                return State(np.exp(-1j * time * self.energies[0] * self._diagonal_hamiltonian) * state,
                             is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                exp_hamiltonian = np.exp(-1j * time * self.energies[0] * self._diagonal_hamiltonian)
                return State(exp_hamiltonian * state * exp_hamiltonian.conj().T,
                             is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)


class HamiltonianRydberg(object):
    def __init__(self, tails_graph: Graph, hard_constraint_graph=None, index: int = 0,
                 energies=(1,), code=qubit, IS_subspace=False):
        """Default is that the first element in transition is the higher energy s."""
        self.index = index
        self.energies = energies
        self.code = code
        self.tails_graph = tails_graph
        self.hard_constraint_graph = hard_constraint_graph
        if self.code.logical_code:
            raise Exception('Logical codes not supported.')
        self.IS_subspace = IS_subspace
        try:
            if self.tails_graph.periodic or self.hard_constraint_graph.periodic:
                raise Exception('Periodic graphs not yet supported.')
        except:
            pass
        # Generate tails matrix
        if self.IS_subspace:
            # Generate sparse mixing Hamiltonian
            assert hard_constraint_graph is not None
            assert isinstance(hard_constraint_graph, Graph)
            if code is not qubit:
                IS, num_IS = hard_constraint_graph.generate_independent_sets_qudit(self.code)
                self._diagonal_hamiltonian = np.zeros((num_IS, 1), dtype=float)
                for k in range(num_IS):
                    # Get state as bit string
                    # Multiply by weights matrix
                    # Sum result
                    weight = 0
                    where_rydberg = np.argwhere(self.hard_constraint_graph.independent_sets[k] == 0)
                    if len(where_rydberg) > 0:
                        where_rydberg = where_rydberg.T[0]
                        for (i, node1) in enumerate(where_rydberg):
                            for (j, node2) in enumerate(where_rydberg[i + 1:]):
                                weight += self.tails_graph.graph[node1][node2]['weight']
                    self._diagonal_hamiltonian[k, 0] = weight
                self._hamiltonian = sparse.csr_matrix(
                    (self._diagonal_hamiltonian.T[0], (np.arange(num_IS), np.arange(num_IS))),
                    shape=(num_IS, num_IS))
            else:
                # We have already solved for this information
                independent_sets = enumerate_independent_sets(self.hard_constraint_graph.graph)
                # Generate a list of integers corresponding to the independent sets in binary
                # All ones
                k = self.hard_constraint_graph.num_independent_sets - 2
                self.mis_size = 0
                self._diagonal_hamiltonian = np.zeros(self.hard_constraint_graph.num_independent_sets, dtype=float)
                self._diagonal_hamiltonian[-1] = 0
                for i in independent_sets:
                    weight = 0
                    where_rydberg = np.array(i)
                    if len(where_rydberg) > 0:
                        for (i, node1) in enumerate(where_rydberg):
                            for (j, node2) in enumerate(where_rydberg[i + 1:]):
                                weight += self.tails_graph.graph[node1][node2]['weight']
                    self._diagonal_hamiltonian[k] = weight
                    k -= 1

                self._hamiltonian = sparse.csr_matrix(
                    (self._diagonal_hamiltonian,
                     (np.arange(self.hard_constraint_graph.num_independent_sets), np.arange(self.hard_constraint_graph.num_independent_sets))),
                    shape=(self.hard_constraint_graph.num_independent_sets, self.hard_constraint_graph.num_independent_sets))
                self._diagonal_hamiltonian = np.expand_dims(self._diagonal_hamiltonian, axis=1)
        else:
            # Use full Hilbert space
            self._hamiltonian = None

    @property
    def hamiltonian(self):
        if self._hamiltonian is None:
            raise Exception
        return self.energies[0] * self._hamiltonian

    def left_multiply(self, state: State):
        return State(self.energies[0] * self._diagonal_hamiltonian * state, is_ket=state.is_ket,
                     IS_subspace=state.IS_subspace,
                     code=state.code, graph=self.hard_constraint_graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            print('Warning: right multiply functionality currently applies the operator and daggers the state.')
            return self.left_multiply(state).conj().T
        elif not self.IS_subspace:
            temp = np.zeros_like(state)
            # For each physical qubit
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.right_multiply(state, [i], self._operator)
                else:
                    ind = self.code.d ** i
                    out = np.zeros_like(state)
                    out = out.reshape(
                        (-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1), self.code.d, ind),
                        order='F')
                    state = state.reshape(
                        (-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1), self.code.d, ind),
                        order='F')
                    out[:, :, :, self.index, :] = state[:, :, :, self.index, :]
                    state = state.reshape(state_shape, order='F')
                    out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                         graph=self.hard_constraint_graph)
        else:
            return State(self.energies[0] * state * self._diagonal_hamiltonian.T, is_ket=state.is_ket,
                         IS_subspace=state.IS_subspace,
                         code=state.code, graph=self.hard_constraint_graph)

    def evolve(self, state: State, time):
        r"""
        Use reshape to efficiently implement evolution under :math:`H_B=\\sum_i X_i`
        """
        if not self.IS_subspace:
            # We don't want to modify the original s
            out = state.copy()
            for i in range(state.number_logical_qudits):
                # Note that self._operator is not necessarily involutary
                out = self.code.rotation(out, [i], self.energies[0] * time, self._operator)
            return out
        else:
            if state.is_ket:
                # Handle dimensions
                return State(np.exp(-1j * time * self.energies[0] * self._diagonal_hamiltonian) * state,
                             is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                             graph=self.hard_constraint_graph)
            else:
                exp_hamiltonian = np.exp(-1j * time * self.energies[0] * self._diagonal_hamiltonian)
                return State(exp_hamiltonian * state * exp_hamiltonian.conj().T,
                             is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                             graph=self.hard_constraint_graph)

5f90c9b74b10a616dee3fc1ac647f391d33dfe45
<mask token>


class HamiltonianGlobalPauli(object):
    <mask token>

    def evolve(self, state: State, alpha):
        if self.hamiltonian is None:
            """Initialize the Hamiltonian only once, as it is costly."""
            self.hamiltonian = tools.tensor_product([self._operator] *
                state.number_logical_qudits)
        return self.code.multiply(np.cos(alpha) * np.identity(state.
            dimension) - 1.0j * np.sin(alpha) * self.hamiltonian)

    def left_multiply(self, state: State):
        all_qubits = list(range(state.number_logical_qudits))
        return self.code.left_multiply(state, all_qubits, [self.pauli] *
            state.number_logical_qudits)

    def right_multiply(self, state: State):
        all_qubits = list(range(state.number_logical_qudits))
        return self.code.right_multiply(state, all_qubits, [self.pauli] *
            state.number_logical_qudits)


class HamiltonianBookatzPenalty(object):

    def __init__(self, code=qubit, energies=(1,)):
        self.code = code
        self.projector = np.identity(self.code.d ** self.code.n
            ) - self.code.code_space_projector
        self.energies = energies

    def evolve(self, state: State, time):
        for i in range(state.number_logical_qudits):
            state = self.code.rotation(state, [i], self.energies[0] * time,
                self.projector, is_idempotent=True)
        return state

    def left_multiply(self, state: State):
        out = np.zeros_like(state, dtype=np.complex128)
        for i in range(state.number_logical_qudits):
            out = out + self.code.left_multiply(state, [i], self.projector)
        return State(self.energies[0] * out, is_ket=state.is_ket,
            IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        out = np.zeros_like(state, dtype=np.complex128)
        for i in range(state.number_logical_qudits):
            out = out + self.code.right_multiply(state, [i], self.projector)
        return State(self.energies[0] * out, is_ket=state.is_ket,
            IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)


class HamiltonianMarvianPenalty(object):

    def __init__(self, Nx, Ny):
        super().__init__()
        self.Nx = Nx
        self.Ny = Ny
        self.n = 3 * Nx * Ny
        hp = np.zeros([2 ** self.n, 2 ** self.n])
        for i in range(int(self.Nx * self.Ny)):
            hp = hp + tools.tensor_product([tools.identity(i * 3), tools.Z(
                ), tools.Z(), tools.identity(self.n - i * 3 - 2)]
                ) + tools.tensor_product([tools.identity(i * 3 + 1), tools.
                X(), tools.X(), tools.identity(self.n - i * 3 - 3)])
        for j in range(self.Ny):
            for k in range(self.Nx):
                if k != self.Nx - 1:
                    hp = hp + tools.tensor_product([tools.identity(j * self
                        .Nx * 3 + k * 3), tools.X(), tools.identity(2),
                        tools.X(), tools.identity(self.n - (j * self.Nx * 3 +
                        k * 3) - 4)]) + tools.tensor_product([tools.
                        identity(j * self.Nx * 3 + k * 3 + 2), tools.Z(),
                        tools.identity(2), tools.Z(), tools.identity(self.n -
                        (j * self.Nx * 3 + k * 3 + 2) - 4)])
                if j != self.Ny - 1:
                    hp = hp + tools.tensor_product([tools.identity(j * self
                        .Nx * 3 + k * 3), tools.X(), tools.identity(3 *
                        self.Nx - 1), tools.X(), tools.identity(self.n - (j *
                        self.Nx * 3 + k * 3) - 3 * self.Nx - 1)]
                        ) + tools.tensor_product([tools.identity(j * self.
                        Nx * 3 + k * 3 + 2), tools.Z(), tools.identity(3 *
                        self.Nx - 1), tools.Z(), tools.identity(self.n - (j *
                        self.Nx * 3 + k * 3 + 2) - 3 * self.Nx - 1)])
        self.hamiltonian = -1 * hp


class HamiltonianHeisenberg(object):

    def __init__(self, graph: Graph, energies=(1, 1), subspace='all', code=
        qubit, IS_subspace=False):
        self.code = code
        self.graph = graph
        self.n = self.graph.n
        self.energies = energies
        self._is_diagonal = True
        assert code is qubit
        if IS_subspace:
            raise NotImplementedError
        self.IS_subspace = IS_subspace
        self._hamiltonian = None
        self.subspace = subspace
        if self.subspace == 'all':
            """Initialize the Hamiltonian."""
            self._hamiltonian_zz = None
            self._hamiltonian_xy = None
        else:
            from scipy.special import comb
            dim = comb(graph.n, (graph.n + self.subspace) / 2)
            if dim % 1 != 0:
                raise Exception('Invalid subspace for number of spins')
            dim = int(dim)
            hamiltonian_zz = np.zeros([dim, 1])
            rows = np.zeros(graph.n * dim, dtype=int)
            columns = np.zeros(graph.n * dim, dtype=int)
            entries = np.zeros(graph.n * dim, dtype=int)
            num_terms = 0
            states = np.zeros((dim, graph.n))
            for i in range(self.code.d ** (self.code.n * self.n)):
                nary = tools.int_to_nary(i, size=graph.n)
                if graph.n - np.sum(nary) == (self.subspace + graph.n) / 2:
                    states[num_terms, ...] = nary
                    num_terms += 1
            num_terms = 0
            for i in range(dim):
                zz = 0
                for a, b in self.graph.edges:
                    if b < a:
                        a, b = b, a
                    if np.abs(states[i, a] - states[i, b]) == 1:
                        zz -= 1
                    else:
                        zz += 1
                    if states[i, a] == 0 and states[i, b] == 1:
                        temp = states[i, ...].copy()
                        temp[a] = 1
                        temp[b] = 0
                        where_matched_temp = np.argwhere(np.sum(np.abs(
                            states - temp), axis=1) == 0).flatten()
                        entries[num_terms] = 1
                        rows[num_terms] = where_matched_temp
                        columns[num_terms] = i
                        entries[num_terms + 1] = 1
                        rows[num_terms + 1] = i
                        columns[num_terms + 1] = where_matched_temp
                        num_terms += 2
                hamiltonian_zz[i, 0] = zz
            self._hamiltonian_zz = sparse.csr_matrix((hamiltonian_zz.
                flatten(), (np.arange(dim), np.arange(dim))), shape=(dim, dim))
            self._hamiltonian_xy = sparse.csr_matrix((entries, (rows,
                columns)), shape=(dim, dim))
            self.states = states

    @property
    def hamiltonian(self):
        if (self._hamiltonian_zz is None or self._hamiltonian_xy is None and
            self.subspace == 'all'):
            z = np.expand_dims(np.diagonal(self.code.Z), axis=0).T

            def my_eye(n):
                return np.ones((np.asarray(self.code.d ** self.code.n) ** n, 1)
                    )
            hamiltonian_zz = np.zeros([self.code.d ** (self.code.n * self.n
                ), 1])
            for a, b in self.graph.edges:
                if b < a:
                    a, b = b, a
                hamiltonian_zz = hamiltonian_zz + self.graph.graph[a][b][
                    'weight'] * tools.tensor_product([my_eye(a), z, my_eye(
                    b - a - 1), z, my_eye(self.n - b - 1)], sparse=not self
                    ._is_diagonal)
            self._hamiltonian_zz = sparse.csr_matrix((hamiltonian_zz.
                flatten(), (np.arange(self.code.d ** (self.code.n * self.n)
                ), np.arange(self.code.d ** (self.code.n * self.n)))),
                shape=(self.code.d ** (self.code.n * self.n), self.code.d **
                (self.code.n * self.n)))
            rows = np.zeros(self.graph.n * self.code.d ** (self.code.n *
                self.n), dtype=int)
            columns = np.zeros(self.graph.n * self.code.d ** (self.code.n *
                self.n), dtype=int)
            entries = np.zeros(self.graph.n * self.code.d ** (self.code.n *
                self.n), dtype=int)
            num_terms = 0
            for i in range(self.code.d ** (self.code.n * self.n)):
                nary = tools.int_to_nary(i, size=self.graph.n)
                for a, b in self.graph.edges:
                    if b < a:
                        a, b = b, a
                    if nary[a] == 0 and nary[b] == 1:
                        temp = nary.copy()
                        temp[a] = 1
                        temp[b] = 0
                        temp = tools.nary_to_int(temp)
                        entries[num_terms] = 1
                        rows[num_terms] = temp
                        columns[num_terms] = i
                        entries[num_terms + 1] = 1
                        rows[num_terms + 1] = i
                        columns[num_terms + 1] = temp
                        num_terms += 2
            self._hamiltonian_xy = sparse.csr_matrix((entries, (rows,
                columns)), shape=(self.code.d ** (self.code.n * self.n), 
                self.code.d ** (self.code.n * self.n)))
        return self.energies[0] * self._hamiltonian_zz + self.energies[1
            ] * self._hamiltonian_xy

    def left_multiply(self, state: State):
        temp = np.zeros(state.shape, dtype=np.complex128)
        for edge in self.graph.edges:
            if self.energies[0] != 0:
                term = self.code.left_multiply(state, [edge[0], edge[1]], [
                    'X', 'X'])
                temp = temp + self.energies[0] * term
                term = self.code.left_multiply(state, [edge[0], edge[1]], [
                    'Y', 'Y'])
                temp = temp + self.energies[0] * term
            if self.energies[1] != 0:
                term = self.code.left_multiply(state, [edge[0], edge[1]], [
                    'Z', 'Z'])
                temp = temp + self.energies[1] * term
        return State(temp, is_ket=state.is_ket, IS_subspace=state.
            IS_subspace, code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        temp = np.zeros(state.shape, dtype=np.complex128)
        for edge in self.graph.edges:
            if self.energies[0] != 0:
                term = self.code.right_multiply(state, [edge[0], edge[1]],
                    ['X', 'X'])
                temp = temp + self.energies[0] * term
                term = self.code.right_multiply(state, [edge[0], edge[1]],
                    ['Y', 'Y'])
                temp = temp + self.energies[0] * term
            if self.energies[1] != 0:
                term = self.code.right_multiply(state, [edge[0], edge[1]],
                    ['Z', 'Z'])
                temp = temp + self.energies[1] * term
        return State(temp, is_ket=state.is_ket, IS_subspace=state.
            IS_subspace, code=state.code, graph=self.graph)

    def evolve(self, state: State, time):
        if not state.is_ket:
            exp_hamiltonian = expm(-1.0j * time * self.hamiltonian)
            return State(exp_hamiltonian @ state @ exp_hamiltonian.conj().T,
                is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=
                state.code, graph=self.graph)
        if state.is_ket:
            return State(expm_multiply(-1.0j * time * self.hamiltonian,
                state), is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                code=state.code, graph=self.graph)

    def cost_function(self, state: State):
        if state.is_ket:
            return np.real(np.vdot(state, self.hamiltonian * state))
        else:
            return np.real(np.squeeze(tools.trace(self.hamiltonian * state)))


class HamiltonianEnergyShift(object):

    def __init__(self, index: int=0, energies=(1,), code=qubit, IS_subspace
        =False, graph=None):
        """Default is that the first element in transition is the higher energy s."""
        self.index = index
        self.graph = graph
        self.energies = energies
        self.code = code
        if not self.code.logical_code:
            if not 0 <= self.index < self.code.d:
                raise Exception('Index exceeds qudit dimension.')
            self._operator = np.zeros((self.code.d, self.code.d))
            self._operator[self.index, self.index] = 1
        else:
            if self.index != 0 and self.index != 1:
                raise Exception(
                    'Logical codes are qubits, so index must be 0 or 1.')
            if self.index == 0:
                self._operator = self.code.Q
            elif self.index == 1:
                self._operator = self.code.P
        self.IS_subspace = IS_subspace
        if self.IS_subspace:
            assert graph is not None
            assert isinstance(graph, Graph)
            if code is not qubit:
                IS, num_IS = graph.independent_sets_qudit(self.code)
                self._diagonal_hamiltonian = np.zeros((num_IS, 1), dtype=float)
                for k in range(num_IS):
                    self._diagonal_hamiltonian[k, 0] = np.sum(IS[k, ...] ==
                        self.index)
                self._hamiltonian = sparse.csr_matrix((self.
                    _diagonal_hamiltonian.T[0], (np.arange(num_IS), np.
                    arange(num_IS))), shape=(num_IS, num_IS))
            else:
                independent_sets = enumerate_independent_sets(graph.graph)
                k = self.graph.num_independent_sets - 2
                self.mis_size = 0
                hamiltonian = np.zeros(self.graph.num_independent_sets,
                    dtype=float)
                hamiltonian[-1] = 0
                for i in independent_sets:
                    hamiltonian[k] = len(i)
                    k -= 1
                self._hamiltonian = sparse.csr_matrix((hamiltonian, (np.
                    arange(self.graph.num_independent_sets), np.arange(self
                    .graph.num_independent_sets))), shape=(self.graph.
                    num_independent_sets, self.graph.num_independent_sets))
        else:
            self._hamiltonian = None

    @property
    def hamiltonian(self):
        if self._hamiltonian is None:
            assert not self.IS_subspace
            try:
                assert self.graph is not None
            except AssertionError:
                print(
                    'self.graph must be not None to generate the Hamiltonian property.'
                    )
            self._hamiltonian = sparse.csr_matrix(((self.code.d * self.code
                .n) ** self.graph.n, (self.code.d * self.code.n) ** self.
                graph.n))
            for i in range(self.graph.n):
                self._hamiltonian = self._hamiltonian + tools.tensor_product([
                    sparse.identity((self.code.d * self.code.n) ** i), self
                    ._operator, sparse.identity((self.code.d * self.code.n) **
                    (self.graph.n - i - 1))], sparse=True)
        return self.energies[0] * self._hamiltonian

    def left_multiply(self, state: State):
        if not self.IS_subspace:
            temp = np.zeros_like(state)
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.left_multiply(state, [i], self.
                        _operator)
                elif not self.code.logical_code:
                    ind = self.code.d ** i
                    out = np.zeros_like(state, dtype=np.complex128)
                    if state.is_ket:
                        state = state.reshape((-1, self.code.d, ind), order='F'
                            )
                        out = out.reshape((-1, self.code.d, ind), order='F')
                        out[:, self.index, :] = state[:, self.index, :]
                        state = state.reshape(state_shape, order='F')
                        out = out.reshape(state_shape, order='F')
                    else:
                        out = out.reshape((-1, self.code.d, self.code.d **
                            (state.number_physical_qudits - 1), self.code.d,
                            ind), order='F')
                        state = state.reshape((-1, self.code.d, self.code.d **
                            (state.number_physical_qudits - 1), self.code.d,
                            ind), order='F')
                        out[:, self.index, :, :, :] = state[:, self.index,
                            :, :, :]
                        state = state.reshape(state_shape, order='F')
                        out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)
        else:
            return State(self.energies[0] * self._diagonal_hamiltonian *
                state, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            print(
                'Warning: right multiply functionality currently applies the operator and daggers the state.'
                )
            return self.left_multiply(state).conj().T
        elif not self.IS_subspace:
            temp = np.zeros_like(state)
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.right_multiply(state, [i], self
                        ._operator)
                else:
                    ind = self.code.d ** i
                    out = np.zeros_like(state)
                    out = out.reshape((-1, self.code.d, self.code.d ** (
                        state.number_physical_qudits - 1), self.code.d, ind
                        ), order='F')
                    state = state.reshape((-1, self.code.d, self.code.d **
                        (state.number_physical_qudits - 1), self.code.d,
                        ind), order='F')
                    out[:, :, :, self.index, :] = state[:, :, :, self.index, :]
                    state = state.reshape(state_shape, order='F')
                    out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)
        else:
            return State(self.energies[0] * state * self.
                _diagonal_hamiltonian.T, is_ket=state.is_ket, IS_subspace=
                state.IS_subspace, code=state.code, graph=self.graph)

    def evolve(self, state: State, time):
        """
        Use reshape to efficiently implement evolution under :math:`H_B=\\\\sum_i X_i`
        """
        if not self.IS_subspace:
            out = state.copy()
            for i in range(state.number_logical_qudits):
                out = self.code.rotation(out, [i], self.energies[0] * time,
                    self._operator)
            return out
        elif state.is_ket:
            return State(np.exp(-1.0j * time * self.energies[0] * self.
                _diagonal_hamiltonian) * state, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)
        else:
            exp_hamiltonian = np.exp(-1.0j * time * self.energies[0] * self
                ._diagonal_hamiltonian)
            return State(exp_hamiltonian * state * exp_hamiltonian.conj().T,
                is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=
                state.code, graph=self.graph)


class HamiltonianRydberg(object):

    def __init__(self, tails_graph: Graph, hard_constraint_graph=None,
        index: int=0, energies=(1,), code=qubit, IS_subspace=False):
        """Default is that the first element in transition is the higher energy s."""
        self.index = index
        self.energies = energies
        self.code = code
        self.tails_graph = tails_graph
        self.hard_constraint_graph = hard_constraint_graph
        if self.code.logical_code:
            raise Exception('Logical codes not supported.')
        self.IS_subspace = IS_subspace
        try:
            if (self.tails_graph.periodic or self.hard_constraint_graph.
                periodic):
                raise Exception('Periodic graphs not yet supported.')
        except:
            pass
        if self.IS_subspace:
            assert hard_constraint_graph is not None
            assert isinstance(hard_constraint_graph, Graph)
            if code is not qubit:
                IS, num_IS = (hard_constraint_graph.
                    generate_independent_sets_qudit(self.code))
                self._diagonal_hamiltonian = np.zeros((num_IS, 1), dtype=float)
                for k in range(num_IS):
                    weight = 0
                    where_rydberg = np.argwhere(self.hard_constraint_graph.
                        independent_sets[k] == 0)
                    if len(where_rydberg) > 0:
                        where_rydberg = where_rydberg.T[0]
                        for i, node1 in enumerate(where_rydberg):
                            for j, node2 in enumerate(where_rydberg[i + 1:]):
                                weight += self.tails_graph.graph[node1][node2][
                                    'weight']
                    self._diagonal_hamiltonian[k, 0] = weight
                self._hamiltonian = sparse.csr_matrix((self.
                    _diagonal_hamiltonian.T[0], (np.arange(num_IS), np.
                    arange(num_IS))), shape=(num_IS, num_IS))
            else:
                independent_sets = enumerate_independent_sets(self.
                    hard_constraint_graph.graph)
                k = self.hard_constraint_graph.num_independent_sets - 2
                self.mis_size = 0
                self._diagonal_hamiltonian = np.zeros(self.
                    hard_constraint_graph.num_independent_sets, dtype=float)
                self._diagonal_hamiltonian[-1] = 0
                for i in independent_sets:
                    weight = 0
                    where_rydberg = np.array(i)
                    if len(where_rydberg) > 0:
                        for i, node1 in enumerate(where_rydberg):
                            for j, node2 in enumerate(where_rydberg[i + 1:]):
                                weight += self.tails_graph.graph[node1][node2][
                                    'weight']
                    self._diagonal_hamiltonian[k] = weight
                    k -= 1
                self._hamiltonian = sparse.csr_matrix((self.
                    _diagonal_hamiltonian, (np.arange(self.
                    hard_constraint_graph.num_independent_sets), np.arange(
                    self.hard_constraint_graph.num_independent_sets))),
                    shape=(self.hard_constraint_graph.num_independent_sets,
                    self.hard_constraint_graph.num_independent_sets))
                self._diagonal_hamiltonian = np.expand_dims(self.
                    _diagonal_hamiltonian, axis=1)
        else:
            self._hamiltonian = None

    @property
    def hamiltonian(self):
        if self._hamiltonian is None:
            raise Exception
        return self.energies[0] * self._hamiltonian

    def left_multiply(self, state: State):
        return State(self.energies[0] * self._diagonal_hamiltonian * state,
            is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.
            code, graph=self.hard_constraint_graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            print(
                'Warning: right multiply functionality currently applies the operator and daggers the state.'
                )
            return self.left_multiply(state).conj().T
        elif not self.IS_subspace:
            temp = np.zeros_like(state)
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.right_multiply(state, [i], self
                        ._operator)
                else:
                    ind = self.code.d ** i
                    out = np.zeros_like(state)
                    out = out.reshape((-1, self.code.d, self.code.d ** (
                        state.number_physical_qudits - 1), self.code.d, ind
                        ), order='F')
                    state = state.reshape((-1, self.code.d, self.code.d **
                        (state.number_physical_qudits - 1), self.code.d,
                        ind), order='F')
                    out[:, :, :, self.index, :] = state[:, :, :, self.index, :]
                    state = state.reshape(state_shape, order='F')
                    out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                hard_constraint_graph)
        else:
            return State(self.energies[0] * state * self.
                _diagonal_hamiltonian.T, is_ket=state.is_ket, IS_subspace=
                state.IS_subspace, code=state.code, graph=self.
                hard_constraint_graph)

    def evolve(self, state: State, time):
        """
        Use reshape to efficiently implement evolution under :math:`H_B=\\\\sum_i X_i`
        """
        if not self.IS_subspace:
            out = state.copy()
            for i in range(state.number_logical_qudits):
                out = self.code.rotation(out, [i], self.energies[0] * time,
                    self._operator)
            return out
        elif state.is_ket:
            return State(np.exp(-1.0j * time * self.energies[0] * self.
                _diagonal_hamiltonian) * state, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                hard_constraint_graph)
        else:
            exp_hamiltonian = np.exp(-1.0j * time * self.energies[0] * self
                ._diagonal_hamiltonian)
            return State(exp_hamiltonian * state * exp_hamiltonian.conj().T,
                is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=
                state.code, graph=self.hard_constraint_graph)

<mask token>


class HamiltonianMaxCut(object):
    <mask token>
    <mask token>
    <mask token>
    <mask token>
    <mask token>
    <mask token>
    <mask token>
    <mask token>
    <mask token>
    <mask token>


class HamiltonianMIS(object):

    def __init__(self, graph: Graph, energies=(1, 1), code=qubit,
        IS_subspace=False):
        """
        Generate a vector corresponding to the diagonal of the MIS Hamiltonian.
        """
        if energies == (1, 1) and IS_subspace:
            energies = 1,
        self.code = code
        self.graph = graph
        self.n = self.graph.n
        self.energies = energies
        self.IS_subspace = IS_subspace
        self.optimization = 'max'
        if not self.IS_subspace:
            if tools.is_diagonal(self.code.Q):
                self._hamiltonian_edge_terms = np.zeros([1, (self.code.d **
                    self.code.n) ** self.n])
                self._hamiltonian_node_terms = np.zeros([1, (self.code.d **
                    self.code.n) ** self.n])
                self._is_diagonal = True
                Q = np.expand_dims(np.diagonal(self.code.Q), axis=0)

                def my_eye(n):
                    return np.ones(np.asarray(self.code.d ** self.code.n) ** n)
            else:
                self._hamiltonian_edge_terms = np.zeros([(self.code.d **
                    self.code.n) ** self.n, (self.code.d ** self.code.n) **
                    self.n])
                self._hamiltonian_node_terms = np.zeros([(self.code.d **
                    self.code.n) ** self.n, (self.code.d ** self.code.n) **
                    self.n])
                Q = np.expand_dims(np.diagonal(qubit.Q), axis=0)

                def my_eye(n):
                    return np.ones(np.asarray(qubit.d ** qubit.n) ** n)
                self._optimum_edge_terms = np.zeros([(qubit.d ** qubit.n) **
                    self.n, (qubit.d ** qubit.n) ** self.n])
                self._optimum_node_terms = np.zeros([(qubit.d ** qubit.n) **
                    self.n, (qubit.d ** qubit.n) ** self.n])
                for i, j in graph.graph.edges:
                    if j < i:
                        i, j = j, i
                    self._optimum_edge_terms = (self._optimum_edge_terms + 
                        graph.graph.edges[i, j]['weight'] * tools.
                        tensor_product([my_eye(i), Q, my_eye(j - i - 1), Q,
                        my_eye(self.n - j - 1)]))
                for i in graph.graph.nodes:
                    self._optimum_node_terms = (self._optimum_node_terms + 
                        graph.graph.nodes[i]['weight'] * tools.
                        tensor_product([my_eye(i), Q, my_eye(self.n - i - 1)]))
                self._is_diagonal = False
                Q = self.code.Q

                def my_eye(n):
                    return np.identity(np.asarray(self.code.d ** self.code.
                        n) ** n)
            for i, j in graph.graph.edges:
                if j < i:
                    i, j = j, i
                self._hamiltonian_edge_terms = (self.
                    _hamiltonian_edge_terms + graph.graph.edges[i, j][
                    'weight'] * tools.tensor_product([my_eye(i), Q, my_eye(
                    j - i - 1), Q, my_eye(self.n - j - 1)]))
            for i in graph.graph.nodes:
                self._hamiltonian_node_terms = (self.
                    _hamiltonian_node_terms + graph.graph.nodes[i]['weight'
                    ] * tools.tensor_product([my_eye(i), Q, my_eye(self.n -
                    i - 1)]))
            self._hamiltonian_node_terms = self._hamiltonian_node_terms.T
            self._hamiltonian_edge_terms = self._hamiltonian_edge_terms.T
            if self._is_diagonal:
                self._optimum_edge_terms = self._hamiltonian_edge_terms
                self._optimum_node_terms = self._hamiltonian_node_terms
                self._diagonal_hamiltonian_edge_terms = (self.
                    _hamiltonian_edge_terms.copy())
                self._diagonal_hamiltonian_node_terms = (self.
                    _hamiltonian_node_terms.copy())
                self._hamiltonian_node_terms = sparse.csr_matrix((self.
                    _hamiltonian_node_terms.flatten(), (np.arange(self.code
                    .d ** (self.code.n * self.n)), np.arange(self.code.d **
                    (self.code.n * self.n)))), shape=(self.code.d ** (self.
                    code.n * self.n), self.code.d ** (self.code.n * self.n)))
                self._hamiltonian_edge_terms = sparse.csr_matrix((self.
                    _hamiltonian_edge_terms.flatten(), (np.arange(self.code
                    .d ** (self.code.n * self.n)), np.arange(self.code.d **
                    (self.code.n * self.n)))), shape=(self.code.d ** (self.
                    code.n * self.n), self.code.d ** (self.code.n * self.n)))
            else:
                self._diagonal_hamiltonian_edge_terms = (self.
                    _hamiltonian_edge_terms.copy())
                self._diagonal_hamiltonian_node_terms = (self.
                    _hamiltonian_node_terms.copy())
                self._hamiltonian_edge_terms = sparse.csr_matrix(self.
                    _hamiltonian_edge_terms)
                self._hamiltonian_node_terms = sparse.csr_matrix(self.
                    _hamiltonian_node_terms)
            self._left_acting_hamiltonian_edge_terms = None
            self._right_acting_hamiltonian_edge_terms = None
        else:
            self._is_diagonal = True
            if not (self.code == qubit or self.code == rydberg):
                raise NotImplementedError(
                    'IS subspace only implemented for qubit and Rydberg codes.'
                    )
            if self.code == qubit:
                node_weights = np.asarray([self.graph.graph.nodes[i][
                    'weight'] for i in range(self.graph.n)])
                independent_sets = enumerate_independent_sets(self.graph.graph)
                k = self.graph.num_independent_sets - 2
                self.mis_size = 0
                C = np.zeros(self.graph.num_independent_sets, dtype=float)
                C[-1] = 0
                for i in independent_sets:
                    C[k] = np.sum([node_weights[j] for j in i])
                    k -= 1
                self._hamiltonian = sparse.csr_matrix((C, (np.arange(self.
                    graph.num_independent_sets), np.arange(self.graph.
                    num_independent_sets))))
                C = np.expand_dims(C, axis=0).T
            elif self.code == rydberg:
                independent_sets, num_IS = (self.graph.
                    generate_independent_sets_qudit(self.code))
                node_weights = np.asarray([self.graph.graph.nodes[i][
                    'weight'] for i in range(self.graph.n)])
                C = np.zeros((num_IS, 1), dtype=np.complex128)
                for k in independent_sets:
                    C[k, 0] = np.sum((independent_sets[k][2] == 0) *
                        node_weights)
            self._diagonal_hamiltonian_node_terms = C
            C = C.flatten()
            self._hamiltonian_node_terms = sparse.csr_matrix((C, (np.arange
                (len(C)), np.arange(len(C)))), shape=(len(C), len(C)))
        self._left_acting_hamiltonian_node_terms = None
        self._right_acting_hamiltonian_node_terms = None

    @property
    def hamiltonian(self):
        if not self.IS_subspace:
            return self.energies[0
                ] * self._hamiltonian_node_terms - self.energies[1
                ] * self._hamiltonian_edge_terms
        else:
            return self.energies[0] * self._hamiltonian_node_terms

    @property
    def evolution_operator(self, vector_space='hilbert'):
        if vector_space != 'hilbert' and vector_space != 'liouville':
            raise Exception(
                'Attribute vector_space must be hilbert or liouville')
        if vector_space == 'liouville':
            if (self._left_acting_hamiltonian_node_terms is None and not
                self.IS_subspace):
                self._left_acting_hamiltonian_edge_terms = sparse.kron(sparse
                    .identity(self._hamiltonian_node_terms.shape[0]), self.
                    _hamiltonian_edge_terms)
                self._right_acting_hamiltonian_edge_terms = sparse.kron(self
                    ._hamiltonian_edge_terms.T, sparse.identity(self.
                    _hamiltonian_edge_terms.shape[0]))
                self._left_acting_hamiltonian_node_terms = sparse.kron(sparse
                    .identity(self._hamiltonian_node_terms.shape[0]), self.
                    _hamiltonian_node_terms)
                self._right_acting_hamiltonian_node_terms = sparse.kron(self
                    ._hamiltonian_node_terms.T, sparse.identity(self.
                    _hamiltonian_node_terms.shape[0]))
            elif self._left_acting_hamiltonian_node_terms is None and self.IS_subspace:
                self._left_acting_hamiltonian_node_terms = sparse.kron(sparse
                    .identity(self._hamiltonian_node_terms.shape[0]), self.
                    _hamiltonian_node_terms)
                self._right_acting_hamiltonian_node_terms = sparse.kron(self
                    ._hamiltonian_node_terms.T, sparse.identity(self.
                    _hamiltonian_node_terms.shape[0]))
            if not self.IS_subspace:
                return -1.0j * (self.energies[0] * self.
                    _left_acting_hamiltonian_node_terms - self.energies[1] *
                    self._left_acting_hamiltonian_edge_terms) + 1.0j * (
                    self.energies[0] * self.
                    _right_acting_hamiltonian_node_terms - self.energies[1] *
                    self._right_acting_hamiltonian_edge_terms)
            elif self.IS_subspace:
                return -1.0j * (self.energies[0] * self.
                    _left_acting_hamiltonian_node_terms) + 1.0j * (self.
                    energies[0] * self._right_acting_hamiltonian_node_terms)
        else:
            return -1.0j * self.hamiltonian

    @property
    def _diagonal_hamiltonian(self):
        if not self.IS_subspace:
            return self.energies[0
                ] * self._diagonal_hamiltonian_node_terms - self.energies[1
                ] * self._diagonal_hamiltonian_edge_terms
        else:
            return self.energies[0] * self._diagonal_hamiltonian_node_terms

    @property
    def optimum(self):
        if self._is_diagonal:
            return np.max(self._diagonal_hamiltonian).real
        else:
            raise NotImplementedError(
                'Optimum unknown for non-diagonal Hamiltonians')

    def evolve(self, state: State, time):
        if state.is_ket:
            if self._is_diagonal:
                return State(np.exp(-1.0j * time * self.
                    _diagonal_hamiltonian) * state, is_ket=state.is_ket,
                    IS_subspace=state.IS_subspace, code=state.code, graph=
                    self.graph)
            else:
                return State(expm_multiply(-1.0j * time * self.hamiltonian,
                    state), is_ket=state.is_ket, IS_subspace=state.
                    IS_subspace, code=state.code, graph=self.graph)
        elif self._is_diagonal:
            return State(np.exp(-1.0j * time * self._diagonal_hamiltonian) *
                state * np.exp(1.0j * time * self._diagonal_hamiltonian).T,
                is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=
                state.code, graph=self.graph)
        else:
            temp = expm(-1.0j * time * self.hamiltonian)
            return State(temp @ state @ temp.conj().T, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)

    def left_multiply(self, state: State):
        if self._is_diagonal:
            return State(self._diagonal_hamiltonian * state, is_ket=state.
                is_ket, IS_subspace=state.IS_subspace, code=state.code,
                graph=self.graph)
        else:
            return State(self.hamiltonian @ state, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            if self._is_diagonal:
                return State(state.conj().T * self._diagonal_hamiltonian.T,
                    is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                    code=state.code, graph=self.graph)
            else:
                return State((state @ self.hamiltonian.T).conj().T, is_ket=
                    state.is_ket, IS_subspace=state.IS_subspace, code=state
                    .code, graph=self.graph)
        elif self._is_diagonal:
            return State(state * self._diagonal_hamiltonian.T, is_ket=state
                .is_ket, IS_subspace=state.IS_subspace, code=state.code,
                graph=self.graph)
        else:
            return State(state @ self.hamiltonian.T, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)

    def cost_function(self, state: State):
        if state.is_ket:
            if self._is_diagonal:
                return np.real(np.vdot(state, self._diagonal_hamiltonian *
                    state))
            else:
                return np.real(np.vdot(state, self.hamiltonian @ state))
        elif self._is_diagonal:
            return np.real(np.squeeze(tools.trace(self.
                _diagonal_hamiltonian * state)))
        else:
            return np.real(np.squeeze(tools.trace(self.hamiltonian @ state)))

    def optimum_overlap(self, state: State):
        if self._is_diagonal:
            optimum_indices = np.argwhere(self._diagonal_hamiltonian ==
                self.optimum).T[0]
            optimum = np.zeros(self._diagonal_hamiltonian.shape)
            optimum[optimum_indices] = 1
        else:
            raise NotImplementedError(
                'Optimum overlap not implemented for non-diagonal Hamiltonians'
                )
        if state.is_ket:
            if self._is_diagonal:
                return np.real(np.vdot(state, optimum * state))
            else:
                return np.real(np.vdot(state, self.hamiltonian @ state))
        elif self._is_diagonal:
            return np.real(np.squeeze(tools.trace(optimum * state)))
        else:
            return np.real(np.squeeze(tools.trace(self.hamiltonian @ state)))

    def approximation_ratio(self, state: State):
        return self.cost_function(state) / self.optimum


class HamiltonianGlobalPauli(object):

    def __init__(self, pauli: str='X', code=qubit):
        self.code = code
        self.pauli = pauli
        if self.pauli == 'X':
            self._operator = self.code.X
        elif self.pauli == 'Y':
            self._operator = self.code.Y
        elif self.pauli == 'Z':
            self._operator = self.code.Z
        self.hamiltonian = None

    def evolve(self, state: State, alpha):
        if self.hamiltonian is None:
            """Initialize the Hamiltonian only once, as it is costly."""
            self.hamiltonian = tools.tensor_product([self._operator] *
                state.number_logical_qudits)
        return self.code.multiply(np.cos(alpha) * np.identity(state.
            dimension) - 1.0j * np.sin(alpha) * self.hamiltonian)

    def left_multiply(self, state: State):
        all_qubits = list(range(state.number_logical_qudits))
        return self.code.left_multiply(state, all_qubits, [self.pauli] *
            state.number_logical_qudits)

    def right_multiply(self, state: State):
        all_qubits = list(range(state.number_logical_qudits))
        return self.code.right_multiply(state, all_qubits, [self.pauli] *
            state.number_logical_qudits)


class HamiltonianBookatzPenalty(object):

    def __init__(self, code=qubit, energies=(1,)):
        self.code = code
        self.projector = np.identity(self.code.d ** self.code.n
            ) - self.code.code_space_projector
        self.energies = energies

    def evolve(self, state: State, time):
        for i in range(state.number_logical_qudits):
            state = self.code.rotation(state, [i], self.energies[0] * time,
                self.projector, is_idempotent=True)
        return state

    def left_multiply(self, state: State):
        out = np.zeros_like(state, dtype=np.complex128)
        for i in range(state.number_logical_qudits):
            out = out + self.code.left_multiply(state, [i], self.projector)
        return State(self.energies[0] * out, is_ket=state.is_ket,
            IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        out = np.zeros_like(state, dtype=np.complex128)
        for i in range(state.number_logical_qudits):
            out = out + self.code.right_multiply(state, [i], self.projector)
        return State(self.energies[0] * out, is_ket=state.is_ket,
            IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)


class HamiltonianMarvianPenalty(object):

    def __init__(self, Nx, Ny):
        super().__init__()
        self.Nx = Nx
        self.Ny = Ny
        self.n = 3 * Nx * Ny
        hp = np.zeros([2 ** self.n, 2 ** self.n])
        for i in range(int(self.Nx * self.Ny)):
            hp = hp + tools.tensor_product([tools.identity(i * 3), tools.Z(
                ), tools.Z(), tools.identity(self.n - i * 3 - 2)]
                ) + tools.tensor_product([tools.identity(i * 3 + 1), tools.
                X(), tools.X(), tools.identity(self.n - i * 3 - 3)])
        for j in range(self.Ny):
            for k in range(self.Nx):
                if k != self.Nx - 1:
                    hp = hp + tools.tensor_product([tools.identity(j * self
                        .Nx * 3 + k * 3), tools.X(), tools.identity(2),
                        tools.X(), tools.identity(self.n - (j * self.Nx * 3 +
                        k * 3) - 4)]) + tools.tensor_product([tools.
                        identity(j * self.Nx * 3 + k * 3 + 2), tools.Z(),
                        tools.identity(2), tools.Z(), tools.identity(self.n -
                        (j * self.Nx * 3 + k * 3 + 2) - 4)])
                if j != self.Ny - 1:
                    hp = hp + tools.tensor_product([tools.identity(j * self
                        .Nx * 3 + k * 3), tools.X(), tools.identity(3 *
                        self.Nx - 1), tools.X(), tools.identity(self.n - (j *
                        self.Nx * 3 + k * 3) - 3 * self.Nx - 1)]
                        ) + tools.tensor_product([tools.identity(j * self.
                        Nx * 3 + k * 3 + 2), tools.Z(), tools.identity(3 *
                        self.Nx - 1), tools.Z(), tools.identity(self.n - (j *
                        self.Nx * 3 + k * 3 + 2) - 3 * self.Nx - 1)])
        self.hamiltonian = -1 * hp


class HamiltonianHeisenberg(object):

    def __init__(self, graph: Graph, energies=(1, 1), subspace='all', code=
        qubit, IS_subspace=False):
        self.code = code
        self.graph = graph
        self.n = self.graph.n
        self.energies = energies
        self._is_diagonal = True
        assert code is qubit
        if IS_subspace:
            raise NotImplementedError
        self.IS_subspace = IS_subspace
        self._hamiltonian = None
        self.subspace = subspace
        if self.subspace == 'all':
            """Initialize the Hamiltonian."""
            self._hamiltonian_zz = None
            self._hamiltonian_xy = None
        else:
            from scipy.special import comb
            dim = comb(graph.n, (graph.n + self.subspace) / 2)
            if dim % 1 != 0:
                raise Exception('Invalid subspace for number of spins')
            dim = int(dim)
            hamiltonian_zz = np.zeros([dim, 1])
            rows = np.zeros(graph.n * dim, dtype=int)
            columns = np.zeros(graph.n * dim, dtype=int)
            entries = np.zeros(graph.n * dim, dtype=int)
            num_terms = 0
            states = np.zeros((dim, graph.n))
            for i in range(self.code.d ** (self.code.n * self.n)):
                nary = tools.int_to_nary(i, size=graph.n)
                if graph.n - np.sum(nary) == (self.subspace + graph.n) / 2:
                    states[num_terms, ...] = nary
                    num_terms += 1
            num_terms = 0
            for i in range(dim):
                zz = 0
                for a, b in self.graph.edges:
                    if b < a:
                        a, b = b, a
                    if np.abs(states[i, a] - states[i, b]) == 1:
                        zz -= 1
                    else:
                        zz += 1
                    if states[i, a] == 0 and states[i, b] == 1:
                        temp = states[i, ...].copy()
                        temp[a] = 1
                        temp[b] = 0
                        where_matched_temp = np.argwhere(np.sum(np.abs(
                            states - temp), axis=1) == 0).flatten()
                        entries[num_terms] = 1
                        rows[num_terms] = where_matched_temp
                        columns[num_terms] = i
                        entries[num_terms + 1] = 1
                        rows[num_terms + 1] = i
                        columns[num_terms + 1] = where_matched_temp
                        num_terms += 2
                hamiltonian_zz[i, 0] = zz
            self._hamiltonian_zz = sparse.csr_matrix((hamiltonian_zz.
                flatten(), (np.arange(dim), np.arange(dim))), shape=(dim, dim))
            self._hamiltonian_xy = sparse.csr_matrix((entries, (rows,
                columns)), shape=(dim, dim))
            self.states = states

    @property
    def hamiltonian(self):
        if (self._hamiltonian_zz is None or self._hamiltonian_xy is None and
            self.subspace == 'all'):
            z = np.expand_dims(np.diagonal(self.code.Z), axis=0).T

            def my_eye(n):
                return np.ones((np.asarray(self.code.d ** self.code.n) ** n, 1)
                    )
            hamiltonian_zz = np.zeros([self.code.d ** (self.code.n * self.n
                ), 1])
            for a, b in self.graph.edges:
                if b < a:
                    a, b = b, a
                hamiltonian_zz = hamiltonian_zz + self.graph.graph[a][b][
                    'weight'] * tools.tensor_product([my_eye(a), z, my_eye(
                    b - a - 1), z, my_eye(self.n - b - 1)], sparse=not self
                    ._is_diagonal)
            self._hamiltonian_zz = sparse.csr_matrix((hamiltonian_zz.
                flatten(), (np.arange(self.code.d ** (self.code.n * self.n)
                ), np.arange(self.code.d ** (self.code.n * self.n)))),
                shape=(self.code.d ** (self.code.n * self.n), self.code.d **
                (self.code.n * self.n)))
            rows = np.zeros(self.graph.n * self.code.d ** (self.code.n *
                self.n), dtype=int)
            columns = np.zeros(self.graph.n * self.code.d ** (self.code.n *
                self.n), dtype=int)
            entries = np.zeros(self.graph.n * self.code.d ** (self.code.n *
                self.n), dtype=int)
            num_terms = 0
            for i in range(self.code.d ** (self.code.n * self.n)):
                nary = tools.int_to_nary(i, size=self.graph.n)
                for a, b in self.graph.edges:
                    if b < a:
                        a, b = b, a
                    if nary[a] == 0 and nary[b] == 1:
                        temp = nary.copy()
                        temp[a] = 1
                        temp[b] = 0
                        temp = tools.nary_to_int(temp)
                        entries[num_terms] = 1
                        rows[num_terms] = temp
                        columns[num_terms] = i
                        entries[num_terms + 1] = 1
                        rows[num_terms + 1] = i
                        columns[num_terms + 1] = temp
                        num_terms += 2
            self._hamiltonian_xy = sparse.csr_matrix((entries, (rows,
                columns)), shape=(self.code.d ** (self.code.n * self.n), 
                self.code.d ** (self.code.n * self.n)))
        return self.energies[0] * self._hamiltonian_zz + self.energies[1
            ] * self._hamiltonian_xy

    def left_multiply(self, state: State):
        temp = np.zeros(state.shape, dtype=np.complex128)
        for edge in self.graph.edges:
            if self.energies[0] != 0:
                term = self.code.left_multiply(state, [edge[0], edge[1]], [
                    'X', 'X'])
                temp = temp + self.energies[0] * term
                term = self.code.left_multiply(state, [edge[0], edge[1]], [
                    'Y', 'Y'])
                temp = temp + self.energies[0] * term
            if self.energies[1] != 0:
                term = self.code.left_multiply(state, [edge[0], edge[1]], [
                    'Z', 'Z'])
                temp = temp + self.energies[1] * term
        return State(temp, is_ket=state.is_ket, IS_subspace=state.
            IS_subspace, code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        temp = np.zeros(state.shape, dtype=np.complex128)
        for edge in self.graph.edges:
            if self.energies[0] != 0:
                term = self.code.right_multiply(state, [edge[0], edge[1]],
                    ['X', 'X'])
                temp = temp + self.energies[0] * term
                term = self.code.right_multiply(state, [edge[0], edge[1]],
                    ['Y', 'Y'])
                temp = temp + self.energies[0] * term
            if self.energies[1] != 0:
                term = self.code.right_multiply(state, [edge[0], edge[1]],
                    ['Z', 'Z'])
                temp = temp + self.energies[1] * term
        return State(temp, is_ket=state.is_ket, IS_subspace=state.
            IS_subspace, code=state.code, graph=self.graph)

    def evolve(self, state: State, time):
        if not state.is_ket:
            exp_hamiltonian = expm(-1.0j * time * self.hamiltonian)
            return State(exp_hamiltonian @ state @ exp_hamiltonian.conj().T,
                is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=
                state.code, graph=self.graph)
        if state.is_ket:
            return State(expm_multiply(-1.0j * time * self.hamiltonian,
                state), is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                code=state.code, graph=self.graph)

    def cost_function(self, state: State):
        if state.is_ket:
            return np.real(np.vdot(state, self.hamiltonian * state))
        else:
            return np.real(np.squeeze(tools.trace(self.hamiltonian * state)))


class HamiltonianEnergyShift(object):

    def __init__(self, index: int=0, energies=(1,), code=qubit, IS_subspace
        =False, graph=None):
        """Default is that the first element in transition is the higher energy s."""
        self.index = index
        self.graph = graph
        self.energies = energies
        self.code = code
        if not self.code.logical_code:
            if not 0 <= self.index < self.code.d:
                raise Exception('Index exceeds qudit dimension.')
            self._operator = np.zeros((self.code.d, self.code.d))
            self._operator[self.index, self.index] = 1
        else:
            if self.index != 0 and self.index != 1:
                raise Exception(
                    'Logical codes are qubits, so index must be 0 or 1.')
            if self.index == 0:
                self._operator = self.code.Q
            elif self.index == 1:
                self._operator = self.code.P
        self.IS_subspace = IS_subspace
        if self.IS_subspace:
            assert graph is not None
            assert isinstance(graph, Graph)
            if code is not qubit:
                IS, num_IS = graph.independent_sets_qudit(self.code)
                self._diagonal_hamiltonian = np.zeros((num_IS, 1), dtype=float)
                for k in range(num_IS):
                    self._diagonal_hamiltonian[k, 0] = np.sum(IS[k, ...] ==
                        self.index)
                self._hamiltonian = sparse.csr_matrix((self.
                    _diagonal_hamiltonian.T[0], (np.arange(num_IS), np.
                    arange(num_IS))), shape=(num_IS, num_IS))
            else:
                independent_sets = enumerate_independent_sets(graph.graph)
                k = self.graph.num_independent_sets - 2
                self.mis_size = 0
                hamiltonian = np.zeros(self.graph.num_independent_sets,
                    dtype=float)
                hamiltonian[-1] = 0
                for i in independent_sets:
                    hamiltonian[k] = len(i)
                    k -= 1
                self._hamiltonian = sparse.csr_matrix((hamiltonian, (np.
                    arange(self.graph.num_independent_sets), np.arange(self
                    .graph.num_independent_sets))), shape=(self.graph.
                    num_independent_sets, self.graph.num_independent_sets))
        else:
            self._hamiltonian = None

    @property
    def hamiltonian(self):
        if self._hamiltonian is None:
            assert not self.IS_subspace
            try:
                assert self.graph is not None
            except AssertionError:
                print(
                    'self.graph must be not None to generate the Hamiltonian property.'
                    )
            self._hamiltonian = sparse.csr_matrix(((self.code.d * self.code
                .n) ** self.graph.n, (self.code.d * self.code.n) ** self.
                graph.n))
            for i in range(self.graph.n):
                self._hamiltonian = self._hamiltonian + tools.tensor_product([
                    sparse.identity((self.code.d * self.code.n) ** i), self
                    ._operator, sparse.identity((self.code.d * self.code.n) **
                    (self.graph.n - i - 1))], sparse=True)
        return self.energies[0] * self._hamiltonian

    def left_multiply(self, state: State):
        if not self.IS_subspace:
            temp = np.zeros_like(state)
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.left_multiply(state, [i], self.
                        _operator)
                elif not self.code.logical_code:
                    ind = self.code.d ** i
                    out = np.zeros_like(state, dtype=np.complex128)
                    if state.is_ket:
                        state = state.reshape((-1, self.code.d, ind), order='F'
                            )
                        out = out.reshape((-1, self.code.d, ind), order='F')
                        out[:, self.index, :] = state[:, self.index, :]
                        state = state.reshape(state_shape, order='F')
                        out = out.reshape(state_shape, order='F')
                    else:
                        out = out.reshape((-1, self.code.d, self.code.d **
                            (state.number_physical_qudits - 1), self.code.d,
                            ind), order='F')
                        state = state.reshape((-1, self.code.d, self.code.d **
                            (state.number_physical_qudits - 1), self.code.d,
                            ind), order='F')
                        out[:, self.index, :, :, :] = state[:, self.index,
                            :, :, :]
                        state = state.reshape(state_shape, order='F')
                        out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)
        else:
            return State(self.energies[0] * self._diagonal_hamiltonian *
                state, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            print(
                'Warning: right multiply functionality currently applies the operator and daggers the state.'
                )
            return self.left_multiply(state).conj().T
        elif not self.IS_subspace:
            temp = np.zeros_like(state)
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.right_multiply(state, [i], self
                        ._operator)
                else:
                    ind = self.code.d ** i
                    out = np.zeros_like(state)
                    out = out.reshape((-1, self.code.d, self.code.d ** (
                        state.number_physical_qudits - 1), self.code.d, ind
                        ), order='F')
                    state = state.reshape((-1, self.code.d, self.code.d **
                        (state.number_physical_qudits - 1), self.code.d,
                        ind), order='F')
                    out[:, :, :, self.index, :] = state[:, :, :, self.index, :]
                    state = state.reshape(state_shape, order='F')
                    out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)
        else:
            return State(self.energies[0] * state * self.
                _diagonal_hamiltonian.T, is_ket=state.is_ket, IS_subspace=
                state.IS_subspace, code=state.code, graph=self.graph)

    def evolve(self, state: State, time):
        """
        Use reshape to efficiently implement evolution under :math:`H_B=\\\\sum_i X_i`
        """
        if not self.IS_subspace:
            out = state.copy()
            for i in range(state.number_logical_qudits):
                out = self.code.rotation(out, [i], self.energies[0] * time,
                    self._operator)
            return out
        elif state.is_ket:
            return State(np.exp(-1.0j * time * self.energies[0] * self.
                _diagonal_hamiltonian) * state, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)
        else:
            exp_hamiltonian = np.exp(-1.0j * time * self.energies[0] * self
                ._diagonal_hamiltonian)
            return State(exp_hamiltonian * state * exp_hamiltonian.conj().T,
                is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=
                state.code, graph=self.graph)


class HamiltonianRydberg(object):

    def __init__(self, tails_graph: Graph, hard_constraint_graph=None,
        index: int=0, energies=(1,), code=qubit, IS_subspace=False):
        """Default is that the first element in transition is the higher energy s."""
        self.index = index
        self.energies = energies
        self.code = code
        self.tails_graph = tails_graph
        self.hard_constraint_graph = hard_constraint_graph
        if self.code.logical_code:
            raise Exception('Logical codes not supported.')
        self.IS_subspace = IS_subspace
        try:
            if (self.tails_graph.periodic or self.hard_constraint_graph.
                periodic):
                raise Exception('Periodic graphs not yet supported.')
        except:
            pass
        if self.IS_subspace:
            assert hard_constraint_graph is not None
            assert isinstance(hard_constraint_graph, Graph)
            if code is not qubit:
                IS, num_IS = (hard_constraint_graph.
                    generate_independent_sets_qudit(self.code))
                self._diagonal_hamiltonian = np.zeros((num_IS, 1), dtype=float)
                for k in range(num_IS):
                    weight = 0
                    where_rydberg = np.argwhere(self.hard_constraint_graph.
                        independent_sets[k] == 0)
                    if len(where_rydberg) > 0:
                        where_rydberg = where_rydberg.T[0]
                        for i, node1 in enumerate(where_rydberg):
                            for j, node2 in enumerate(where_rydberg[i + 1:]):
                                weight += self.tails_graph.graph[node1][node2][
                                    'weight']
                    self._diagonal_hamiltonian[k, 0] = weight
                self._hamiltonian = sparse.csr_matrix((self.
                    _diagonal_hamiltonian.T[0], (np.arange(num_IS), np.
                    arange(num_IS))), shape=(num_IS, num_IS))
            else:
                independent_sets = enumerate_independent_sets(self.
                    hard_constraint_graph.graph)
                k = self.hard_constraint_graph.num_independent_sets - 2
                self.mis_size = 0
                self._diagonal_hamiltonian = np.zeros(self.
                    hard_constraint_graph.num_independent_sets, dtype=float)
                self._diagonal_hamiltonian[-1] = 0
                for i in independent_sets:
                    weight = 0
                    where_rydberg = np.array(i)
                    if len(where_rydberg) > 0:
                        for i, node1 in enumerate(where_rydberg):
                            for j, node2 in enumerate(where_rydberg[i + 1:]):
                                weight += self.tails_graph.graph[node1][node2][
                                    'weight']
                    self._diagonal_hamiltonian[k] = weight
                    k -= 1
                self._hamiltonian = sparse.csr_matrix((self.
                    _diagonal_hamiltonian, (np.arange(self.
                    hard_constraint_graph.num_independent_sets), np.arange(
                    self.hard_constraint_graph.num_independent_sets))),
                    shape=(self.hard_constraint_graph.num_independent_sets,
                    self.hard_constraint_graph.num_independent_sets))
                self._diagonal_hamiltonian = np.expand_dims(self.
                    _diagonal_hamiltonian, axis=1)
        else:
            self._hamiltonian = None

    @property
    def hamiltonian(self):
        if self._hamiltonian is None:
            raise Exception
        return self.energies[0] * self._hamiltonian

    def left_multiply(self, state: State):
        return State(self.energies[0] * self._diagonal_hamiltonian * state,
            is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.
            code, graph=self.hard_constraint_graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            print(
                'Warning: right multiply functionality currently applies the operator and daggers the state.'
                )
            return self.left_multiply(state).conj().T
        elif not self.IS_subspace:
            temp = np.zeros_like(state)
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.right_multiply(state, [i], self
                        ._operator)
                else:
                    ind = self.code.d ** i
                    out = np.zeros_like(state)
                    out = out.reshape((-1, self.code.d, self.code.d ** (
                        state.number_physical_qudits - 1), self.code.d, ind
                        ), order='F')
                    state = state.reshape((-1, self.code.d, self.code.d **
                        (state.number_physical_qudits - 1), self.code.d,
                        ind), order='F')
                    out[:, :, :, self.index, :] = state[:, :, :, self.index, :]
                    state = state.reshape(state_shape, order='F')
                    out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                hard_constraint_graph)
        else:
            return State(self.energies[0] * state * self.
                _diagonal_hamiltonian.T, is_ket=state.is_ket, IS_subspace=
                state.IS_subspace, code=state.code, graph=self.
                hard_constraint_graph)

    def evolve(self, state: State, time):
        """
        Use reshape to efficiently implement evolution under :math:`H_B=\\\\sum_i X_i`
        """
        if not self.IS_subspace:
            out = state.copy()
            for i in range(state.number_logical_qudits):
                out = self.code.rotation(out, [i], self.energies[0] * time,
                    self._operator)
            return out
        elif state.is_ket:
            return State(np.exp(-1.0j * time * self.energies[0] * self.
                _diagonal_hamiltonian) * state, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                hard_constraint_graph)
        else:
            exp_hamiltonian = np.exp(-1.0j * time * self.energies[0] * self
                ._diagonal_hamiltonian)
            return State(exp_hamiltonian * state * exp_hamiltonian.conj().T,
                is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=
                state.code, graph=self.hard_constraint_graph)

<mask token>


class HamiltonianMaxCut(object):
    <mask token>
    <mask token>
    <mask token>
    <mask token>
    <mask token>
    <mask token>
    <mask token>

    def cost_function(self, state: State):
        if state.is_ket:
            if self._is_diagonal:
                return np.real(np.vdot(state, self._diagonal_hamiltonian *
                    state))
            else:
                return np.real(np.vdot(state, self.hamiltonian @ state))
        elif self._is_diagonal:
            return np.real(np.squeeze(tools.trace(self.
                _diagonal_hamiltonian * state)))
        else:
            return np.real(np.squeeze(tools.trace(self.hamiltonian @ state)))
    <mask token>
    <mask token>


class HamiltonianMIS(object):

    def __init__(self, graph: Graph, energies=(1, 1), code=qubit,
        IS_subspace=False):
        """
        Generate a vector corresponding to the diagonal of the MIS Hamiltonian.
        """
        if energies == (1, 1) and IS_subspace:
            energies = 1,
        self.code = code
        self.graph = graph
        self.n = self.graph.n
        self.energies = energies
        self.IS_subspace = IS_subspace
        self.optimization = 'max'
        if not self.IS_subspace:
            if tools.is_diagonal(self.code.Q):
                self._hamiltonian_edge_terms = np.zeros([1, (self.code.d **
                    self.code.n) ** self.n])
                self._hamiltonian_node_terms = np.zeros([1, (self.code.d **
                    self.code.n) ** self.n])
                self._is_diagonal = True
                Q = np.expand_dims(np.diagonal(self.code.Q), axis=0)

                def my_eye(n):
                    return np.ones(np.asarray(self.code.d ** self.code.n) ** n)
            else:
                self._hamiltonian_edge_terms = np.zeros([(self.code.d **
                    self.code.n) ** self.n, (self.code.d ** self.code.n) **
                    self.n])
                self._hamiltonian_node_terms = np.zeros([(self.code.d **
                    self.code.n) ** self.n, (self.code.d ** self.code.n) **
                    self.n])
                Q = np.expand_dims(np.diagonal(qubit.Q), axis=0)

                def my_eye(n):
                    return np.ones(np.asarray(qubit.d ** qubit.n) ** n)
                self._optimum_edge_terms = np.zeros([(qubit.d ** qubit.n) **
                    self.n, (qubit.d ** qubit.n) ** self.n])
                self._optimum_node_terms = np.zeros([(qubit.d ** qubit.n) **
                    self.n, (qubit.d ** qubit.n) ** self.n])
                for i, j in graph.graph.edges:
                    if j < i:
                        i, j = j, i
                    self._optimum_edge_terms = (self._optimum_edge_terms + 
                        graph.graph.edges[i, j]['weight'] * tools.
                        tensor_product([my_eye(i), Q, my_eye(j - i - 1), Q,
                        my_eye(self.n - j - 1)]))
                for i in graph.graph.nodes:
                    self._optimum_node_terms = (self._optimum_node_terms + 
                        graph.graph.nodes[i]['weight'] * tools.
                        tensor_product([my_eye(i), Q, my_eye(self.n - i - 1)]))
                self._is_diagonal = False
                Q = self.code.Q

                def my_eye(n):
                    return np.identity(np.asarray(self.code.d ** self.code.
                        n) ** n)
            for i, j in graph.graph.edges:
                if j < i:
                    i, j = j, i
                self._hamiltonian_edge_terms = (self.
                    _hamiltonian_edge_terms + graph.graph.edges[i, j][
                    'weight'] * tools.tensor_product([my_eye(i), Q, my_eye(
                    j - i - 1), Q, my_eye(self.n - j - 1)]))
            for i in graph.graph.nodes:
                self._hamiltonian_node_terms = (self.
                    _hamiltonian_node_terms + graph.graph.nodes[i]['weight'
                    ] * tools.tensor_product([my_eye(i), Q, my_eye(self.n -
                    i - 1)]))
            self._hamiltonian_node_terms = self._hamiltonian_node_terms.T
            self._hamiltonian_edge_terms = self._hamiltonian_edge_terms.T
            if self._is_diagonal:
                self._optimum_edge_terms = self._hamiltonian_edge_terms
                self._optimum_node_terms = self._hamiltonian_node_terms
                self._diagonal_hamiltonian_edge_terms = (self.
                    _hamiltonian_edge_terms.copy())
                self._diagonal_hamiltonian_node_terms = (self.
                    _hamiltonian_node_terms.copy())
                self._hamiltonian_node_terms = sparse.csr_matrix((self.
                    _hamiltonian_node_terms.flatten(), (np.arange(self.code
                    .d ** (self.code.n * self.n)), np.arange(self.code.d **
                    (self.code.n * self.n)))), shape=(self.code.d ** (self.
                    code.n * self.n), self.code.d ** (self.code.n * self.n)))
                self._hamiltonian_edge_terms = sparse.csr_matrix((self.
                    _hamiltonian_edge_terms.flatten(), (np.arange(self.code
                    .d ** (self.code.n * self.n)), np.arange(self.code.d **
                    (self.code.n * self.n)))), shape=(self.code.d ** (self.
                    code.n * self.n), self.code.d ** (self.code.n * self.n)))
            else:
                self._diagonal_hamiltonian_edge_terms = (self.
                    _hamiltonian_edge_terms.copy())
                self._diagonal_hamiltonian_node_terms = (self.
                    _hamiltonian_node_terms.copy())
                self._hamiltonian_edge_terms = sparse.csr_matrix(self.
                    _hamiltonian_edge_terms)
                self._hamiltonian_node_terms = sparse.csr_matrix(self.
                    _hamiltonian_node_terms)
            self._left_acting_hamiltonian_edge_terms = None
            self._right_acting_hamiltonian_edge_terms = None
        else:
            self._is_diagonal = True
            if not (self.code == qubit or self.code == rydberg):
                raise NotImplementedError(
                    'IS subspace only implemented for qubit and Rydberg codes.'
                    )
            if self.code == qubit:
                node_weights = np.asarray([self.graph.graph.nodes[i][
                    'weight'] for i in range(self.graph.n)])
                independent_sets = enumerate_independent_sets(self.graph.graph)
                k = self.graph.num_independent_sets - 2
                self.mis_size = 0
                C = np.zeros(self.graph.num_independent_sets, dtype=float)
                C[-1] = 0
                for i in independent_sets:
                    C[k] = np.sum([node_weights[j] for j in i])
                    k -= 1
                self._hamiltonian = sparse.csr_matrix((C, (np.arange(self.
                    graph.num_independent_sets), np.arange(self.graph.
                    num_independent_sets))))
                C = np.expand_dims(C, axis=0).T
            elif self.code == rydberg:
                independent_sets, num_IS = (self.graph.
                    generate_independent_sets_qudit(self.code))
                node_weights = np.asarray([self.graph.graph.nodes[i][
                    'weight'] for i in range(self.graph.n)])
                C = np.zeros((num_IS, 1), dtype=np.complex128)
                for k in independent_sets:
                    C[k, 0] = np.sum((independent_sets[k][2] == 0) *
                        node_weights)
            self._diagonal_hamiltonian_node_terms = C
            C = C.flatten()
            self._hamiltonian_node_terms = sparse.csr_matrix((C, (np.arange
                (len(C)), np.arange(len(C)))), shape=(len(C), len(C)))
        self._left_acting_hamiltonian_node_terms = None
        self._right_acting_hamiltonian_node_terms = None

    @property
    def hamiltonian(self):
        if not self.IS_subspace:
            return self.energies[0
                ] * self._hamiltonian_node_terms - self.energies[1
                ] * self._hamiltonian_edge_terms
        else:
            return self.energies[0] * self._hamiltonian_node_terms

    @property
    def evolution_operator(self, vector_space='hilbert'):
        if vector_space != 'hilbert' and vector_space != 'liouville':
            raise Exception(
                'Attribute vector_space must be hilbert or liouville')
        if vector_space == 'liouville':
            if (self._left_acting_hamiltonian_node_terms is None and not
                self.IS_subspace):
                self._left_acting_hamiltonian_edge_terms = sparse.kron(sparse
                    .identity(self._hamiltonian_node_terms.shape[0]), self.
                    _hamiltonian_edge_terms)
                self._right_acting_hamiltonian_edge_terms = sparse.kron(self
                    ._hamiltonian_edge_terms.T, sparse.identity(self.
                    _hamiltonian_edge_terms.shape[0]))
                self._left_acting_hamiltonian_node_terms = sparse.kron(sparse
                    .identity(self._hamiltonian_node_terms.shape[0]), self.
                    _hamiltonian_node_terms)
                self._right_acting_hamiltonian_node_terms = sparse.kron(self
                    ._hamiltonian_node_terms.T, sparse.identity(self.
                    _hamiltonian_node_terms.shape[0]))
            elif self._left_acting_hamiltonian_node_terms is None and self.IS_subspace:
                self._left_acting_hamiltonian_node_terms = sparse.kron(sparse
                    .identity(self._hamiltonian_node_terms.shape[0]), self.
                    _hamiltonian_node_terms)
                self._right_acting_hamiltonian_node_terms = sparse.kron(self
                    ._hamiltonian_node_terms.T, sparse.identity(self.
                    _hamiltonian_node_terms.shape[0]))
            if not self.IS_subspace:
                return -1.0j * (self.energies[0] * self.
                    _left_acting_hamiltonian_node_terms - self.energies[1] *
                    self._left_acting_hamiltonian_edge_terms) + 1.0j * (
                    self.energies[0] * self.
                    _right_acting_hamiltonian_node_terms - self.energies[1] *
                    self._right_acting_hamiltonian_edge_terms)
            elif self.IS_subspace:
                return -1.0j * (self.energies[0] * self.
                    _left_acting_hamiltonian_node_terms) + 1.0j * (self.
                    energies[0] * self._right_acting_hamiltonian_node_terms)
        else:
            return -1.0j * self.hamiltonian

    @property
    def _diagonal_hamiltonian(self):
        if not self.IS_subspace:
            return self.energies[0
                ] * self._diagonal_hamiltonian_node_terms - self.energies[1
                ] * self._diagonal_hamiltonian_edge_terms
        else:
            return self.energies[0] * self._diagonal_hamiltonian_node_terms

    @property
    def optimum(self):
        if self._is_diagonal:
            return np.max(self._diagonal_hamiltonian).real
        else:
            raise NotImplementedError(
                'Optimum unknown for non-diagonal Hamiltonians')

    def evolve(self, state: State, time):
        if state.is_ket:
            if self._is_diagonal:
                return State(np.exp(-1.0j * time * self.
                    _diagonal_hamiltonian) * state, is_ket=state.is_ket,
                    IS_subspace=state.IS_subspace, code=state.code, graph=
                    self.graph)
            else:
                return State(expm_multiply(-1.0j * time * self.hamiltonian,
                    state), is_ket=state.is_ket, IS_subspace=state.
                    IS_subspace, code=state.code, graph=self.graph)
        elif self._is_diagonal:
            return State(np.exp(-1.0j * time * self._diagonal_hamiltonian) *
                state * np.exp(1.0j * time * self._diagonal_hamiltonian).T,
                is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=
                state.code, graph=self.graph)
        else:
            temp = expm(-1.0j * time * self.hamiltonian)
            return State(temp @ state @ temp.conj().T, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)

    def left_multiply(self, state: State):
        if self._is_diagonal:
            return State(self._diagonal_hamiltonian * state, is_ket=state.
                is_ket, IS_subspace=state.IS_subspace, code=state.code,
                graph=self.graph)
        else:
            return State(self.hamiltonian @ state, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            if self._is_diagonal:
                return State(state.conj().T * self._diagonal_hamiltonian.T,
                    is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                    code=state.code, graph=self.graph)
            else:
                return State((state @ self.hamiltonian.T).conj().T, is_ket=
                    state.is_ket, IS_subspace=state.IS_subspace, code=state
                    .code, graph=self.graph)
        elif self._is_diagonal:
            return State(state * self._diagonal_hamiltonian.T, is_ket=state
                .is_ket, IS_subspace=state.IS_subspace, code=state.code,
                graph=self.graph)
        else:
            return State(state @ self.hamiltonian.T, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)

    def cost_function(self, state: State):
        if state.is_ket:
            if self._is_diagonal:
                return np.real(np.vdot(state, self._diagonal_hamiltonian *
                    state))
            else:
                return np.real(np.vdot(state, self.hamiltonian @ state))
        elif self._is_diagonal:
            return np.real(np.squeeze(tools.trace(self.
                _diagonal_hamiltonian * state)))
        else:
            return np.real(np.squeeze(tools.trace(self.hamiltonian @ state)))

    def optimum_overlap(self, state: State):
        if self._is_diagonal:
            optimum_indices = np.argwhere(self._diagonal_hamiltonian ==
                self.optimum).T[0]
            optimum = np.zeros(self._diagonal_hamiltonian.shape)
            optimum[optimum_indices] = 1
        else:
            raise NotImplementedError(
                'Optimum overlap not implemented for non-diagonal Hamiltonians'
                )
        if state.is_ket:
            if self._is_diagonal:
                return np.real(np.vdot(state, optimum * state))
            else:
                return np.real(np.vdot(state, self.hamiltonian @ state))
        elif self._is_diagonal:
            return np.real(np.squeeze(tools.trace(optimum * state)))
        else:
            return np.real(np.squeeze(tools.trace(self.hamiltonian @ state)))

    def approximation_ratio(self, state: State):
        return self.cost_function(state) / self.optimum


class HamiltonianGlobalPauli(object):

    def __init__(self, pauli: str='X', code=qubit):
        self.code = code
        self.pauli = pauli
        if self.pauli == 'X':
            self._operator = self.code.X
        elif self.pauli == 'Y':
            self._operator = self.code.Y
        elif self.pauli == 'Z':
            self._operator = self.code.Z
        self.hamiltonian = None

    def evolve(self, state: State, alpha):
        if self.hamiltonian is None:
            """Initialize the Hamiltonian only once, as it is costly."""
            self.hamiltonian = tools.tensor_product([self._operator] *
                state.number_logical_qudits)
        return self.code.multiply(np.cos(alpha) * np.identity(state.
            dimension) - 1.0j * np.sin(alpha) * self.hamiltonian)

    def left_multiply(self, state: State):
        all_qubits = list(range(state.number_logical_qudits))
        return self.code.left_multiply(state, all_qubits, [self.pauli] *
            state.number_logical_qudits)

    def right_multiply(self, state: State):
        all_qubits = list(range(state.number_logical_qudits))
        return self.code.right_multiply(state, all_qubits, [self.pauli] *
            state.number_logical_qudits)


class HamiltonianBookatzPenalty(object):

    def __init__(self, code=qubit, energies=(1,)):
        self.code = code
        self.projector = np.identity(self.code.d ** self.code.n
            ) - self.code.code_space_projector
        self.energies = energies

    def evolve(self, state: State, time):
        for i in range(state.number_logical_qudits):
            state = self.code.rotation(state, [i], self.energies[0] * time,
                self.projector, is_idempotent=True)
        return state

    def left_multiply(self, state: State):
        out = np.zeros_like(state, dtype=np.complex128)
        for i in range(state.number_logical_qudits):
            out = out + self.code.left_multiply(state, [i], self.projector)
        return State(self.energies[0] * out, is_ket=state.is_ket,
            IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        out = np.zeros_like(state, dtype=np.complex128)
        for i in range(state.number_logical_qudits):
            out = out + self.code.right_multiply(state, [i], self.projector)
        return State(self.energies[0] * out, is_ket=state.is_ket,
            IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)


class HamiltonianMarvianPenalty(object):

    def __init__(self, Nx, Ny):
        super().__init__()
        self.Nx = Nx
        self.Ny = Ny
        self.n = 3 * Nx * Ny
        hp = np.zeros([2 ** self.n, 2 ** self.n])
        for i in range(int(self.Nx * self.Ny)):
            hp = hp + tools.tensor_product([tools.identity(i * 3), tools.Z(
                ), tools.Z(), tools.identity(self.n - i * 3 - 2)]
                ) + tools.tensor_product([tools.identity(i * 3 + 1), tools.
                X(), tools.X(), tools.identity(self.n - i * 3 - 3)])
        for j in range(self.Ny):
            for k in range(self.Nx):
                if k != self.Nx - 1:
                    hp = hp + tools.tensor_product([tools.identity(j * self
                        .Nx * 3 + k * 3), tools.X(), tools.identity(2),
                        tools.X(), tools.identity(self.n - (j * self.Nx * 3 +
                        k * 3) - 4)]) + tools.tensor_product([tools.
                        identity(j * self.Nx * 3 + k * 3 + 2), tools.Z(),
                        tools.identity(2), tools.Z(), tools.identity(self.n -
                        (j * self.Nx * 3 + k * 3 + 2) - 4)])
                if j != self.Ny - 1:
                    hp = hp + tools.tensor_product([tools.identity(j * self
                        .Nx * 3 + k * 3), tools.X(), tools.identity(3 *
                        self.Nx - 1), tools.X(), tools.identity(self.n - (j *
                        self.Nx * 3 + k * 3) - 3 * self.Nx - 1)]
                        ) + tools.tensor_product([tools.identity(j * self.
                        Nx * 3 + k * 3 + 2), tools.Z(), tools.identity(3 *
                        self.Nx - 1), tools.Z(), tools.identity(self.n - (j *
                        self.Nx * 3 + k * 3 + 2) - 3 * self.Nx - 1)])
        self.hamiltonian = -1 * hp


class HamiltonianHeisenberg(object):

    def __init__(self, graph: Graph, energies=(1, 1), subspace='all', code=
        qubit, IS_subspace=False):
        self.code = code
        self.graph = graph
        self.n = self.graph.n
        self.energies = energies
        self._is_diagonal = True
        assert code is qubit
        if IS_subspace:
            raise NotImplementedError
        self.IS_subspace = IS_subspace
        self._hamiltonian = None
        self.subspace = subspace
        if self.subspace == 'all':
            """Initialize the Hamiltonian."""
            self._hamiltonian_zz = None
            self._hamiltonian_xy = None
        else:
            from scipy.special import comb
            dim = comb(graph.n, (graph.n + self.subspace) / 2)
            if dim % 1 != 0:
                raise Exception('Invalid subspace for number of spins')
            dim = int(dim)
            hamiltonian_zz = np.zeros([dim, 1])
            rows = np.zeros(graph.n * dim, dtype=int)
            columns = np.zeros(graph.n * dim, dtype=int)
            entries = np.zeros(graph.n * dim, dtype=int)
            num_terms = 0
            states = np.zeros((dim, graph.n))
            for i in range(self.code.d ** (self.code.n * self.n)):
                nary = tools.int_to_nary(i, size=graph.n)
                if graph.n - np.sum(nary) == (self.subspace + graph.n) / 2:
                    states[num_terms, ...] = nary
                    num_terms += 1
            num_terms = 0
            for i in range(dim):
                zz = 0
                for a, b in self.graph.edges:
                    if b < a:
                        a, b = b, a
                    if np.abs(states[i, a] - states[i, b]) == 1:
                        zz -= 1
                    else:
                        zz += 1
                    if states[i, a] == 0 and states[i, b] == 1:
                        temp = states[i, ...].copy()
                        temp[a] = 1
                        temp[b] = 0
                        where_matched_temp = np.argwhere(np.sum(np.abs(
                            states - temp), axis=1) == 0).flatten()
                        entries[num_terms] = 1
                        rows[num_terms] = where_matched_temp
                        columns[num_terms] = i
                        entries[num_terms + 1] = 1
                        rows[num_terms + 1] = i
                        columns[num_terms + 1] = where_matched_temp
                        num_terms += 2
                hamiltonian_zz[i, 0] = zz
            self._hamiltonian_zz = sparse.csr_matrix((hamiltonian_zz.
                flatten(), (np.arange(dim), np.arange(dim))), shape=(dim, dim))
            self._hamiltonian_xy = sparse.csr_matrix((entries, (rows,
                columns)), shape=(dim, dim))
            self.states = states

    @property
    def hamiltonian(self):
        if (self._hamiltonian_zz is None or self._hamiltonian_xy is None and
            self.subspace == 'all'):
            z = np.expand_dims(np.diagonal(self.code.Z), axis=0).T

            def my_eye(n):
                return np.ones((np.asarray(self.code.d ** self.code.n) ** n, 1)
                    )
            hamiltonian_zz = np.zeros([self.code.d ** (self.code.n * self.n
                ), 1])
            for a, b in self.graph.edges:
                if b < a:
                    a, b = b, a
                hamiltonian_zz = hamiltonian_zz + self.graph.graph[a][b][
                    'weight'] * tools.tensor_product([my_eye(a), z, my_eye(
                    b - a - 1), z, my_eye(self.n - b - 1)], sparse=not self
                    ._is_diagonal)
            self._hamiltonian_zz = sparse.csr_matrix((hamiltonian_zz.
                flatten(), (np.arange(self.code.d ** (self.code.n * self.n)
                ), np.arange(self.code.d ** (self.code.n * self.n)))),
                shape=(self.code.d ** (self.code.n * self.n), self.code.d **
                (self.code.n * self.n)))
            rows = np.zeros(self.graph.n * self.code.d ** (self.code.n *
                self.n), dtype=int)
            columns = np.zeros(self.graph.n * self.code.d ** (self.code.n *
                self.n), dtype=int)
            entries = np.zeros(self.graph.n * self.code.d ** (self.code.n *
                self.n), dtype=int)
            num_terms = 0
            for i in range(self.code.d ** (self.code.n * self.n)):
                nary = tools.int_to_nary(i, size=self.graph.n)
                for a, b in self.graph.edges:
                    if b < a:
                        a, b = b, a
                    if nary[a] == 0 and nary[b] == 1:
                        temp = nary.copy()
                        temp[a] = 1
                        temp[b] = 0
                        temp = tools.nary_to_int(temp)
                        entries[num_terms] = 1
                        rows[num_terms] = temp
                        columns[num_terms] = i
                        entries[num_terms + 1] = 1
                        rows[num_terms + 1] = i
                        columns[num_terms + 1] = temp
                        num_terms += 2
            self._hamiltonian_xy = sparse.csr_matrix((entries, (rows,
                columns)), shape=(self.code.d ** (self.code.n * self.n), 
                self.code.d ** (self.code.n * self.n)))
        return self.energies[0] * self._hamiltonian_zz + self.energies[1
            ] * self._hamiltonian_xy

    def left_multiply(self, state: State):
        temp = np.zeros(state.shape, dtype=np.complex128)
        for edge in self.graph.edges:
            if self.energies[0] != 0:
                term = self.code.left_multiply(state, [edge[0], edge[1]], [
                    'X', 'X'])
                temp = temp + self.energies[0] * term
                term = self.code.left_multiply(state, [edge[0], edge[1]], [
                    'Y', 'Y'])
                temp = temp + self.energies[0] * term
            if self.energies[1] != 0:
                term = self.code.left_multiply(state, [edge[0], edge[1]], [
                    'Z', 'Z'])
                temp = temp + self.energies[1] * term
        return State(temp, is_ket=state.is_ket, IS_subspace=state.
            IS_subspace, code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        temp = np.zeros(state.shape, dtype=np.complex128)
        for edge in self.graph.edges:
            if self.energies[0] != 0:
                term = self.code.right_multiply(state, [edge[0], edge[1]],
                    ['X', 'X'])
                temp = temp + self.energies[0] * term
                term = self.code.right_multiply(state, [edge[0], edge[1]],
                    ['Y', 'Y'])
                temp = temp + self.energies[0] * term
            if self.energies[1] != 0:
                term = self.code.right_multiply(state, [edge[0], edge[1]],
                    ['Z', 'Z'])
                temp = temp + self.energies[1] * term
        return State(temp, is_ket=state.is_ket, IS_subspace=state.
            IS_subspace, code=state.code, graph=self.graph)

    def evolve(self, state: State, time):
        if not state.is_ket:
            exp_hamiltonian = expm(-1.0j * time * self.hamiltonian)
            return State(exp_hamiltonian @ state @ exp_hamiltonian.conj().T,
                is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=
                state.code, graph=self.graph)
        if state.is_ket:
            return State(expm_multiply(-1.0j * time * self.hamiltonian,
                state), is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                code=state.code, graph=self.graph)

    def cost_function(self, state: State):
        if state.is_ket:
            return np.real(np.vdot(state, self.hamiltonian * state))
        else:
            return np.real(np.squeeze(tools.trace(self.hamiltonian * state)))


class HamiltonianEnergyShift(object):

    def __init__(self, index: int=0, energies=(1,), code=qubit, IS_subspace
        =False, graph=None):
        """Default is that the first element in transition is the higher energy s."""
        self.index = index
        self.graph = graph
        self.energies = energies
        self.code = code
        if not self.code.logical_code:
            if not 0 <= self.index < self.code.d:
                raise Exception('Index exceeds qudit dimension.')
            self._operator = np.zeros((self.code.d, self.code.d))
            self._operator[self.index, self.index] = 1
        else:
            if self.index != 0 and self.index != 1:
                raise Exception(
                    'Logical codes are qubits, so index must be 0 or 1.')
            if self.index == 0:
                self._operator = self.code.Q
            elif self.index == 1:
                self._operator = self.code.P
        self.IS_subspace = IS_subspace
        if self.IS_subspace:
            assert graph is not None
            assert isinstance(graph, Graph)
            if code is not qubit:
                IS, num_IS = graph.independent_sets_qudit(self.code)
                self._diagonal_hamiltonian = np.zeros((num_IS, 1), dtype=float)
                for k in range(num_IS):
                    self._diagonal_hamiltonian[k, 0] = np.sum(IS[k, ...] ==
                        self.index)
                self._hamiltonian = sparse.csr_matrix((self.
                    _diagonal_hamiltonian.T[0], (np.arange(num_IS), np.
                    arange(num_IS))), shape=(num_IS, num_IS))
            else:
                independent_sets = enumerate_independent_sets(graph.graph)
                k = self.graph.num_independent_sets - 2
                self.mis_size = 0
                hamiltonian = np.zeros(self.graph.num_independent_sets,
                    dtype=float)
                hamiltonian[-1] = 0
                for i in independent_sets:
                    hamiltonian[k] = len(i)
                    k -= 1
                self._hamiltonian = sparse.csr_matrix((hamiltonian, (np.
                    arange(self.graph.num_independent_sets), np.arange(self
                    .graph.num_independent_sets))), shape=(self.graph.
                    num_independent_sets, self.graph.num_independent_sets))
        else:
            self._hamiltonian = None

    @property
    def hamiltonian(self):
        if self._hamiltonian is None:
            assert not self.IS_subspace
            try:
                assert self.graph is not None
            except AssertionError:
                print(
                    'self.graph must be not None to generate the Hamiltonian property.'
                    )
            self._hamiltonian = sparse.csr_matrix(((self.code.d * self.code
                .n) ** self.graph.n, (self.code.d * self.code.n) ** self.
                graph.n))
            for i in range(self.graph.n):
                self._hamiltonian = self._hamiltonian + tools.tensor_product([
                    sparse.identity((self.code.d * self.code.n) ** i), self
                    ._operator, sparse.identity((self.code.d * self.code.n) **
                    (self.graph.n - i - 1))], sparse=True)
        return self.energies[0] * self._hamiltonian

    def left_multiply(self, state: State):
        if not self.IS_subspace:
            temp = np.zeros_like(state)
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.left_multiply(state, [i], self.
                        _operator)
                elif not self.code.logical_code:
                    ind = self.code.d ** i
                    out = np.zeros_like(state, dtype=np.complex128)
                    if state.is_ket:
                        state = state.reshape((-1, self.code.d, ind), order='F'
                            )
                        out = out.reshape((-1, self.code.d, ind), order='F')
                        out[:, self.index, :] = state[:, self.index, :]
                        state = state.reshape(state_shape, order='F')
                        out = out.reshape(state_shape, order='F')
                    else:
                        out = out.reshape((-1, self.code.d, self.code.d **
                            (state.number_physical_qudits - 1), self.code.d,
                            ind), order='F')
                        state = state.reshape((-1, self.code.d, self.code.d **
                            (state.number_physical_qudits - 1), self.code.d,
                            ind), order='F')
                        out[:, self.index, :, :, :] = state[:, self.index,
                            :, :, :]
                        state = state.reshape(state_shape, order='F')
                        out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)
        else:
            return State(self.energies[0] * self._diagonal_hamiltonian *
                state, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            print(
                'Warning: right multiply functionality currently applies the operator and daggers the state.'
                )
            return self.left_multiply(state).conj().T
        elif not self.IS_subspace:
            temp = np.zeros_like(state)
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.right_multiply(state, [i], self
                        ._operator)
                else:
                    ind = self.code.d ** i
                    out = np.zeros_like(state)
                    out = out.reshape((-1, self.code.d, self.code.d ** (
                        state.number_physical_qudits - 1), self.code.d, ind
                        ), order='F')
                    state = state.reshape((-1, self.code.d, self.code.d **
                        (state.number_physical_qudits - 1), self.code.d,
                        ind), order='F')
                    out[:, :, :, self.index, :] = state[:, :, :, self.index, :]
                    state = state.reshape(state_shape, order='F')
                    out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)
        else:
            return State(self.energies[0] * state * self.
                _diagonal_hamiltonian.T, is_ket=state.is_ket, IS_subspace=
                state.IS_subspace, code=state.code, graph=self.graph)

    def evolve(self, state: State, time):
        """
        Use reshape to efficiently implement evolution under :math:`H_B=\\\\sum_i X_i`
        """
        if not self.IS_subspace:
            out = state.copy()
            for i in range(state.number_logical_qudits):
                out = self.code.rotation(out, [i], self.energies[0] * time,
                    self._operator)
            return out
        elif state.is_ket:
            return State(np.exp(-1.0j * time * self.energies[0] * self.
                _diagonal_hamiltonian) * state, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)
        else:
            exp_hamiltonian = np.exp(-1.0j * time * self.energies[0] * self
                ._diagonal_hamiltonian)
            return State(exp_hamiltonian * state * exp_hamiltonian.conj().T,
                is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=
                state.code, graph=self.graph)


class HamiltonianRydberg(object):

    def __init__(self, tails_graph: Graph, hard_constraint_graph=None,
        index: int=0, energies=(1,), code=qubit, IS_subspace=False):
        """Default is that the first element in transition is the higher energy s."""
        self.index = index
        self.energies = energies
        self.code = code
        self.tails_graph = tails_graph
        self.hard_constraint_graph = hard_constraint_graph
        if self.code.logical_code:
            raise Exception('Logical codes not supported.')
        self.IS_subspace = IS_subspace
        try:
            if (self.tails_graph.periodic or self.hard_constraint_graph.
                periodic):
                raise Exception('Periodic graphs not yet supported.')
        except:
            pass
        if self.IS_subspace:
            assert hard_constraint_graph is not None
            assert isinstance(hard_constraint_graph, Graph)
            if code is not qubit:
                IS, num_IS = (hard_constraint_graph.
                    generate_independent_sets_qudit(self.code))
                self._diagonal_hamiltonian = np.zeros((num_IS, 1), dtype=float)
                for k in range(num_IS):
                    weight = 0
                    where_rydberg = np.argwhere(self.hard_constraint_graph.
                        independent_sets[k] == 0)
                    if len(where_rydberg) > 0:
                        where_rydberg = where_rydberg.T[0]
                        for i, node1 in enumerate(where_rydberg):
                            for j, node2 in enumerate(where_rydberg[i + 1:]):
                                weight += self.tails_graph.graph[node1][node2][
                                    'weight']
                    self._diagonal_hamiltonian[k, 0] = weight
                self._hamiltonian = sparse.csr_matrix((self.
                    _diagonal_hamiltonian.T[0], (np.arange(num_IS), np.
                    arange(num_IS))), shape=(num_IS, num_IS))
            else:
                independent_sets = enumerate_independent_sets(self.
                    hard_constraint_graph.graph)
                k = self.hard_constraint_graph.num_independent_sets - 2
                self.mis_size = 0
                self._diagonal_hamiltonian = np.zeros(self.
                    hard_constraint_graph.num_independent_sets, dtype=float)
                self._diagonal_hamiltonian[-1] = 0
                for i in independent_sets:
                    weight = 0
                    where_rydberg = np.array(i)
                    if len(where_rydberg) > 0:
                        for i, node1 in enumerate(where_rydberg):
                            for j, node2 in enumerate(where_rydberg[i + 1:]):
                                weight += self.tails_graph.graph[node1][node2][
                                    'weight']
                    self._diagonal_hamiltonian[k] = weight
                    k -= 1
                self._hamiltonian = sparse.csr_matrix((self.
                    _diagonal_hamiltonian, (np.arange(self.
                    hard_constraint_graph.num_independent_sets), np.arange(
                    self.hard_constraint_graph.num_independent_sets))),
                    shape=(self.hard_constraint_graph.num_independent_sets,
                    self.hard_constraint_graph.num_independent_sets))
                self._diagonal_hamiltonian = np.expand_dims(self.
                    _diagonal_hamiltonian, axis=1)
        else:
            self._hamiltonian = None

    @property
    def hamiltonian(self):
        if self._hamiltonian is None:
            raise Exception
        return self.energies[0] * self._hamiltonian

    def left_multiply(self, state: State):
        return State(self.energies[0] * self._diagonal_hamiltonian * state,
            is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.
            code, graph=self.hard_constraint_graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            print(
                'Warning: right multiply functionality currently applies the operator and daggers the state.'
                )
            return self.left_multiply(state).conj().T
        elif not self.IS_subspace:
            temp = np.zeros_like(state)
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.right_multiply(state, [i], self
                        ._operator)
                else:
                    ind = self.code.d ** i
                    out = np.zeros_like(state)
                    out = out.reshape((-1, self.code.d, self.code.d ** (
                        state.number_physical_qudits - 1), self.code.d, ind
                        ), order='F')
                    state = state.reshape((-1, self.code.d, self.code.d **
                        (state.number_physical_qudits - 1), self.code.d,
                        ind), order='F')
                    out[:, :, :, self.index, :] = state[:, :, :, self.index, :]
                    state = state.reshape(state_shape, order='F')
                    out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                hard_constraint_graph)
        else:
            return State(self.energies[0] * state * self.
                _diagonal_hamiltonian.T, is_ket=state.is_ket, IS_subspace=
                state.IS_subspace, code=state.code, graph=self.
                hard_constraint_graph)

    def evolve(self, state: State, time):
        """
        Use reshape to efficiently implement evolution under :math:`H_B=\\\\sum_i X_i`
        """
        if not self.IS_subspace:
            out = state.copy()
            for i in range(state.number_logical_qudits):
                out = self.code.rotation(out, [i], self.energies[0] * time,
                    self._operator)
            return out
        elif state.is_ket:
            return State(np.exp(-1.0j * time * self.energies[0] * self.
                _diagonal_hamiltonian) * state, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                hard_constraint_graph)
        else:
            exp_hamiltonian = np.exp(-1.0j * time * self.energies[0] * self
                ._diagonal_hamiltonian)
            return State(exp_hamiltonian * state * exp_hamiltonian.conj().T,
                is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=
                state.code, graph=self.hard_constraint_graph)

<mask token>


class HamiltonianMaxCut(object):

    def __init__(self, graph: Graph, code=qubit, energies=(1,),
        cost_function=True, use_Z2_symmetry=False):
        """
        Generate a vector corresponding to the diagonal of the MaxCut Hamiltonian.
        """
        self.code = code
        self.energies = energies
        self.graph = graph
        self.optimization = 'max'
        self.n = self.graph.n
        if use_Z2_symmetry:
            c = np.zeros([self.code.d ** (self.code.n * (self.n - 1)), 1])
        else:
            c = np.zeros([self.code.d ** (self.code.n * self.n), 1])
        if tools.is_diagonal(self.code.Z):
            self._is_diagonal = True
            z = np.expand_dims(np.diagonal(self.code.Z), axis=0).T

            def my_eye(n):
                return np.ones((np.asarray(self.code.d ** self.code.n) ** n, 1)
                    )
        else:
            self._is_diagonal = False
            z = np.expand_dims(np.diagonal(qubit.Z), axis=0).T

            def my_eye(n):
                return np.ones((np.asarray(qubit.d) ** n, 1))
            for a, b in self.graph.edges:
                if b < a:
                    a, b = b, a
                if use_Z2_symmetry:
                    if cost_function:
                        if a == min(self.graph.nodes):
                            c = c - 1 / 2 * graph.graph[a][b]['weight'] * (
                                tools.tensor_product([my_eye(b - 1), z,
                                my_eye(self.n - b - 1)]) - my_eye(self.n - 1))
                        else:
                            c = c - 1 / 2 * graph.graph[a][b]['weight'] * (
                                tools.tensor_product([my_eye(a - 1), z,
                                my_eye(b - a - 1), z, my_eye(self.n - b - 1
                                )]) - my_eye(self.n - 1))
                    elif a == min(self.graph.nodes):
                        c = c + graph.graph[a][b]['weight'
                            ] * tools.tensor_product([my_eye(b - 1), z,
                            my_eye(self.n - b - 1)])
                    else:
                        c = c + graph.graph[a][b]['weight'
                            ] * tools.tensor_product([my_eye(a - 1), z,
                            my_eye(b - a - 1), z, my_eye(self.n - b - 1)])
                elif cost_function:
                    c = c - 1 / 2 * graph.graph[a][b]['weight'] * (tools.
                        tensor_product([my_eye(a), z, my_eye(b - a - 1), z,
                        my_eye(self.n - b - 1)]) - my_eye(self.n))
                else:
                    c = c + graph.graph[a][b]['weight'] * tools.tensor_product(
                        [my_eye(a), z, my_eye(b - a - 1), z, my_eye(self.n -
                        b - 1)])
            self._optimum = np.max(c).real
            c = sparse.csr_matrix((self.code.d ** (self.code.n * self.n), 
                self.code.d ** (self.code.n * self.n)))
            z = sparse.csr_matrix(self.code.Z)

            def my_eye(n):
                return sparse.csr_matrix(np.ones(np.asarray(z.shape[0]) **
                    n), (np.asarray(z.shape[0]) ** n, np.asarray(z.shape[0]
                    ) ** n))
        for a, b in self.graph.edges:
            if b < a:
                a, b = b, a
            if cost_function:
                if use_Z2_symmetry:
                    if a == min(self.graph.nodes):
                        c = c - 1 / 2 * graph.graph[a][b]['weight'] * (tools
                            .tensor_product([my_eye(b - 1), z, my_eye(self.
                            n - b - 1)]) - my_eye(self.n - 1))
                    else:
                        c = c - 1 / 2 * graph.graph[a][b]['weight'] * (tools
                            .tensor_product([my_eye(a - 1), z, my_eye(b - a -
                            1), z, my_eye(self.n - b - 1)]) - my_eye(self.n -
                            1))
                else:
                    c = c - 1 / 2 * graph.graph[a][b]['weight'] * (tools.
                        tensor_product([my_eye(a), z, my_eye(b - a - 1), z,
                        my_eye(self.n - b - 1)], sparse=not self.
                        _is_diagonal) - my_eye(self.n))
            elif use_Z2_symmetry:
                if a == min(self.graph.nodes):
                    c = c + graph.graph[a][b]['weight'] * tools.tensor_product(
                        [my_eye(b - 1), z, my_eye(self.n - b - 1)])
                else:
                    c = c + graph.graph[a][b]['weight'] * tools.tensor_product(
                        [my_eye(a - 1), z, my_eye(b - a - 1), z, my_eye(
                        self.n - b - 1)])
            else:
                c = c + graph.graph[a][b]['weight'] * tools.tensor_product([
                    my_eye(a), z, my_eye(b - a - 1), z, my_eye(self.n - b -
                    1)], sparse=not self._is_diagonal)
        if self._is_diagonal:
            self._diagonal_hamiltonian = c
            self._optimum = np.max(c).real
            if use_Z2_symmetry:
                c = sparse.csr_matrix((c.flatten(), (np.arange(self.code.d **
                    (self.code.n * (self.n - 1))), np.arange(self.code.d **
                    (self.code.n * (self.n - 1))))), shape=(self.code.d **
                    (self.code.n * (self.n - 1)), self.code.d ** (self.code
                    .n * (self.n - 1))))
            else:
                c = sparse.csr_matrix((c.flatten(), (np.arange(self.code.d **
                    (self.code.n * self.n)), np.arange(self.code.d ** (self
                    .code.n * self.n)))), shape=(self.code.d ** (self.code.
                    n * self.n), self.code.d ** (self.code.n * self.n)))
        else:
            c = sparse.csr_matrix(c)
        self._hamiltonian = c
        self._left_acting_hamiltonian = None
        self._right_acting_hamiltonian = None

    @property
    def hamiltonian(self):
        return self.energies[0] * self._hamiltonian

    @property
    def evolution_operator(self, vector_space='hilbert'):
        if vector_space != 'hilbert' and vector_space != 'liouville':
            raise Exception(
                'Attribute vector_space must be hilbert or liouville')
        if vector_space == 'liouville':
            if self._left_acting_hamiltonian is None:
                self._left_acting_hamiltonian = sparse.kron(sparse.identity
                    (self._hamiltonian.shape[0]), self._hamiltonian)
                self._right_acting_hamiltonian = sparse.kron(self.
                    _hamiltonian.T, sparse.identity(self._hamiltonian.shape[0])
                    )
            return -1.0j * self.energies[0
                ] * self._left_acting_hamiltonian + 1.0j * self.energies[0
                ] * self._right_acting_hamiltonian
        else:
            return -1.0j * self.hamiltonian

    @property
    def optimum(self):
        return self.energies[0] * self._optimum

    def evolve(self, state: State, time):
        if state.is_ket:
            if self._is_diagonal:
                return State(np.exp(-1.0j * time * self.
                    _diagonal_hamiltonian) * state, is_ket=state.is_ket,
                    IS_subspace=state.IS_subspace, code=state.code, graph=
                    self.graph)
            else:
                return State(expm_multiply(-1.0j * time * self.hamiltonian,
                    state), is_ket=state.is_ket, IS_subspace=state.
                    IS_subspace, code=state.code, graph=self.graph)
        elif self._is_diagonal:
            return State(np.exp(-1.0j * time * self._diagonal_hamiltonian) *
                state * np.exp(1.0j * time * self._diagonal_hamiltonian).T,
                is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=
                state.code, graph=self.graph)
        else:
            temp = expm(-1.0j * time * self.hamiltonian)
            return State(temp @ state @ temp.conj().T, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)

    def left_multiply(self, state: State):
        if self._is_diagonal:
            return State(self._diagonal_hamiltonian * state, is_ket=state.
                is_ket, IS_subspace=state.IS_subspace, code=state.code,
                graph=self.graph)
        else:
            return State(self.hamiltonian @ state, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            if self._is_diagonal:
                return State(state.conj().T * self._diagonal_hamiltonian.T,
                    is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                    code=state.code, graph=self.graph)
            else:
                return State((state @ self.hamiltonian.T).conj().T, is_ket=
                    state.is_ket, IS_subspace=state.IS_subspace, code=state
                    .code, graph=self.graph)
        elif self._is_diagonal:
            return State(state * self._diagonal_hamiltonian.T, is_ket=state
                .is_ket, IS_subspace=state.IS_subspace, code=state.code,
                graph=self.graph)
        else:
            return State(state @ self.hamiltonian.T, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)

    def cost_function(self, state: State):
        if state.is_ket:
            if self._is_diagonal:
                return np.real(np.vdot(state, self._diagonal_hamiltonian *
                    state))
            else:
                return np.real(np.vdot(state, self.hamiltonian @ state))
        elif self._is_diagonal:
            return np.real(np.squeeze(tools.trace(self.
                _diagonal_hamiltonian * state)))
        else:
            return np.real(np.squeeze(tools.trace(self.hamiltonian @ state)))
    <mask token>

    def approximation_ratio(self, state: State):
        return self.cost_function(state) / self.optimum


class HamiltonianMIS(object):

    def __init__(self, graph: Graph, energies=(1, 1), code=qubit,
        IS_subspace=False):
        """
        Generate a vector corresponding to the diagonal of the MIS Hamiltonian.
        """
        if energies == (1, 1) and IS_subspace:
            energies = 1,
        self.code = code
        self.graph = graph
        self.n = self.graph.n
        self.energies = energies
        self.IS_subspace = IS_subspace
        self.optimization = 'max'
        if not self.IS_subspace:
            if tools.is_diagonal(self.code.Q):
                self._hamiltonian_edge_terms = np.zeros([1, (self.code.d **
                    self.code.n) ** self.n])
                self._hamiltonian_node_terms = np.zeros([1, (self.code.d **
                    self.code.n) ** self.n])
                self._is_diagonal = True
                Q = np.expand_dims(np.diagonal(self.code.Q), axis=0)

                def my_eye(n):
                    return np.ones(np.asarray(self.code.d ** self.code.n) ** n)
            else:
                self._hamiltonian_edge_terms = np.zeros([(self.code.d **
                    self.code.n) ** self.n, (self.code.d ** self.code.n) **
                    self.n])
                self._hamiltonian_node_terms = np.zeros([(self.code.d **
                    self.code.n) ** self.n, (self.code.d ** self.code.n) **
                    self.n])
                Q = np.expand_dims(np.diagonal(qubit.Q), axis=0)

                def my_eye(n):
                    return np.ones(np.asarray(qubit.d ** qubit.n) ** n)
                self._optimum_edge_terms = np.zeros([(qubit.d ** qubit.n) **
                    self.n, (qubit.d ** qubit.n) ** self.n])
                self._optimum_node_terms = np.zeros([(qubit.d ** qubit.n) **
                    self.n, (qubit.d ** qubit.n) ** self.n])
                for i, j in graph.graph.edges:
                    if j < i:
                        i, j = j, i
                    self._optimum_edge_terms = (self._optimum_edge_terms + 
                        graph.graph.edges[i, j]['weight'] * tools.
                        tensor_product([my_eye(i), Q, my_eye(j - i - 1), Q,
                        my_eye(self.n - j - 1)]))
                for i in graph.graph.nodes:
                    self._optimum_node_terms = (self._optimum_node_terms + 
                        graph.graph.nodes[i]['weight'] * tools.
                        tensor_product([my_eye(i), Q, my_eye(self.n - i - 1)]))
                self._is_diagonal = False
                Q = self.code.Q

                def my_eye(n):
                    return np.identity(np.asarray(self.code.d ** self.code.
                        n) ** n)
            for i, j in graph.graph.edges:
                if j < i:
                    i, j = j, i
                self._hamiltonian_edge_terms = (self.
                    _hamiltonian_edge_terms + graph.graph.edges[i, j][
                    'weight'] * tools.tensor_product([my_eye(i), Q, my_eye(
                    j - i - 1), Q, my_eye(self.n - j - 1)]))
            for i in graph.graph.nodes:
                self._hamiltonian_node_terms = (self.
                    _hamiltonian_node_terms + graph.graph.nodes[i]['weight'
                    ] * tools.tensor_product([my_eye(i), Q, my_eye(self.n -
                    i - 1)]))
            self._hamiltonian_node_terms = self._hamiltonian_node_terms.T
            self._hamiltonian_edge_terms = self._hamiltonian_edge_terms.T
            if self._is_diagonal:
                self._optimum_edge_terms = self._hamiltonian_edge_terms
                self._optimum_node_terms = self._hamiltonian_node_terms
                self._diagonal_hamiltonian_edge_terms = (self.
                    _hamiltonian_edge_terms.copy())
                self._diagonal_hamiltonian_node_terms = (self.
                    _hamiltonian_node_terms.copy())
                self._hamiltonian_node_terms = sparse.csr_matrix((self.
                    _hamiltonian_node_terms.flatten(), (np.arange(self.code
                    .d ** (self.code.n * self.n)), np.arange(self.code.d **
                    (self.code.n * self.n)))), shape=(self.code.d ** (self.
                    code.n * self.n), self.code.d ** (self.code.n * self.n)))
                self._hamiltonian_edge_terms = sparse.csr_matrix((self.
                    _hamiltonian_edge_terms.flatten(), (np.arange(self.code
                    .d ** (self.code.n * self.n)), np.arange(self.code.d **
                    (self.code.n * self.n)))), shape=(self.code.d ** (self.
                    code.n * self.n), self.code.d ** (self.code.n * self.n)))
            else:
                self._diagonal_hamiltonian_edge_terms = (self.
                    _hamiltonian_edge_terms.copy())
                self._diagonal_hamiltonian_node_terms = (self.
                    _hamiltonian_node_terms.copy())
                self._hamiltonian_edge_terms = sparse.csr_matrix(self.
                    _hamiltonian_edge_terms)
                self._hamiltonian_node_terms = sparse.csr_matrix(self.
                    _hamiltonian_node_terms)
            self._left_acting_hamiltonian_edge_terms = None
            self._right_acting_hamiltonian_edge_terms = None
        else:
            self._is_diagonal = True
            if not (self.code == qubit or self.code == rydberg):
                raise NotImplementedError(
                    'IS subspace only implemented for qubit and Rydberg codes.'
                    )
            if self.code == qubit:
                node_weights = np.asarray([self.graph.graph.nodes[i][
                    'weight'] for i in range(self.graph.n)])
                independent_sets = enumerate_independent_sets(self.graph.graph)
                k = self.graph.num_independent_sets - 2
                self.mis_size = 0
                C = np.zeros(self.graph.num_independent_sets, dtype=float)
                C[-1] = 0
                for i in independent_sets:
                    C[k] = np.sum([node_weights[j] for j in i])
                    k -= 1
                self._hamiltonian = sparse.csr_matrix((C, (np.arange(self.
                    graph.num_independent_sets), np.arange(self.graph.
                    num_independent_sets))))
                C = np.expand_dims(C, axis=0).T
            elif self.code == rydberg:
                independent_sets, num_IS = (self.graph.
                    generate_independent_sets_qudit(self.code))
                node_weights = np.asarray([self.graph.graph.nodes[i][
                    'weight'] for i in range(self.graph.n)])
                C = np.zeros((num_IS, 1), dtype=np.complex128)
                for k in independent_sets:
                    C[k, 0] = np.sum((independent_sets[k][2] == 0) *
                        node_weights)
            self._diagonal_hamiltonian_node_terms = C
            C = C.flatten()
            self._hamiltonian_node_terms = sparse.csr_matrix((C, (np.arange
                (len(C)), np.arange(len(C)))), shape=(len(C), len(C)))
        self._left_acting_hamiltonian_node_terms = None
        self._right_acting_hamiltonian_node_terms = None

    @property
    def hamiltonian(self):
        if not self.IS_subspace:
            return self.energies[0
                ] * self._hamiltonian_node_terms - self.energies[1
                ] * self._hamiltonian_edge_terms
        else:
            return self.energies[0] * self._hamiltonian_node_terms

    @property
    def evolution_operator(self, vector_space='hilbert'):
        if vector_space != 'hilbert' and vector_space != 'liouville':
            raise Exception(
                'Attribute vector_space must be hilbert or liouville')
        if vector_space == 'liouville':
            if (self._left_acting_hamiltonian_node_terms is None and not
                self.IS_subspace):
                self._left_acting_hamiltonian_edge_terms = sparse.kron(sparse
                    .identity(self._hamiltonian_node_terms.shape[0]), self.
                    _hamiltonian_edge_terms)
                self._right_acting_hamiltonian_edge_terms = sparse.kron(self
                    ._hamiltonian_edge_terms.T, sparse.identity(self.
                    _hamiltonian_edge_terms.shape[0]))
                self._left_acting_hamiltonian_node_terms = sparse.kron(sparse
                    .identity(self._hamiltonian_node_terms.shape[0]), self.
                    _hamiltonian_node_terms)
                self._right_acting_hamiltonian_node_terms = sparse.kron(self
                    ._hamiltonian_node_terms.T, sparse.identity(self.
                    _hamiltonian_node_terms.shape[0]))
            elif self._left_acting_hamiltonian_node_terms is None and self.IS_subspace:
                self._left_acting_hamiltonian_node_terms = sparse.kron(sparse
                    .identity(self._hamiltonian_node_terms.shape[0]), self.
                    _hamiltonian_node_terms)
                self._right_acting_hamiltonian_node_terms = sparse.kron(self
                    ._hamiltonian_node_terms.T, sparse.identity(self.
                    _hamiltonian_node_terms.shape[0]))
            if not self.IS_subspace:
                return -1.0j * (self.energies[0] * self.
                    _left_acting_hamiltonian_node_terms - self.energies[1] *
                    self._left_acting_hamiltonian_edge_terms) + 1.0j * (
                    self.energies[0] * self.
                    _right_acting_hamiltonian_node_terms - self.energies[1] *
                    self._right_acting_hamiltonian_edge_terms)
            elif self.IS_subspace:
                return -1.0j * (self.energies[0] * self.
                    _left_acting_hamiltonian_node_terms) + 1.0j * (self.
                    energies[0] * self._right_acting_hamiltonian_node_terms)
        else:
            return -1.0j * self.hamiltonian

    @property
    def _diagonal_hamiltonian(self):
        if not self.IS_subspace:
            return self.energies[0
                ] * self._diagonal_hamiltonian_node_terms - self.energies[1
                ] * self._diagonal_hamiltonian_edge_terms
        else:
            return self.energies[0] * self._diagonal_hamiltonian_node_terms

    @property
    def optimum(self):
        if self._is_diagonal:
            return np.max(self._diagonal_hamiltonian).real
        else:
            raise NotImplementedError(
                'Optimum unknown for non-diagonal Hamiltonians')

    def evolve(self, state: State, time):
        if state.is_ket:
            if self._is_diagonal:
                return State(np.exp(-1.0j * time * self.
                    _diagonal_hamiltonian) * state, is_ket=state.is_ket,
                    IS_subspace=state.IS_subspace, code=state.code, graph=
                    self.graph)
            else:
                return State(expm_multiply(-1.0j * time * self.hamiltonian,
                    state), is_ket=state.is_ket, IS_subspace=state.
                    IS_subspace, code=state.code, graph=self.graph)
        elif self._is_diagonal:
            return State(np.exp(-1.0j * time * self._diagonal_hamiltonian) *
                state * np.exp(1.0j * time * self._diagonal_hamiltonian).T,
                is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=
                state.code, graph=self.graph)
        else:
            temp = expm(-1.0j * time * self.hamiltonian)
            return State(temp @ state @ temp.conj().T, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)

    def left_multiply(self, state: State):
        if self._is_diagonal:
            return State(self._diagonal_hamiltonian * state, is_ket=state.
                is_ket, IS_subspace=state.IS_subspace, code=state.code,
                graph=self.graph)
        else:
            return State(self.hamiltonian @ state, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            if self._is_diagonal:
                return State(state.conj().T * self._diagonal_hamiltonian.T,
                    is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                    code=state.code, graph=self.graph)
            else:
                return State((state @ self.hamiltonian.T).conj().T, is_ket=
                    state.is_ket, IS_subspace=state.IS_subspace, code=state
                    .code, graph=self.graph)
        elif self._is_diagonal:
            return State(state * self._diagonal_hamiltonian.T, is_ket=state
                .is_ket, IS_subspace=state.IS_subspace, code=state.code,
                graph=self.graph)
        else:
            return State(state @ self.hamiltonian.T, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)

    def cost_function(self, state: State):
        if state.is_ket:
            if self._is_diagonal:
                return np.real(np.vdot(state, self._diagonal_hamiltonian *
                    state))
            else:
                return np.real(np.vdot(state, self.hamiltonian @ state))
        elif self._is_diagonal:
            return np.real(np.squeeze(tools.trace(self.
                _diagonal_hamiltonian * state)))
        else:
            return np.real(np.squeeze(tools.trace(self.hamiltonian @ state)))

    def optimum_overlap(self, state: State):
        if self._is_diagonal:
            optimum_indices = np.argwhere(self._diagonal_hamiltonian ==
                self.optimum).T[0]
            optimum = np.zeros(self._diagonal_hamiltonian.shape)
            optimum[optimum_indices] = 1
        else:
            raise NotImplementedError(
                'Optimum overlap not implemented for non-diagonal Hamiltonians'
                )
        if state.is_ket:
            if self._is_diagonal:
                return np.real(np.vdot(state, optimum * state))
            else:
                return np.real(np.vdot(state, self.hamiltonian @ state))
        elif self._is_diagonal:
            return np.real(np.squeeze(tools.trace(optimum * state)))
        else:
            return np.real(np.squeeze(tools.trace(self.hamiltonian @ state)))

    def approximation_ratio(self, state: State):
        return self.cost_function(state) / self.optimum


class HamiltonianGlobalPauli(object):

    def __init__(self, pauli: str='X', code=qubit):
        self.code = code
        self.pauli = pauli
        if self.pauli == 'X':
            self._operator = self.code.X
        elif self.pauli == 'Y':
            self._operator = self.code.Y
        elif self.pauli == 'Z':
            self._operator = self.code.Z
        self.hamiltonian = None

    def evolve(self, state: State, alpha):
        if self.hamiltonian is None:
            """Initialize the Hamiltonian only once, as it is costly."""
            self.hamiltonian = tools.tensor_product([self._operator] *
                state.number_logical_qudits)
        return self.code.multiply(np.cos(alpha) * np.identity(state.
            dimension) - 1.0j * np.sin(alpha) * self.hamiltonian)

    def left_multiply(self, state: State):
        all_qubits = list(range(state.number_logical_qudits))
        return self.code.left_multiply(state, all_qubits, [self.pauli] *
            state.number_logical_qudits)

    def right_multiply(self, state: State):
        all_qubits = list(range(state.number_logical_qudits))
        return self.code.right_multiply(state, all_qubits, [self.pauli] *
            state.number_logical_qudits)


class HamiltonianBookatzPenalty(object):

    def __init__(self, code=qubit, energies=(1,)):
        self.code = code
        self.projector = np.identity(self.code.d ** self.code.n
            ) - self.code.code_space_projector
        self.energies = energies

    def evolve(self, state: State, time):
        for i in range(state.number_logical_qudits):
            state = self.code.rotation(state, [i], self.energies[0] * time,
                self.projector, is_idempotent=True)
        return state

    def left_multiply(self, state: State):
        out = np.zeros_like(state, dtype=np.complex128)
        for i in range(state.number_logical_qudits):
            out = out + self.code.left_multiply(state, [i], self.projector)
        return State(self.energies[0] * out, is_ket=state.is_ket,
            IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        out = np.zeros_like(state, dtype=np.complex128)
        for i in range(state.number_logical_qudits):
            out = out + self.code.right_multiply(state, [i], self.projector)
        return State(self.energies[0] * out, is_ket=state.is_ket,
            IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)


class HamiltonianMarvianPenalty(object):

    def __init__(self, Nx, Ny):
        super().__init__()
        self.Nx = Nx
        self.Ny = Ny
        self.n = 3 * Nx * Ny
        hp = np.zeros([2 ** self.n, 2 ** self.n])
        for i in range(int(self.Nx * self.Ny)):
            hp = hp + tools.tensor_product([tools.identity(i * 3), tools.Z(
                ), tools.Z(), tools.identity(self.n - i * 3 - 2)]
                ) + tools.tensor_product([tools.identity(i * 3 + 1), tools.
                X(), tools.X(), tools.identity(self.n - i * 3 - 3)])
        for j in range(self.Ny):
            for k in range(self.Nx):
                if k != self.Nx - 1:
                    hp = hp + tools.tensor_product([tools.identity(j * self
                        .Nx * 3 + k * 3), tools.X(), tools.identity(2),
                        tools.X(), tools.identity(self.n - (j * self.Nx * 3 +
                        k * 3) - 4)]) + tools.tensor_product([tools.
                        identity(j * self.Nx * 3 + k * 3 + 2), tools.Z(),
                        tools.identity(2), tools.Z(), tools.identity(self.n -
                        (j * self.Nx * 3 + k * 3 + 2) - 4)])
                if j != self.Ny - 1:
                    hp = hp + tools.tensor_product([tools.identity(j * self
                        .Nx * 3 + k * 3), tools.X(), tools.identity(3 *
                        self.Nx - 1), tools.X(), tools.identity(self.n - (j *
                        self.Nx * 3 + k * 3) - 3 * self.Nx - 1)]
                        ) + tools.tensor_product([tools.identity(j * self.
                        Nx * 3 + k * 3 + 2), tools.Z(), tools.identity(3 *
                        self.Nx - 1), tools.Z(), tools.identity(self.n - (j *
                        self.Nx * 3 + k * 3 + 2) - 3 * self.Nx - 1)])
        self.hamiltonian = -1 * hp


class HamiltonianHeisenberg(object):

    def __init__(self, graph: Graph, energies=(1, 1), subspace='all', code=
        qubit, IS_subspace=False):
        self.code = code
        self.graph = graph
        self.n = self.graph.n
        self.energies = energies
        self._is_diagonal = True
        assert code is qubit
        if IS_subspace:
            raise NotImplementedError
        self.IS_subspace = IS_subspace
        self._hamiltonian = None
        self.subspace = subspace
        if self.subspace == 'all':
            """Initialize the Hamiltonian."""
            self._hamiltonian_zz = None
            self._hamiltonian_xy = None
        else:
            from scipy.special import comb
            dim = comb(graph.n, (graph.n + self.subspace) / 2)
            if dim % 1 != 0:
                raise Exception('Invalid subspace for number of spins')
            dim = int(dim)
            hamiltonian_zz = np.zeros([dim, 1])
            rows = np.zeros(graph.n * dim, dtype=int)
            columns = np.zeros(graph.n * dim, dtype=int)
            entries = np.zeros(graph.n * dim, dtype=int)
            num_terms = 0
            states = np.zeros((dim, graph.n))
            for i in range(self.code.d ** (self.code.n * self.n)):
                nary = tools.int_to_nary(i, size=graph.n)
                if graph.n - np.sum(nary) == (self.subspace + graph.n) / 2:
                    states[num_terms, ...] = nary
                    num_terms += 1
            num_terms = 0
            for i in range(dim):
                zz = 0
                for a, b in self.graph.edges:
                    if b < a:
                        a, b = b, a
                    if np.abs(states[i, a] - states[i, b]) == 1:
                        zz -= 1
                    else:
                        zz += 1
                    if states[i, a] == 0 and states[i, b] == 1:
                        temp = states[i, ...].copy()
                        temp[a] = 1
                        temp[b] = 0
                        where_matched_temp = np.argwhere(np.sum(np.abs(
                            states - temp), axis=1) == 0).flatten()
                        entries[num_terms] = 1
                        rows[num_terms] = where_matched_temp
                        columns[num_terms] = i
                        entries[num_terms + 1] = 1
                        rows[num_terms + 1] = i
                        columns[num_terms + 1] = where_matched_temp
                        num_terms += 2
                hamiltonian_zz[i, 0] = zz
            self._hamiltonian_zz = sparse.csr_matrix((hamiltonian_zz.
                flatten(), (np.arange(dim), np.arange(dim))), shape=(dim, dim))
            self._hamiltonian_xy = sparse.csr_matrix((entries, (rows,
                columns)), shape=(dim, dim))
            self.states = states

    @property
    def hamiltonian(self):
        if (self._hamiltonian_zz is None or self._hamiltonian_xy is None and
            self.subspace == 'all'):
            z = np.expand_dims(np.diagonal(self.code.Z), axis=0).T

            def my_eye(n):
                return np.ones((np.asarray(self.code.d ** self.code.n) ** n, 1)
                    )
            hamiltonian_zz = np.zeros([self.code.d ** (self.code.n * self.n
                ), 1])
            for a, b in self.graph.edges:
                if b < a:
                    a, b = b, a
                hamiltonian_zz = hamiltonian_zz + self.graph.graph[a][b][
                    'weight'] * tools.tensor_product([my_eye(a), z, my_eye(
                    b - a - 1), z, my_eye(self.n - b - 1)], sparse=not self
                    ._is_diagonal)
            self._hamiltonian_zz = sparse.csr_matrix((hamiltonian_zz.
                flatten(), (np.arange(self.code.d ** (self.code.n * self.n)
                ), np.arange(self.code.d ** (self.code.n * self.n)))),
                shape=(self.code.d ** (self.code.n * self.n), self.code.d **
                (self.code.n * self.n)))
            rows = np.zeros(self.graph.n * self.code.d ** (self.code.n *
                self.n), dtype=int)
            columns = np.zeros(self.graph.n * self.code.d ** (self.code.n *
                self.n), dtype=int)
            entries = np.zeros(self.graph.n * self.code.d ** (self.code.n *
                self.n), dtype=int)
            num_terms = 0
            for i in range(self.code.d ** (self.code.n * self.n)):
                nary = tools.int_to_nary(i, size=self.graph.n)
                for a, b in self.graph.edges:
                    if b < a:
                        a, b = b, a
                    if nary[a] == 0 and nary[b] == 1:
                        temp = nary.copy()
                        temp[a] = 1
                        temp[b] = 0
                        temp = tools.nary_to_int(temp)
                        entries[num_terms] = 1
                        rows[num_terms] = temp
                        columns[num_terms] = i
                        entries[num_terms + 1] = 1
                        rows[num_terms + 1] = i
                        columns[num_terms + 1] = temp
                        num_terms += 2
            self._hamiltonian_xy = sparse.csr_matrix((entries, (rows,
                columns)), shape=(self.code.d ** (self.code.n * self.n), 
                self.code.d ** (self.code.n * self.n)))
        return self.energies[0] * self._hamiltonian_zz + self.energies[1
            ] * self._hamiltonian_xy

    def left_multiply(self, state: State):
        temp = np.zeros(state.shape, dtype=np.complex128)
        for edge in self.graph.edges:
            if self.energies[0] != 0:
                term = self.code.left_multiply(state, [edge[0], edge[1]], [
                    'X', 'X'])
                temp = temp + self.energies[0] * term
                term = self.code.left_multiply(state, [edge[0], edge[1]], [
                    'Y', 'Y'])
                temp = temp + self.energies[0] * term
            if self.energies[1] != 0:
                term = self.code.left_multiply(state, [edge[0], edge[1]], [
                    'Z', 'Z'])
                temp = temp + self.energies[1] * term
        return State(temp, is_ket=state.is_ket, IS_subspace=state.
            IS_subspace, code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        temp = np.zeros(state.shape, dtype=np.complex128)
        for edge in self.graph.edges:
            if self.energies[0] != 0:
                term = self.code.right_multiply(state, [edge[0], edge[1]],
                    ['X', 'X'])
                temp = temp + self.energies[0] * term
                term = self.code.right_multiply(state, [edge[0], edge[1]],
                    ['Y', 'Y'])
                temp = temp + self.energies[0] * term
            if self.energies[1] != 0:
                term = self.code.right_multiply(state, [edge[0], edge[1]],
                    ['Z', 'Z'])
                temp = temp + self.energies[1] * term
        return State(temp, is_ket=state.is_ket, IS_subspace=state.
            IS_subspace, code=state.code, graph=self.graph)

    def evolve(self, state: State, time):
        if not state.is_ket:
            exp_hamiltonian = expm(-1.0j * time * self.hamiltonian)
            return State(exp_hamiltonian @ state @ exp_hamiltonian.conj().T,
                is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=
                state.code, graph=self.graph)
        if state.is_ket:
            return State(expm_multiply(-1.0j * time * self.hamiltonian,
                state), is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                code=state.code, graph=self.graph)

    def cost_function(self, state: State):
        if state.is_ket:
            return np.real(np.vdot(state, self.hamiltonian * state))
        else:
            return np.real(np.squeeze(tools.trace(self.hamiltonian * state)))


class HamiltonianEnergyShift(object):

    def __init__(self, index: int=0, energies=(1,), code=qubit, IS_subspace
        =False, graph=None):
        """Default is that the first element in transition is the higher energy s."""
        self.index = index
        self.graph = graph
        self.energies = energies
        self.code = code
        if not self.code.logical_code:
            if not 0 <= self.index < self.code.d:
                raise Exception('Index exceeds qudit dimension.')
            self._operator = np.zeros((self.code.d, self.code.d))
            self._operator[self.index, self.index] = 1
        else:
            if self.index != 0 and self.index != 1:
                raise Exception(
                    'Logical codes are qubits, so index must be 0 or 1.')
            if self.index == 0:
                self._operator = self.code.Q
            elif self.index == 1:
                self._operator = self.code.P
        self.IS_subspace = IS_subspace
        if self.IS_subspace:
            assert graph is not None
            assert isinstance(graph, Graph)
            if code is not qubit:
                IS, num_IS = graph.independent_sets_qudit(self.code)
                self._diagonal_hamiltonian = np.zeros((num_IS, 1), dtype=float)
                for k in range(num_IS):
                    self._diagonal_hamiltonian[k, 0] = np.sum(IS[k, ...] ==
                        self.index)
                self._hamiltonian = sparse.csr_matrix((self.
                    _diagonal_hamiltonian.T[0], (np.arange(num_IS), np.
                    arange(num_IS))), shape=(num_IS, num_IS))
            else:
                independent_sets = enumerate_independent_sets(graph.graph)
                k = self.graph.num_independent_sets - 2
                self.mis_size = 0
                hamiltonian = np.zeros(self.graph.num_independent_sets,
                    dtype=float)
                hamiltonian[-1] = 0
                for i in independent_sets:
                    hamiltonian[k] = len(i)
                    k -= 1
                self._hamiltonian = sparse.csr_matrix((hamiltonian, (np.
                    arange(self.graph.num_independent_sets), np.arange(self
                    .graph.num_independent_sets))), shape=(self.graph.
                    num_independent_sets, self.graph.num_independent_sets))
        else:
            self._hamiltonian = None

    @property
    def hamiltonian(self):
        if self._hamiltonian is None:
            assert not self.IS_subspace
            try:
                assert self.graph is not None
            except AssertionError:
                print(
                    'self.graph must be not None to generate the Hamiltonian property.'
                    )
            self._hamiltonian = sparse.csr_matrix(((self.code.d * self.code
                .n) ** self.graph.n, (self.code.d * self.code.n) ** self.
                graph.n))
            for i in range(self.graph.n):
                self._hamiltonian = self._hamiltonian + tools.tensor_product([
                    sparse.identity((self.code.d * self.code.n) ** i), self
                    ._operator, sparse.identity((self.code.d * self.code.n) **
                    (self.graph.n - i - 1))], sparse=True)
        return self.energies[0] * self._hamiltonian

    def left_multiply(self, state: State):
        if not self.IS_subspace:
            temp = np.zeros_like(state)
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.left_multiply(state, [i], self.
                        _operator)
                elif not self.code.logical_code:
                    ind = self.code.d ** i
                    out = np.zeros_like(state, dtype=np.complex128)
                    if state.is_ket:
                        state = state.reshape((-1, self.code.d, ind), order='F'
                            )
                        out = out.reshape((-1, self.code.d, ind), order='F')
                        out[:, self.index, :] = state[:, self.index, :]
                        state = state.reshape(state_shape, order='F')
                        out = out.reshape(state_shape, order='F')
                    else:
                        out = out.reshape((-1, self.code.d, self.code.d **
                            (state.number_physical_qudits - 1), self.code.d,
                            ind), order='F')
                        state = state.reshape((-1, self.code.d, self.code.d **
                            (state.number_physical_qudits - 1), self.code.d,
                            ind), order='F')
                        out[:, self.index, :, :, :] = state[:, self.index,
                            :, :, :]
                        state = state.reshape(state_shape, order='F')
                        out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)
        else:
            return State(self.energies[0] * self._diagonal_hamiltonian *
                state, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            print(
                'Warning: right multiply functionality currently applies the operator and daggers the state.'
                )
            return self.left_multiply(state).conj().T
        elif not self.IS_subspace:
            temp = np.zeros_like(state)
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.right_multiply(state, [i], self
                        ._operator)
                else:
                    ind = self.code.d ** i
                    out = np.zeros_like(state)
                    out = out.reshape((-1, self.code.d, self.code.d ** (
                        state.number_physical_qudits - 1), self.code.d, ind
                        ), order='F')
                    state = state.reshape((-1, self.code.d, self.code.d **
                        (state.number_physical_qudits - 1), self.code.d,
                        ind), order='F')
                    out[:, :, :, self.index, :] = state[:, :, :, self.index, :]
                    state = state.reshape(state_shape, order='F')
                    out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)
        else:
            return State(self.energies[0] * state * self.
                _diagonal_hamiltonian.T, is_ket=state.is_ket, IS_subspace=
                state.IS_subspace, code=state.code, graph=self.graph)

    def evolve(self, state: State, time):
        """
        Use reshape to efficiently implement evolution under :math:`H_B=\\\\sum_i X_i`
        """
        if not self.IS_subspace:
            out = state.copy()
            for i in range(state.number_logical_qudits):
                out = self.code.rotation(out, [i], self.energies[0] * time,
                    self._operator)
            return out
        elif state.is_ket:
            return State(np.exp(-1.0j * time * self.energies[0] * self.
                _diagonal_hamiltonian) * state, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                graph)
        else:
            exp_hamiltonian = np.exp(-1.0j * time * self.energies[0] * self
                ._diagonal_hamiltonian)
            return State(exp_hamiltonian * state * exp_hamiltonian.conj().T,
                is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=
                state.code, graph=self.graph)


class HamiltonianRydberg(object):

    def __init__(self, tails_graph: Graph, hard_constraint_graph=None,
        index: int=0, energies=(1,), code=qubit, IS_subspace=False):
        """Default is that the first element in transition is the higher energy s."""
        self.index = index
        self.energies = energies
        self.code = code
        self.tails_graph = tails_graph
        self.hard_constraint_graph = hard_constraint_graph
        if self.code.logical_code:
            raise Exception('Logical codes not supported.')
        self.IS_subspace = IS_subspace
        try:
            if (self.tails_graph.periodic or self.hard_constraint_graph.
                periodic):
                raise Exception('Periodic graphs not yet supported.')
        except:
            pass
        if self.IS_subspace:
            assert hard_constraint_graph is not None
            assert isinstance(hard_constraint_graph, Graph)
            if code is not qubit:
                IS, num_IS = (hard_constraint_graph.
                    generate_independent_sets_qudit(self.code))
                self._diagonal_hamiltonian = np.zeros((num_IS, 1), dtype=float)
                for k in range(num_IS):
                    weight = 0
                    where_rydberg = np.argwhere(self.hard_constraint_graph.
                        independent_sets[k] == 0)
                    if len(where_rydberg) > 0:
                        where_rydberg = where_rydberg.T[0]
                        for i, node1 in enumerate(where_rydberg):
                            for j, node2 in enumerate(where_rydberg[i + 1:]):
                                weight += self.tails_graph.graph[node1][node2][
                                    'weight']
                    self._diagonal_hamiltonian[k, 0] = weight
                self._hamiltonian = sparse.csr_matrix((self.
                    _diagonal_hamiltonian.T[0], (np.arange(num_IS), np.
                    arange(num_IS))), shape=(num_IS, num_IS))
            else:
                independent_sets = enumerate_independent_sets(self.
                    hard_constraint_graph.graph)
                k = self.hard_constraint_graph.num_independent_sets - 2
                self.mis_size = 0
                self._diagonal_hamiltonian = np.zeros(self.
                    hard_constraint_graph.num_independent_sets, dtype=float)
                self._diagonal_hamiltonian[-1] = 0
                for i in independent_sets:
                    weight = 0
                    where_rydberg = np.array(i)
                    if len(where_rydberg) > 0:
                        for i, node1 in enumerate(where_rydberg):
                            for j, node2 in enumerate(where_rydberg[i + 1:]):
                                weight += self.tails_graph.graph[node1][node2][
                                    'weight']
                    self._diagonal_hamiltonian[k] = weight
                    k -= 1
                self._hamiltonian = sparse.csr_matrix((self.
                    _diagonal_hamiltonian, (np.arange(self.
                    hard_constraint_graph.num_independent_sets), np.arange(
                    self.hard_constraint_graph.num_independent_sets))),
                    shape=(self.hard_constraint_graph.num_independent_sets,
                    self.hard_constraint_graph.num_independent_sets))
                self._diagonal_hamiltonian = np.expand_dims(self.
                    _diagonal_hamiltonian, axis=1)
        else:
            self._hamiltonian = None

    @property
    def hamiltonian(self):
        if self._hamiltonian is None:
            raise Exception
        return self.energies[0] * self._hamiltonian

    def left_multiply(self, state: State):
        return State(self.energies[0] * self._diagonal_hamiltonian * state,
            is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.
            code, graph=self.hard_constraint_graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            print(
                'Warning: right multiply functionality currently applies the operator and daggers the state.'
                )
            return self.left_multiply(state).conj().T
        elif not self.IS_subspace:
            temp = np.zeros_like(state)
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.right_multiply(state, [i], self
                        ._operator)
                else:
                    ind = self.code.d ** i
                    out = np.zeros_like(state)
                    out = out.reshape((-1, self.code.d, self.code.d ** (
                        state.number_physical_qudits - 1), self.code.d, ind
                        ), order='F')
                    state = state.reshape((-1, self.code.d, self.code.d **
                        (state.number_physical_qudits - 1), self.code.d,
                        ind), order='F')
                    out[:, :, :, self.index, :] = state[:, :, :, self.index, :]
                    state = state.reshape(state_shape, order='F')
                    out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                hard_constraint_graph)
        else:
            return State(self.energies[0] * state * self.
                _diagonal_hamiltonian.T, is_ket=state.is_ket, IS_subspace=
                state.IS_subspace, code=state.code, graph=self.
                hard_constraint_graph)

    def evolve(self, state: State, time):
        """
        Use reshape to efficiently implement evolution under :math:`H_B=\\\\sum_i X_i`
        """
        if not self.IS_subspace:
            out = state.copy()
            for i in range(state.number_logical_qudits):
                out = self.code.rotation(out, [i], self.energies[0] * time,
                    self._operator)
            return out
        elif state.is_ket:
            return State(np.exp(-1.0j * time * self.energies[0] * self.
                _diagonal_hamiltonian) * state, is_ket=state.is_ket,
                IS_subspace=state.IS_subspace, code=state.code, graph=self.
                hard_constraint_graph)
        else:
            exp_hamiltonian = np.exp(-1.0j * time * self.energies[0] * self
                ._diagonal_hamiltonian)
            return State(exp_hamiltonian * state * exp_hamiltonian.conj().T,
                is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=
                state.code, graph=self.hard_constraint_graph)

import numpy as np
from qsim.codes import qubit, rydberg
from qsim.codes.quantum_state import State
from qsim import tools
from scipy.linalg import expm
import scipy.sparse as sparse
from scipy.sparse.linalg import expm_multiply
from qsim.graph_algorithms.graph import Graph, enumerate_independent_sets


class HamiltonianDriver(object):
    def __init__(self, transition: tuple = (0, 1), energies: tuple = (1,), pauli='X', code=qubit, graph=None,
                 IS_subspace=False):
        """Default is that the first element in transition is the higher energy s."""
        self.transition = transition
        self.energies = energies
        assert pauli in ['X', 'Y', 'Z']
        self.pauli = pauli
        self.code = code
        self.graph = graph
        if self.pauli == 'X' and not self.code.logical_code:
            self._operator = np.zeros((self.code.d, self.code.d))
            self._operator[self.transition[1], self.transition[0]] = 1
            self._operator[self.transition[0], self.transition[1]] = 1
        elif self.pauli == 'Y' and not self.code.logical_code:
            self._operator = np.zeros((self.code.d, self.code.d))
            self._operator[self.transition[1], self.transition[0]] = 1j
            self._operator[self.transition[0], self.transition[1]] = -1j
        elif self.pauli == 'Z' and not self.code.logical_code:
            self._operator = np.zeros((self.code.d, self.code.d))
            self._operator[self.transition[0], self.transition[0]] = 1
            self._operator[self.transition[1], self.transition[1]] = -1
        # If a logical code, we should use the normal qubit operators because we assume the code is a qubit
        elif self.pauli == 'X' and self.code.logical_code:
            self._operator = self.code.X
        elif self.pauli == 'Y' and self.code.logical_code:
            self._operator = self.code.Y
        elif self.pauli == 'Z' and self.code.logical_code:
            self._operator = self.code.Z
        self.IS_subspace = IS_subspace

        if self.IS_subspace:
            # Generate sparse mixing Hamiltonian
            assert isinstance(graph, Graph)
            if code is not qubit:
                IS, num_IS = graph.generate_independent_sets_qudit(self.code)
                if self.pauli == 'Z':
                    self._diagonal_hamiltonian = np.zeros((num_IS, 1))
                    for k in range(num_IS):
                        self._diagonal_hamiltonian[k, 0] = np.sum(IS[k, ...] == self.transition[0]) - np.sum(
                            IS[k, ...] == self.transition[1])

                    self._csr_hamiltonian = sparse.csr_matrix((self._diagonal_hamiltonian.T[0], (np.arange(num_IS),
                                                                                                 np.arange(num_IS))))

                    self._hamiltonian = self._csr_hamiltonian

                else:
                    # For each IS, look at spin flips generated by the laser
                    # Over-allocate space
                    rows = np.zeros(graph.n * num_IS, dtype=int)
                    columns = np.zeros(graph.n * num_IS, dtype=int)
                    entries = np.zeros(graph.n * num_IS, dtype=int)
                    num_terms = 0
                    for i in range(num_IS):
                        for j in range(graph.n):
                            if IS[i, j] == self.transition[0]:
                                # Flip spin at this location
                                # Get binary representation
                                temp = IS[i, ...].copy()
                                temp[j] = self.transition[1]
                                where_matched = (np.argwhere(np.sum(np.abs(IS - temp), axis=1) == 0).flatten())
                                if len(where_matched) > 0:
                                    # This is a valid spin flip by removing a node
                                    rows[num_terms] = where_matched[0]
                                    columns[num_terms] = i
                                    if self.pauli == 'X':
                                        entries[num_terms] = 1
                                    elif self.pauli == 'Y':
                                        # entries[num_terms] = -1j
                                        entries[num_terms] = 1j
                                    num_terms += 1
                    """else:
                        for (i, key) in enumerate(self.graph.independent_sets_dict):
                            num_neighbors = len(self.graph.independent_sets_dict[key][1])
                            rows[num_terms:num_terms + num_neighbors] = self.graph.independent_sets_dict[key][1]
                            columns[num_terms:num_terms + num_neighbors] = np.ones(num_neighbors) * \
                                                                           self.graph.independent_sets_dict[key][0]
                            if self.pauli == 'X':
                                entries[num_terms:num_terms + num_neighbors] = 1
                            elif self.pauli == 'Y':
                                entries[num_terms:num_terms + num_neighbors] = -1j
                            num_terms += num_neighbors"""

                    # Cut off the excess in the arrays
                    columns = columns[:2 * num_terms]
                    rows = rows[:2 * num_terms]
                    entries = entries[:2 * num_terms]
                    # Populate the second half of the entries according to self.pauli
                    if self.pauli == 'X':
                        columns[num_terms:2 * num_terms] = rows[:num_terms]
                        rows[num_terms:2 * num_terms] = columns[:num_terms]
                        entries[num_terms:2 * num_terms] = entries[:num_terms]
                    elif self.pauli == 'Y':
                        columns[num_terms:2 * num_terms] = rows[:num_terms]
                        rows[num_terms:2 * num_terms] = columns[:num_terms]
                        entries[num_terms:2 * num_terms] = -1 * entries[:num_terms]
                    # Now, construct the Hamiltonian
                    self._csr_hamiltonian = sparse.csr_matrix((entries, (rows, columns)), shape=(num_IS, num_IS))
                    self._hamiltonian = self._csr_hamiltonian
            else:
                # Use graph generator functions
                if self.pauli == 'Z':
                    independent_sets = enumerate_independent_sets(graph.graph)
                    # Generate a list of integers corresponding to the independent sets in binary
                    # All ones
                    k = self.graph.num_independent_sets - 2
                    self.mis_size = 0
                    hamiltonian = np.zeros(self.graph.num_independent_sets, dtype=float)
                    hamiltonian[-1] = -1 * self.graph.n
                    for i in independent_sets:
                        hamiltonian[k] = len(i) - (self.graph.n - len(i))
                        k -= 1

                    self._hamiltonian = sparse.csr_matrix(
                        (hamiltonian,
                         (np.arange(self.graph.num_independent_sets), np.arange(self.graph.num_independent_sets))),
                        shape=(self.graph.num_independent_sets, self.graph.num_independent_sets))

                else:
                    independent_sets = enumerate_independent_sets(graph.graph)
                    # Generate a list of integers corresponding to the independent sets in binary
                    previous_size = 0
                    self.mis_size = 0
                    independent_sets_dict = {(): self.graph.num_independent_sets - 1}
                    rows = []
                    columns = []
                    entries = []
                    k = self.graph.num_independent_sets - 2
                    for i in independent_sets:
                        current_size = len(i)
                        if current_size - previous_size > 1:
                            previous_size = current_size - 1
                            # Clear out the dictionary with terms we can't connect to
                            for key in list(independent_sets_dict):
                                if len(key) != previous_size:
                                    independent_sets_dict.pop(key)
                        independent_sets_dict[tuple(i)] = k
                        for (j, node) in enumerate(i):
                            i_removed = i.copy()
                            i_removed.pop(j)
                            index = independent_sets_dict[tuple(i_removed)]
                            # Index is the current independent set with a single node removed
                            rows.append(k)
                            columns.append(index)
                            rows.append(index)
                            columns.append(k)
                            if self.pauli == 'Y':
                                entries.append(-1j)
                                entries.append(1j)
                            else:
                                entries.append(1)
                                entries.append(1)
                        k -= 1
                    # Now, construct the Hamiltonian
                    self._csr_hamiltonian = sparse.csr_matrix((entries, (rows, columns)),
                                                              shape=(self.graph.num_independent_sets,
                                                                     self.graph.num_independent_sets))
                    self._hamiltonian = self._csr_hamiltonian
        else:
            self._hamiltonian = None
        self._left_acting_hamiltonian = None
        self._right_acting_hamiltonian = None

    @property
    def hamiltonian(self):
        if self._hamiltonian is None:
            assert not self.IS_subspace
            try:
                assert self.graph is not None
            except AssertionError:
                print('self.graph must be not None to generate the Hamiltonian property.')
            self._hamiltonian = sparse.csr_matrix(((self.code.d * self.code.n) ** self.graph.n,
                                                   (self.code.d * self.code.n) ** self.graph.n))
            for i in range(self.graph.n):
                self._hamiltonian = self._hamiltonian + tools.tensor_product(
                    [sparse.identity((self.code.d * self.code.n) ** i),
                     self._operator,
                     sparse.identity((self.code.d * self.code.n) ** (self.graph.n - i - 1))],
                    sparse=True)
        return self.energies[0] * self._hamiltonian

    @property
    def evolution_operator(self, vector_space='hilbert'):
        if vector_space != 'hilbert' and vector_space != 'liouville':
            raise Exception('Attribute vector_space must be hilbert or liouville')
        if vector_space == 'liouville':
            if self._left_acting_hamiltonian is None:
                self._left_acting_hamiltonian = sparse.kron(sparse.identity(self._hamiltonian.shape[0]),
                                                            self._hamiltonian)
                self._right_acting_hamiltonian = sparse.kron(self._hamiltonian.T,
                                                             sparse.identity(self._hamiltonian.shape[0]))

            return -1j * self.energies[0] * self._left_acting_hamiltonian + 1j * self.energies[0] * \
                   self._right_acting_hamiltonian
        else:
            return -1j * self.hamiltonian

    def left_multiply(self, state: State):
        if not self.IS_subspace:
            temp = np.zeros_like(state)
            # For each logical qubit
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    if self.pauli == 'X':
                        temp = temp + self.code.left_multiply(state, [i], ['X'])
                    elif self.pauli == 'Y':
                        temp = temp + self.code.left_multiply(state, [i], ['Y'])
                    elif self.pauli == 'Z':
                        temp = temp + self.code.left_multiply(state, [i], ['Z'])
                elif not self.code.logical_code:
                    ind = self.code.d ** i
                    out = np.zeros_like(state, dtype=np.complex128)
                    if state.is_ket:
                        state = state.reshape((-1, self.code.d, ind), order='F')
                        # Note index start from the right (sN,...,s3,s2,s1)
                        out = out.reshape((-1, self.code.d, ind), order='F')
                        if self.pauli == 'X':  # Sigma_X
                            # We want to exchange two indices
                            out[:, [self.transition[0], self.transition[1]], :] = \
                                state[:, [self.transition[1], self.transition[0]], :]
                        elif self.pauli == 'Y':  # Sigma_Y
                            out[:, [self.transition[0], self.transition[1]], :] = \
                                state[:, [self.transition[1], self.transition[0]], :]
                            out[:, self.transition[0], :] = -1j * out[:, self.transition[0], :]
                            out[:, self.transition[1], :] = 1j * out[:, self.transition[1], :]
                        elif self.pauli == 'Z':  # Sigma_Z
                            out[:, [self.transition[0], self.transition[1]], :] = \
                                state[:, [self.transition[0], self.transition[1]], :]
                            out[:, self.transition[1], :] = -1 * out[:, self.transition[1], :]
                        state = state.reshape(state_shape, order='F')
                        out = out.reshape(state_shape, order='F')
                    else:
                        out = out.reshape((-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1),
                                           self.code.d, ind), order='F')
                        state = state.reshape((-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1),
                                               self.code.d, ind), order='F')

                        if self.pauli == 'X':  # Sigma_X
                            out[:, [self.transition[0], self.transition[1]], :, :, :] = \
                                state[:, [self.transition[1], self.transition[0]], :, :, :]
                        elif self.pauli == 'Y':  # Sigma_Y
                            out[:, [self.transition[0], self.transition[1]], :, :, :] = \
                                state[:, [self.transition[1], self.transition[0]], :, :, :]
                            out[:, self.transition[0], :, :, :] = -1j * out[:, self.transition[0], :, :, :]
                            out[:, self.transition[1], :, :, :] = 1j * out[:, self.transition[1], :, :, :]
                        elif self.pauli == 'Z':  # Sigma_Z
                            out[:, [self.transition[0], self.transition[1]], :, :, :] = \
                                state[:, [self.transition[0], self.transition[1]], :, :, :]
                            out[:, self.transition[1], :, :, :] = -1 * out[:, self.transition[1], :, :, :]
                        state = state.reshape(state_shape, order='F')
                        out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                         graph=self.graph)
        else:
            # Handle dimensions
            if self.pauli == 'Z' and not self.code.logical_code:  # In this case, the Hamiltonian is diagonal
                return State(self.energies[0] * self._diagonal_hamiltonian * state, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                return State(self.energies[0] * self._csr_hamiltonian @ state, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            print('Warning: right multiply functionality currently applies the operator and daggers the state.')
            return self.left_multiply(state).conj().T
        if not self.IS_subspace:
            temp = np.zeros_like(state)
            # For each physical qubit
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    if self.pauli == 'X':
                        temp = temp + self.code.right_multiply(state, [i], ['X'])
                    elif self.pauli == 'Y':
                        temp = temp + self.code.right_multiply(state, [i], ['Y'])
                    elif self.pauli == 'Z':
                        temp = temp + self.code.right_multiply(state, [i], ['Z'])
                elif not self.code.logical_code:
                    ind = self.code.d ** i
                    out = np.zeros_like(state)
                    out = out.reshape(
                        (-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1), self.code.d, ind),
                        order='F')
                    state = state.reshape(
                        (-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1), self.code.d, ind),
                        order='F')
                    if self.pauli == 'X' and not self.code.logical_code:  # Sigma_X
                        out[:, :, :, [self.transition[0], self.transition[1]], :] = state[:, :, :, [self.transition[1],
                                                                                                    self.transition[0]],
                                                                                    :]
                    elif self.pauli == 'Y' and not self.code.logical_code:  # Sigma_Y
                        out[:, :, :, [self.transition[0], self.transition[1]], :] = state[:, :, :, [self.transition[1],
                                                                                                    self.transition[0]],
                                                                                    :]
                        out[:, :, :, self.transition[0], :] = -1j * out[:, :, :, self.transition[0], :]
                        out[:, :, :, self.transition[1], :] = 1j * out[:, :, :, self.transition[1], :]
                    elif self.pauli == 'Z' and not self.code.logical_code:  # Sigma_Z
                        out[:, :, :, [self.transition[0], self.transition[1]], :] = state[:, :, :, [self.transition[0],
                                                                                                    self.transition[1]],
                                                                                    :]
                        out[:, :, :, self.transition[1], :] = -1 * state[:, :, :, self.transition[1], :]
                        state = state.reshape(state_shape, order='F')
                    state = state.reshape(state_shape, order='F')
                    out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                         graph=self.graph)
        else:
            if self.pauli == 'Z' and not self.code.logical_code:  # In this case, the Hamiltonian is diagonal
                return State(state * self.hamiltonian.T, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                             code=state.code, graph=self.graph)
            else:
                return State(state @ self.hamiltonian.T.conj(), is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                             code=state.code, graph=self.graph)

    def evolve(self, state: State, time):
        r"""
        Use reshape to efficiently implement evolution under :math:`H_B=\\sum_i X_i`
        """
        if not self.IS_subspace:
            # We don't want to modify the original s
            out = state.copy()
            for i in range(state.number_logical_qudits):
                # Note that self._operator is not necessarily involutary
                if self.pauli == 'X':
                    out = self.code.rotation(out, [i], self.energies[0] * time, self._operator)
                elif self.pauli == 'Y':
                    out = self.code.rotation(out, [i], self.energies[0] * time, self._operator)
                elif self.pauli == 'Z':
                    out = self.code.rotation(out, [i], self.energies[0] * time, self._operator)
            return out
        else:
            if state.is_ket:
                # Handle dimensions
                if self.hamiltonian.shape[1] == 1:
                    return State(np.exp(-1j * time * self.hamiltonian) * state, is_ket=state.is_ket,
                                 IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
                else:
                    return State(expm_multiply(-1j * time * self.hamiltonian, state),
                                 is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                if self.hamiltonian.shape[1] == 1:
                    exp_hamiltonian = np.exp(-1j * time * self.hamiltonian)
                    return State(exp_hamiltonian * state * exp_hamiltonian.conj().T,
                                 is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
                else:
                    exp_hamiltonian = expm(-1j * time * self.hamiltonian)
                    return State(exp_hamiltonian @ state @ exp_hamiltonian.conj().T,
                                 is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)


class HamiltonianMaxCut(object):
    def __init__(self, graph: Graph, code=qubit, energies=(1,), cost_function=True, use_Z2_symmetry=False):
        # If MIS is true, create an MIS Hamiltonian. Otherwise, make a MaxCut Hamiltonian
        r"""
        Generate a vector corresponding to the diagonal of the MaxCut Hamiltonian.
        """
        self.code = code
        self.energies = energies
        # Make sure all edges have weight attribute; default to 1

        self.graph = graph
        self.optimization = 'max'
        self.n = self.graph.n
        if use_Z2_symmetry:
            c = np.zeros([self.code.d ** (self.code.n * (self.n - 1)), 1])
        else:
            c = np.zeros([self.code.d ** (self.code.n * self.n), 1])
        if tools.is_diagonal(self.code.Z):
            self._is_diagonal = True
            z = np.expand_dims(np.diagonal(self.code.Z), axis=0).T

            def my_eye(n):
                return np.ones((np.asarray(self.code.d ** self.code.n) ** n, 1))
        else:
            self._is_diagonal = False
            # Compute the optimum first. We don't care that this takes extra time, since it only needs to run once
            z = np.expand_dims(np.diagonal(qubit.Z), axis=0).T

            def my_eye(n):
                return np.ones((np.asarray(qubit.d) ** n, 1))

            for a, b in self.graph.edges:
                if b < a:
                    a, b = b, a
                if use_Z2_symmetry:
                    if cost_function:
                        if a == min(self.graph.nodes):
                            c = c - 1 / 2 * graph.graph[a][b]['weight'] * (tools.tensor_product(
                                [my_eye(b - 1), z, my_eye(self.n - b - 1)]) - my_eye(self.n - 1))
                        else:
                            c = c - 1 / 2 * graph.graph[a][b]['weight'] * (tools.tensor_product(
                                [my_eye(a - 1), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)]) - my_eye(self.n - 1))
                    else:
                        if a == min(self.graph.nodes):
                            c = c + graph.graph[a][b]['weight'] * (tools.tensor_product(
                                [my_eye(b - 1), z, my_eye(self.n - b - 1)]))
                        else:
                            c = c + graph.graph[a][b]['weight'] * (tools.tensor_product(
                                [my_eye(a - 1), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)]))
                else:
                    if cost_function:
                        c = c - 1 / 2 * graph.graph[a][b]['weight'] * (tools.tensor_product(
                            [my_eye(a), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)]) - my_eye(self.n))
                    else:
                        c = c + graph.graph[a][b]['weight'] * tools.tensor_product(
                            [my_eye(a), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)])
            self._optimum = np.max(c).real
            c = sparse.csr_matrix((self.code.d ** (self.code.n * self.n), self.code.d ** (self.code.n * self.n)))

            z = sparse.csr_matrix(self.code.Z)

            def my_eye(n):
                return sparse.csr_matrix(np.ones(np.asarray(z.shape[0]) ** n),
                                         (np.asarray(z.shape[0]) ** n, np.asarray(z.shape[0]) ** n))

        for a, b in self.graph.edges:
            if b < a:
                a, b = b, a

            if cost_function:
                if use_Z2_symmetry:
                    if a == min(self.graph.nodes):
                        c = c - 1 / 2 * graph.graph[a][b]['weight'] * (tools.tensor_product(
                            [my_eye(b - 1), z, my_eye(self.n - b - 1)]) - my_eye(self.n - 1))
                    else:
                        c = c - 1 / 2 * graph.graph[a][b]['weight'] * (tools.tensor_product(
                            [my_eye(a - 1), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)]) - my_eye(self.n - 1))
                else:
                    c = c - 1 / 2 * graph.graph[a][b]['weight'] * (tools.tensor_product(
                        [my_eye(a), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)],
                        sparse=(not self._is_diagonal)) - my_eye(
                        self.n))
            else:
                if use_Z2_symmetry:
                    if a == min(self.graph.nodes):
                        c = c + graph.graph[a][b]['weight'] * (tools.tensor_product(
                            [my_eye(b - 1), z, my_eye(self.n - b - 1)]))
                    else:
                        c = c + graph.graph[a][b]['weight'] * (tools.tensor_product(
                            [my_eye(a - 1), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)]))

                else:
                    c = c + graph.graph[a][b]['weight'] * (tools.tensor_product(
                        [my_eye(a), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)],
                        sparse=(not self._is_diagonal)))
        if self._is_diagonal:
            self._diagonal_hamiltonian = c
            self._optimum = np.max(c).real
            if use_Z2_symmetry:
                c = sparse.csr_matrix((c.flatten(), (np.arange(self.code.d ** (self.code.n * (self.n - 1))),
                                                     np.arange(self.code.d ** (self.code.n * (self.n - 1))))),
                                      shape=(self.code.d ** (self.code.n * (self.n - 1)),
                                             self.code.d ** (self.code.n * (self.n - 1))))
            else:
                c = sparse.csr_matrix((c.flatten(), (np.arange(self.code.d ** (self.code.n * self.n)),
                                                     np.arange(self.code.d ** (self.code.n * self.n)))),
                                      shape=(self.code.d ** (self.code.n * self.n),
                                             self.code.d ** (self.code.n * self.n)))
        else:
            # c is already the right shape, just convert it to a csr matrix
            c = sparse.csr_matrix(c)
        self._hamiltonian = c
        self._left_acting_hamiltonian = None
        self._right_acting_hamiltonian = None

    @property
    def hamiltonian(self):
        return self.energies[0] * self._hamiltonian

    @property
    def evolution_operator(self, vector_space='hilbert'):
        if vector_space != 'hilbert' and vector_space != 'liouville':
            raise Exception('Attribute vector_space must be hilbert or liouville')
        if vector_space == 'liouville':
            if self._left_acting_hamiltonian is None:
                self._left_acting_hamiltonian = sparse.kron(sparse.identity(self._hamiltonian.shape[0]),
                                                            self._hamiltonian)
                self._right_acting_hamiltonian = sparse.kron(self._hamiltonian.T,
                                                             sparse.identity(self._hamiltonian.shape[0]))

            return -1j * self.energies[0] * self._left_acting_hamiltonian + 1j * self.energies[0] * \
                   self._right_acting_hamiltonian
        else:
            return -1j * self.hamiltonian

    @property
    def optimum(self):
        # Optimum for non-diagonal Hamiltonians can be found by computing the optimum in the standard basis,
        # which is done in self.__init__()
        return self.energies[0] * self._optimum

    def evolve(self, state: State, time):
        if state.is_ket:
            if self._is_diagonal:
                # It's quicker to exponentiate a diagonal array than use expm_multiply
                return State(np.exp(-1j * time * self._diagonal_hamiltonian) * state, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                return State(expm_multiply(-1j * time * self.hamiltonian, state), is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
        else:
            if self._is_diagonal:
                return State(np.exp(-1j * time * self._diagonal_hamiltonian) * state * np.exp(
                    1j * time * self._diagonal_hamiltonian).T, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                             code=state.code, graph=self.graph)
            else:
                temp = expm(-1j * time * self.hamiltonian)
                return State(temp @ state @ temp.conj().T, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                             code=state.code, graph=self.graph)

    def left_multiply(self, state: State):
        if self._is_diagonal:
            return State(self._diagonal_hamiltonian * state, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                         code=state.code, graph=self.graph)
        else:
            return State(self.hamiltonian @ state, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                         code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        # Already real, so you don't need to conjugate
        if state.is_ket:
            if self._is_diagonal:
                return State(state.conj().T * self._diagonal_hamiltonian.T, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                return State((state @ self.hamiltonian.T).conj().T, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
        else:
            if self._is_diagonal:
                # Density matrices are already Hermitian, so you don't need to dagger
                return State(state * self._diagonal_hamiltonian.T, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                return State(state @ self.hamiltonian.T, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

    def cost_function(self, state: State):
        # Need to project into the IS subspace
        # Returns <s|C|s>
        if state.is_ket:
            if self._is_diagonal:
                return np.real(np.vdot(state, self._diagonal_hamiltonian * state))
            else:
                return np.real(np.vdot(state, self.hamiltonian @ state))
        else:
            # Density matrix
            if self._is_diagonal:
                return np.real(np.squeeze(tools.trace(self._diagonal_hamiltonian * state)))
            else:
                return np.real(np.squeeze(tools.trace(self.hamiltonian @ state)))

    def optimum_overlap(self, state: State):
        # Returns \sum_i <s|opt_i><opt_i|s>
        if self._is_diagonal:
            optimum_indices = np.argwhere(self._diagonal_hamiltonian == self.optimum).T[0]
            # Construct an operator that is zero everywhere except at the optimum
            optimum = np.zeros(self._diagonal_hamiltonian.shape)
            optimum[optimum_indices] = 1
        else:
            # The plan for this is to basically use the code for _is_diagonal to identify the logical qubit subspaces
            # which encode the optimum. Then, make an operator that's the identity in those subspaces
            raise NotImplementedError('Optimum overlap not implemented for non-diagonal Hamiltonians')
        if state.is_ket:
            return np.real(np.vdot(state, optimum * state))
        else:
            # Density matrix
            return np.real(np.squeeze(tools.trace(optimum * state)))

    def approximation_ratio(self, state: State):
        # Returns <s|C|s>/optimum
        return self.cost_function(state) / self.optimum


class HamiltonianMIS(object):
    def __init__(self, graph: Graph, energies=(1, 1), code=qubit, IS_subspace=False):
        r"""
        Generate a vector corresponding to the diagonal of the MIS Hamiltonian.
        """
        if energies == (1, 1) and IS_subspace:
            energies = (1,)
        self.code = code
        self.graph = graph
        self.n = self.graph.n
        self.energies = energies
        self.IS_subspace = IS_subspace
        self.optimization = 'max'
        if not self.IS_subspace:
            # Store node and edge terms separately so the Hamiltonian can be dynamically updated when energies
            # are changed

            if tools.is_diagonal(self.code.Q):
                self._hamiltonian_edge_terms = np.zeros([1, (self.code.d ** self.code.n) ** self.n])
                self._hamiltonian_node_terms = np.zeros([1, (self.code.d ** self.code.n) ** self.n])
                self._is_diagonal = True

                Q = np.expand_dims(np.diagonal(self.code.Q), axis=0)

                def my_eye(n):
                    return np.ones(np.asarray(self.code.d ** self.code.n) ** n)
            else:
                # TODO: generate a sparse matrix instead
                self._hamiltonian_edge_terms = np.zeros([(self.code.d ** self.code.n) ** self.n,
                                                         (self.code.d ** self.code.n) ** self.n])
                self._hamiltonian_node_terms = np.zeros([(self.code.d ** self.code.n) ** self.n,
                                                         (self.code.d ** self.code.n) ** self.n])

                Q = np.expand_dims(np.diagonal(qubit.Q), axis=0)

                def my_eye(n):
                    return np.ones(np.asarray(qubit.d ** qubit.n) ** n)

                self._optimum_edge_terms = np.zeros([(qubit.d ** qubit.n) ** self.n,
                                                     (qubit.d ** qubit.n) ** self.n])
                self._optimum_node_terms = np.zeros([(qubit.d ** qubit.n) ** self.n,
                                                     (qubit.d ** qubit.n) ** self.n])

                for i, j in graph.graph.edges:
                    if j < i:
                        i, j = j, i
                    self._optimum_edge_terms = self._optimum_edge_terms + graph.graph.edges[(i, j)]['weight'] * \
                                               tools.tensor_product(
                                                   [my_eye(i), Q, my_eye(j - i - 1), Q, my_eye(self.n - j - 1)])

                for i in graph.graph.nodes:
                    self._optimum_node_terms = self._optimum_node_terms + graph.graph.nodes[i]['weight'] * \
                                               tools.tensor_product([my_eye(i), Q, my_eye(self.n - i - 1)])

                self._is_diagonal = False
                Q = self.code.Q

                def my_eye(n):
                    return np.identity(np.asarray(self.code.d ** self.code.n) ** n)
            for i, j in graph.graph.edges:
                if j < i:
                    i, j = j, i
                self._hamiltonian_edge_terms = self._hamiltonian_edge_terms + graph.graph.edges[(i, j)]['weight'] * \
                                               tools.tensor_product(
                                                   [my_eye(i), Q, my_eye(j - i - 1), Q, my_eye(self.n - j - 1)])
            for i in graph.graph.nodes:
                self._hamiltonian_node_terms = self._hamiltonian_node_terms + graph.graph.nodes[i]['weight'] * \
                                               tools.tensor_product([my_eye(i), Q, my_eye(self.n - i - 1)])
            self._hamiltonian_node_terms = self._hamiltonian_node_terms.T
            self._hamiltonian_edge_terms = self._hamiltonian_edge_terms.T
            if self._is_diagonal:
                self._optimum_edge_terms = self._hamiltonian_edge_terms
                self._optimum_node_terms = self._hamiltonian_node_terms
                self._diagonal_hamiltonian_edge_terms = self._hamiltonian_edge_terms.copy()
                self._diagonal_hamiltonian_node_terms = self._hamiltonian_node_terms.copy()
                self._hamiltonian_node_terms = sparse.csr_matrix(
                    (self._hamiltonian_node_terms.flatten(), (np.arange(self.code.d ** (self.code.n * self.n)),
                                                              np.arange(self.code.d ** (self.code.n * self.n)))),
                    shape=(self.code.d ** (self.code.n * self.n),
                           self.code.d ** (self.code.n * self.n)))
                self._hamiltonian_edge_terms = sparse.csr_matrix(
                    (self._hamiltonian_edge_terms.flatten(), (np.arange(self.code.d ** (self.code.n * self.n)),
                                                              np.arange(self.code.d ** (self.code.n * self.n)))),
                    shape=(self.code.d ** (self.code.n * self.n),
                           self.code.d ** (self.code.n * self.n)))
                # TODO: what happens to _optimum_node_terms if not _is_diagonal
            else:
                self._diagonal_hamiltonian_edge_terms = self._hamiltonian_edge_terms.copy()
                self._diagonal_hamiltonian_node_terms = self._hamiltonian_node_terms.copy()
                self._hamiltonian_edge_terms = sparse.csr_matrix(self._hamiltonian_edge_terms)
                self._hamiltonian_node_terms = sparse.csr_matrix(self._hamiltonian_node_terms)
            self._left_acting_hamiltonian_edge_terms = None
            self._right_acting_hamiltonian_edge_terms = None

        else:
            self._is_diagonal = True
            if not (self.code == qubit or self.code == rydberg):
                raise NotImplementedError("IS subspace only implemented for qubit and Rydberg codes.")
            # Don't generate anything that depends on the entire Hilbert space as to save space

            # These are your independent sets of the original graphs, ordered by node and size
            if self.code == qubit:
                node_weights = np.asarray([self.graph.graph.nodes[i]['weight'] for i in range(self.graph.n)])
                independent_sets = enumerate_independent_sets(self.graph.graph)
                # Generate a list of integers corresponding to the independent sets in binary
                # All ones
                k = self.graph.num_independent_sets - 2
                self.mis_size = 0
                C = np.zeros(self.graph.num_independent_sets, dtype=float)
                C[-1] = 0
                for i in independent_sets:
                    C[k] = np.sum([node_weights[j] for j in i])
                    k -= 1
                self._hamiltonian = sparse.csr_matrix((C, (np.arange(self.graph.num_independent_sets),
                                                           np.arange(self.graph.num_independent_sets))))

                C = np.expand_dims(C, axis=0).T

            # Otherwise, we need to include the possibility that we are in one of many ground space states
            elif self.code == rydberg:
                # TODO: fix this to reflect the new way of notating independent sets!
                # Count the number of elements in the ground space and map to their representation in ternary
                # Determine how large to make the array
                independent_sets, num_IS = self.graph.generate_independent_sets_qudit(self.code)
                # Generate Hamiltonian from independent sets
                node_weights = np.asarray([self.graph.graph.nodes[i]['weight'] for i in range(self.graph.n)])
                C = np.zeros((num_IS, 1), dtype=np.complex128)
                for k in independent_sets:
                    C[k, 0] = np.sum((independent_sets[k][2] == 0) * node_weights)
            self._diagonal_hamiltonian_node_terms = C
            C = C.flatten()

            self._hamiltonian_node_terms = sparse.csr_matrix((
                C, (np.arange(len(C)), np.arange(len(C)))), shape=(len(C), len(C)))

        self._left_acting_hamiltonian_node_terms = None
        self._right_acting_hamiltonian_node_terms = None

    @property
    def hamiltonian(self):
        if not self.IS_subspace:
            return self.energies[0] * self._hamiltonian_node_terms - self.energies[1] * self._hamiltonian_edge_terms
        else:
            return self.energies[0] * self._hamiltonian_node_terms

    @property
    def evolution_operator(self, vector_space='hilbert'):
        if vector_space != 'hilbert' and vector_space != 'liouville':
            raise Exception('Attribute vector_space must be hilbert or liouville')
        if vector_space == 'liouville':
            if self._left_acting_hamiltonian_node_terms is None and not self.IS_subspace:
                self._left_acting_hamiltonian_edge_terms = sparse.kron(sparse.identity(
                    self._hamiltonian_node_terms.shape[0]), self._hamiltonian_edge_terms)
                self._right_acting_hamiltonian_edge_terms = sparse.kron(
                    self._hamiltonian_edge_terms.T, sparse.identity(self._hamiltonian_edge_terms.shape[0]))
                self._left_acting_hamiltonian_node_terms = sparse.kron(sparse.identity(
                    self._hamiltonian_node_terms.shape[0]), self._hamiltonian_node_terms)
                self._right_acting_hamiltonian_node_terms = sparse.kron(
                    self._hamiltonian_node_terms.T, sparse.identity(self._hamiltonian_node_terms.shape[0]))
            elif self._left_acting_hamiltonian_node_terms is None and self.IS_subspace:
                self._left_acting_hamiltonian_node_terms = sparse.kron(sparse.identity(
                    self._hamiltonian_node_terms.shape[0]), self._hamiltonian_node_terms)
                self._right_acting_hamiltonian_node_terms = sparse.kron(
                    self._hamiltonian_node_terms.T, sparse.identity(self._hamiltonian_node_terms.shape[0]))
            if not self.IS_subspace:
                return -1j * (self.energies[0] * self._left_acting_hamiltonian_node_terms - self.energies[1] *
                              self._left_acting_hamiltonian_edge_terms) + 1j * \
                       (self.energies[0] * self._right_acting_hamiltonian_node_terms - self.energies[1] *
                        self._right_acting_hamiltonian_edge_terms)
            elif self.IS_subspace:
                return -1j * (self.energies[0] * self._left_acting_hamiltonian_node_terms) + 1j * \
                       (self.energies[0] * self._right_acting_hamiltonian_node_terms)
        else:
            return -1j * self.hamiltonian

    @property
    def _diagonal_hamiltonian(self):
        if not self.IS_subspace:
            return self.energies[0] * self._diagonal_hamiltonian_node_terms - self.energies[
                1] * self._diagonal_hamiltonian_edge_terms
        else:
            return self.energies[0] * self._diagonal_hamiltonian_node_terms

    @property
    def optimum(self):
        # This needs to be recomputed because the optimum depends on the energies
        # TODO: figure out what to compute if not _is_diagonal
        if self._is_diagonal:
            return np.max(self._diagonal_hamiltonian).real
        else:
            raise NotImplementedError('Optimum unknown for non-diagonal Hamiltonians')

    def evolve(self, state: State, time):
        if state.is_ket:
            if self._is_diagonal:
                return State(np.exp(-1j * time * self._diagonal_hamiltonian) * state, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                return State(expm_multiply(-1j * time * self.hamiltonian, state), is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
        else:
            if self._is_diagonal:
                return State(np.exp(-1j * time * self._diagonal_hamiltonian) * state * np.exp(
                    1j * time * self._diagonal_hamiltonian).T, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                             code=state.code, graph=self.graph)
            else:
                temp = expm(-1j * time * self.hamiltonian)
                return State(temp @ state @ temp.conj().T, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                             code=state.code, graph=self.graph)

    def left_multiply(self, state: State):
        if self._is_diagonal:
            return State(self._diagonal_hamiltonian * state, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                         code=state.code, graph=self.graph)
        else:
            return State(self.hamiltonian @ state, is_ket=state.is_ket, IS_subspace=state.IS_subspace,
                         code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        # Already real, so you don't need to conjugate
        if state.is_ket:
            if self._is_diagonal:
                return State(state.conj().T * self._diagonal_hamiltonian.T, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                return State((state @ self.hamiltonian.T).conj().T, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
        else:
            if self._is_diagonal:
                # Density matrices are already Hermitian, so you don't need to dagger
                return State(state * self._diagonal_hamiltonian.T, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                return State(state @ self.hamiltonian.T, is_ket=state.is_ket,
                             IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

    def cost_function(self, state: State):
        # Returns <s|C|s>
        if state.is_ket:
            if self._is_diagonal:
                return np.real(np.vdot(state, self._diagonal_hamiltonian * state))
            else:
                return np.real(np.vdot(state, self.hamiltonian @ state))
        else:
            # Density matrix
            if self._is_diagonal:
                return np.real(np.squeeze(tools.trace(self._diagonal_hamiltonian * state)))
            else:
                return np.real(np.squeeze(tools.trace(self.hamiltonian @ state)))

    def optimum_overlap(self, state: State):
        # Returns \sum_i <s|opt_i><opt_i|s>
        if self._is_diagonal:
            optimum_indices = np.argwhere(self._diagonal_hamiltonian == self.optimum).T[0]
            # Construct an operator that is zero everywhere except at the optimum
            optimum = np.zeros(self._diagonal_hamiltonian.shape)
            optimum[optimum_indices] = 1
        else:
            raise NotImplementedError('Optimum overlap not implemented for non-diagonal Hamiltonians')
        if state.is_ket:
            if self._is_diagonal:
                return np.real(np.vdot(state, optimum * state))
            else:
                return np.real(np.vdot(state, self.hamiltonian @ state))
        else:
            # Density matrix
            if self._is_diagonal:
                return np.real(np.squeeze(tools.trace(optimum * state)))
            else:
                return np.real(np.squeeze(tools.trace(self.hamiltonian @ state)))

    def approximation_ratio(self, state: State):
        # Returns <s|C|s>/optimum
        return self.cost_function(state) / self.optimum


class HamiltonianGlobalPauli(object):
    def __init__(self, pauli: str = 'X', code=qubit):
        self.code = code
        self.pauli = pauli
        if self.pauli == 'X':
            self._operator = self.code.X
        elif self.pauli == 'Y':
            self._operator = self.code.Y
        elif self.pauli == 'Z':
            self._operator = self.code.Z
        self.hamiltonian = None

    def evolve(self, state: State, alpha):
        if self.hamiltonian is None:
            """Initialize the Hamiltonian only once, as it is costly."""
            # TODO: make this a sparse matrix!
            self.hamiltonian = tools.tensor_product([self._operator] * state.number_logical_qudits)
        return self.code.multiply(np.cos(alpha) * np.identity(state.dimension) - 1j * np.sin(alpha) * self.hamiltonian)

    def left_multiply(self, state: State):
        all_qubits = list(range(state.number_logical_qudits))
        return self.code.left_multiply(state, all_qubits, [self.pauli] * state.number_logical_qudits)

    def right_multiply(self, state: State):
        all_qubits = list(range(state.number_logical_qudits))
        return self.code.right_multiply(state, all_qubits, [self.pauli] * state.number_logical_qudits)


class HamiltonianBookatzPenalty(object):
    def __init__(self, code=qubit, energies=(1,)):
        self.code = code
        self.projector = np.identity(self.code.d ** self.code.n) - self.code.code_space_projector
        self.energies = energies

    def evolve(self, state: State, time):
        # Term for a single qubit
        for i in range(state.number_logical_qudits):
            state = self.code.rotation(state, [i], self.energies[0] * time, self.projector, is_idempotent=True)
        return state

    def left_multiply(self, state: State):
        out = np.zeros_like(state, dtype=np.complex128)
        for i in range(state.number_logical_qudits):
            out = out + self.code.left_multiply(state, [i], self.projector)
        return State(self.energies[0] * out, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                     graph=self.graph)

    def right_multiply(self, state: State):
        out = np.zeros_like(state, dtype=np.complex128)
        for i in range(state.number_logical_qudits):
            out = out + self.code.right_multiply(state, [i], self.projector)
        return State(self.energies[0] * out, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                     graph=self.graph)


class HamiltonianMarvianPenalty(object):
    def __init__(self, Nx, Ny):
        super().__init__()
        self.Nx = Nx
        self.Ny = Ny
        self.n = 3 * Nx * Ny
        # Generate Hamiltonian
        # Two by two geometry (can be generalized in the future)
        hp = np.zeros([2 ** self.n, 2 ** self.n])
        for i in range(int(self.Nx * self.Ny)):
            # Add gauge interactions within a single logical qubit
            hp = hp + tools.tensor_product(
                [tools.identity(i * 3), tools.Z(), tools.Z(),
                 tools.identity(self.n - i * 3 - 2)]) + tools.tensor_product(
                [tools.identity(i * 3 + 1), tools.X(), tools.X(), tools.identity(self.n - i * 3 - 3)])
        # Between rows
        for j in range(self.Ny):
            # j is the number of rows
            for k in range(self.Nx):
                # k is the number of columns
                # Need to deal with edge effects
                # Add gauge interactions within a single logical qubit
                if k != self.Nx - 1:
                    # Along the same row
                    hp = hp + tools.tensor_product(
                        [tools.identity(j * self.Nx * 3 + k * 3), tools.X(), tools.identity(2), tools.X(),
                         tools.identity(self.n - (j * self.Nx * 3 + k * 3) - 4)]) + \
                         tools.tensor_product(
                             [tools.identity(j * self.Nx * 3 + k * 3 + 2), tools.Z(), tools.identity(2), tools.Z(),
                              tools.identity(self.n - (j * self.Nx * 3 + k * 3 + 2) - 4)])
                    # Along the same column
                if j != self.Ny - 1:
                    hp = hp + tools.tensor_product(
                        [tools.identity(j * self.Nx * 3 + k * 3), tools.X(), tools.identity(3 * self.Nx - 1),
                         tools.X(),
                         tools.identity(self.n - (j * self.Nx * 3 + k * 3) - 3 * self.Nx - 1)]) + \
                         tools.tensor_product(
                             [tools.identity(j * self.Nx * 3 + k * 3 + 2), tools.Z(),
                              tools.identity(3 * self.Nx - 1),
                              tools.Z(), tools.identity(self.n - (j * self.Nx * 3 + k * 3 + 2) - 3 * self.Nx - 1)])
        self.hamiltonian = -1 * hp


class HamiltonianHeisenberg(object):
    def __init__(self, graph: Graph, energies=(1, 1), subspace='all', code=qubit, IS_subspace=False):
        self.code = code
        self.graph = graph
        self.n = self.graph.n
        self.energies = energies
        self._is_diagonal = True
        assert code is qubit
        if IS_subspace:
            raise NotImplementedError
        self.IS_subspace = IS_subspace
        self._hamiltonian = None
        self.subspace = subspace
        if self.subspace == 'all':
            """Initialize the Hamiltonian."""
            self._hamiltonian_zz = None
            self._hamiltonian_xy = None
        else:
            # For each IS, look at spin flips generated by the laser
            # Over-allocate space
            from scipy.special import comb
            dim = comb(graph.n, (graph.n + self.subspace) / 2)
            if dim % 1 != 0:
                raise Exception('Invalid subspace for number of spins')
            dim = int(dim)
            hamiltonian_zz = np.zeros([dim, 1])
            rows = np.zeros(graph.n * dim, dtype=int)
            columns = np.zeros(graph.n * dim, dtype=int)
            entries = np.zeros(graph.n * dim, dtype=int)
            num_terms = 0
            states = np.zeros((dim, graph.n))
            for i in range(self.code.d ** (self.code.n * self.n)):
                nary = tools.int_to_nary(i, size=graph.n)
                if graph.n - np.sum(nary) == (self.subspace + graph.n) / 2:
                    states[num_terms, ...] = nary
                    num_terms += 1
            num_terms = 0
            for i in range(dim):
                zz = 0
                for a, b in self.graph.edges:
                    if b < a:
                        a, b = b, a
                    if np.abs(states[i, a] - states[i, b]) == 1:
                        zz -= 1
                    else:
                        zz += 1
                    if states[i, a] == 0 and states[i, b] == 1:
                        # Flip spin at this location
                        # Get binary representation
                        temp = states[i, ...].copy()
                        temp[a] = 1
                        temp[b] = 0
                        where_matched_temp = (np.argwhere(np.sum(np.abs(states - temp), axis=1) == 0).flatten())
                        entries[num_terms] = 1
                        rows[num_terms] = where_matched_temp
                        columns[num_terms] = i
                        entries[num_terms + 1] = 1
                        rows[num_terms + 1] = i
                        columns[num_terms + 1] = where_matched_temp
                        num_terms += 2
                hamiltonian_zz[i, 0] = zz
            # Now, construct the Hamiltonian
            self._hamiltonian_zz = sparse.csr_matrix(
                (hamiltonian_zz.flatten(), (np.arange(dim),
                                            np.arange(dim))), shape=(dim, dim))
            self._hamiltonian_xy = sparse.csr_matrix((entries, (rows, columns)), shape=(dim, dim))
            self.states = states

    @property
    def hamiltonian(self):
        if self._hamiltonian_zz is None or self._hamiltonian_xy is None and self.subspace == 'all':
            z = np.expand_dims(np.diagonal(self.code.Z), axis=0).T

            def my_eye(n):
                return np.ones((np.asarray(self.code.d ** self.code.n) ** n, 1))

            hamiltonian_zz = np.zeros([self.code.d ** (self.code.n * self.n), 1])
            for a, b in self.graph.edges:
                if b < a:
                    a, b = b, a
                hamiltonian_zz = hamiltonian_zz + self.graph.graph[a][b]['weight'] * (tools.tensor_product(
                    [my_eye(a), z, my_eye(b - a - 1), z, my_eye(self.n - b - 1)],
                    sparse=(not self._is_diagonal)))
            self._hamiltonian_zz = sparse.csr_matrix(
                (hamiltonian_zz.flatten(), (np.arange(self.code.d ** (self.code.n * self.n)),
                                            np.arange(self.code.d ** (self.code.n * self.n)))),
                shape=(self.code.d ** (self.code.n * self.n),
                       self.code.d ** (self.code.n * self.n)))

            # For each IS, look at spin flips generated by the laser
            # Over-allocate space
            rows = np.zeros(self.graph.n * self.code.d ** (self.code.n * self.n), dtype=int)
            columns = np.zeros(self.graph.n * self.code.d ** (self.code.n * self.n), dtype=int)
            entries = np.zeros(self.graph.n * self.code.d ** (self.code.n * self.n), dtype=int)
            num_terms = 0
            for i in range(self.code.d ** (self.code.n * self.n)):
                nary = tools.int_to_nary(i, size=self.graph.n)
                for a, b in self.graph.edges:
                    if b < a:
                        a, b = b, a
                    if nary[a] == 0 and nary[b] == 1:
                        # Flip spin at this location
                        # Get binary representation
                        temp = nary.copy()
                        temp[a] = 1
                        temp[b] = 0
                        temp = tools.nary_to_int(temp)
                        entries[num_terms] = 1
                        rows[num_terms] = temp
                        columns[num_terms] = i
                        entries[num_terms + 1] = 1
                        rows[num_terms + 1] = i
                        columns[num_terms + 1] = temp
                        num_terms += 2
            # Now, construct the Hamiltonian
            self._hamiltonian_xy = sparse.csr_matrix((entries, (rows, columns)),
                                                     shape=(self.code.d ** (self.code.n * self.n),
                                                            self.code.d ** (self.code.n * self.n)))
        return self.energies[0] * self._hamiltonian_zz + self.energies[1] * self._hamiltonian_xy

    def left_multiply(self, state: State):
        temp = np.zeros(state.shape, dtype=np.complex128)
        for edge in self.graph.edges:
            if self.energies[0] != 0:
                term = self.code.left_multiply(state, [edge[0], edge[1]], ['X', 'X'])
                temp = temp + self.energies[0] * term
                term = self.code.left_multiply(state, [edge[0], edge[1]], ['Y', 'Y'])
                temp = temp + self.energies[0] * term
            if self.energies[1] != 0:
                term = self.code.left_multiply(state, [edge[0], edge[1]], ['Z', 'Z'])
                temp = temp + self.energies[1] * term
        return State(temp, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        temp = np.zeros(state.shape, dtype=np.complex128)
        for edge in self.graph.edges:
            if self.energies[0] != 0:
                term = self.code.right_multiply(state, [edge[0], edge[1]], ['X', 'X'])
                temp = temp + self.energies[0] * term
                term = self.code.right_multiply(state, [edge[0], edge[1]], ['Y', 'Y'])
                temp = temp + self.energies[0] * term
            if self.energies[1] != 0:
                term = self.code.right_multiply(state, [edge[0], edge[1]], ['Z', 'Z'])
                temp = temp + self.energies[1] * term
        return State(temp, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

    def evolve(self, state: State, time):
        if not state.is_ket:
            exp_hamiltonian = expm(-1j * time * self.hamiltonian)
            return State(exp_hamiltonian @ state @ exp_hamiltonian.conj().T,
                         is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
        if state.is_ket:
            return State(expm_multiply(-1j * time * self.hamiltonian, state), is_ket=state.is_ket,
                         IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)

    def cost_function(self, state: State):
        # Need to project into the IS subspace
        # Returns <s|C|s>
        if state.is_ket:
            return np.real(np.vdot(state, self.hamiltonian * state))
        else:
            # Density matrix
            return np.real(np.squeeze(tools.trace(self.hamiltonian * state)))


class HamiltonianEnergyShift(object):
    def __init__(self, index: int = 0, energies=(1,), code=qubit, IS_subspace=False, graph=None):
        """Default is that the first element in transition is the higher energy s."""
        self.index = index
        self.graph = graph
        self.energies = energies
        self.code = code
        if not self.code.logical_code:
            if not 0 <= self.index < self.code.d:
                raise Exception('Index exceeds qudit dimension.')
            self._operator = np.zeros((self.code.d, self.code.d))
            self._operator[self.index, self.index] = 1
        else:
            if self.index != 0 and self.index != 1:
                raise Exception('Logical codes are qubits, so index must be 0 or 1.')
            if self.index == 0:
                self._operator = self.code.Q
            elif self.index == 1:
                self._operator = self.code.P
        self.IS_subspace = IS_subspace
        if self.IS_subspace:
            # Generate sparse mixing Hamiltonian
            assert graph is not None
            assert isinstance(graph, Graph)
            if code is not qubit:
                IS, num_IS = graph.independent_sets_qudit(self.code)
                self._diagonal_hamiltonian = np.zeros((num_IS, 1), dtype=float)
                for k in range(num_IS):
                    self._diagonal_hamiltonian[k, 0] = np.sum(IS[k, ...] == self.index)
                self._hamiltonian = sparse.csr_matrix(
                    (self._diagonal_hamiltonian.T[0], (np.arange(num_IS), np.arange(num_IS))),
                    shape=(num_IS, num_IS))
            else:
                # We have already solved for this information
                independent_sets = enumerate_independent_sets(graph.graph)
                # Generate a list of integers corresponding to the independent sets in binary
                # All ones
                k = self.graph.num_independent_sets - 2
                self.mis_size = 0
                hamiltonian = np.zeros(self.graph.num_independent_sets, dtype=float)
                hamiltonian[-1] = 0
                for i in independent_sets:
                    hamiltonian[k] = len(i)
                    k -= 1

                self._hamiltonian = sparse.csr_matrix(
                    (hamiltonian,
                     (np.arange(self.graph.num_independent_sets), np.arange(self.graph.num_independent_sets))),
                    shape=(self.graph.num_independent_sets, self.graph.num_independent_sets))
        else:
            # Use full Hilbert space
            self._hamiltonian = None

    @property
    def hamiltonian(self):
        if self._hamiltonian is None:
            assert not self.IS_subspace
            try:
                assert self.graph is not None
            except AssertionError:
                print('self.graph must be not None to generate the Hamiltonian property.')
            self._hamiltonian = sparse.csr_matrix(((self.code.d * self.code.n) ** self.graph.n,
                                                   (self.code.d * self.code.n) ** self.graph.n))
            for i in range(self.graph.n):
                self._hamiltonian = self._hamiltonian + tools.tensor_product(
                    [sparse.identity((self.code.d * self.code.n) ** i),
                     self._operator,
                     sparse.identity((self.code.d * self.code.n) ** (self.graph.n - i - 1))],
                    sparse=True)
        return self.energies[0] * self._hamiltonian

    def left_multiply(self, state: State):
        if not self.IS_subspace:
            temp = np.zeros_like(state)
            # For each logical qubit
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.left_multiply(state, [i], self._operator)
                elif not self.code.logical_code:
                    ind = self.code.d ** i
                    out = np.zeros_like(state, dtype=np.complex128)
                    if state.is_ket:
                        state = state.reshape((-1, self.code.d, ind), order='F')
                        # Note index start from the right (sN,...,s3,s2,s1)
                        out = out.reshape((-1, self.code.d, ind), order='F')
                        out[:, self.index, :] = state[:, self.index, :]
                        state = state.reshape(state_shape, order='F')
                        out = out.reshape(state_shape, order='F')
                    else:
                        out = out.reshape((-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1),
                                           self.code.d, ind), order='F')
                        state = state.reshape((-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1),
                                               self.code.d, ind), order='F')
                        out[:, self.index, :, :, :] = state[:, self.index, :, :, :]
                        state = state.reshape(state_shape, order='F')
                        out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                         graph=self.graph)
        else:
            # Handle dimensions
            return State(self.energies[0] * self._diagonal_hamiltonian * state, is_ket=state.is_ket,
                         IS_subspace=state.IS_subspace,
                         code=state.code, graph=self.graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            print('Warning: right multiply functionality currently applies the operator and daggers the state.')
            return self.left_multiply(state).conj().T
        elif not self.IS_subspace:
            temp = np.zeros_like(state)
            # For each physical qubit
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.right_multiply(state, [i], self._operator)
                else:
                    ind = self.code.d ** i
                    out = np.zeros_like(state)
                    out = out.reshape(
                        (-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1), self.code.d, ind),
                        order='F')
                    state = state.reshape(
                        (-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1), self.code.d, ind),
                        order='F')
                    out[:, :, :, self.index, :] = state[:, :, :, self.index, :]
                    state = state.reshape(state_shape, order='F')
                    out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                         graph=self.graph)
        else:
            return State(self.energies[0] * state * self._diagonal_hamiltonian.T, is_ket=state.is_ket,
                         IS_subspace=state.IS_subspace,
                         code=state.code, graph=self.graph)

    def evolve(self, state: State, time):
        r"""
        Use reshape to efficiently implement evolution under :math:`H_B=\\sum_i X_i`
        """
        if not self.IS_subspace:
            # We don't want to modify the original s
            out = state.copy()
            for i in range(state.number_logical_qudits):
                # Note that self._operator is not necessarily involutary
                out = self.code.rotation(out, [i], self.energies[0] * time, self._operator)
            return out
        else:
            if state.is_ket:
                # Handle dimensions
                return State(np.exp(-1j * time * self.energies[0] * self._diagonal_hamiltonian) * state,
                             is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)
            else:
                exp_hamiltonian = np.exp(-1j * time * self.energies[0] * self._diagonal_hamiltonian)
                return State(exp_hamiltonian * state * exp_hamiltonian.conj().T,
                             is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code, graph=self.graph)


class HamiltonianRydberg(object):
    def __init__(self, tails_graph: Graph, hard_constraint_graph=None, index: int = 0,
                 energies=(1,), code=qubit, IS_subspace=False):
        """Default is that the first element in transition is the higher energy s."""
        self.index = index
        self.energies = energies
        self.code = code
        self.tails_graph = tails_graph
        self.hard_constraint_graph = hard_constraint_graph
        if self.code.logical_code:
            raise Exception('Logical codes not supported.')
        self.IS_subspace = IS_subspace
        try:
            if self.tails_graph.periodic or self.hard_constraint_graph.periodic:
                raise Exception('Periodic graphs not yet supported.')
        except:
            pass
        # Generate tails matrix
        if self.IS_subspace:
            # Generate sparse mixing Hamiltonian
            assert hard_constraint_graph is not None
            assert isinstance(hard_constraint_graph, Graph)
            if code is not qubit:
                IS, num_IS = hard_constraint_graph.generate_independent_sets_qudit(self.code)
                self._diagonal_hamiltonian = np.zeros((num_IS, 1), dtype=float)
                for k in range(num_IS):
                    # Get state as bit string
                    # Multiply by weights matrix
                    # Sum result
                    weight = 0
                    where_rydberg = np.argwhere(self.hard_constraint_graph.independent_sets[k] == 0)
                    if len(where_rydberg) > 0:
                        where_rydberg = where_rydberg.T[0]
                        for (i, node1) in enumerate(where_rydberg):
                            for (j, node2) in enumerate(where_rydberg[i + 1:]):
                                weight += self.tails_graph.graph[node1][node2]['weight']
                    self._diagonal_hamiltonian[k, 0] = weight
                self._hamiltonian = sparse.csr_matrix(
                    (self._diagonal_hamiltonian.T[0], (np.arange(num_IS), np.arange(num_IS))),
                    shape=(num_IS, num_IS))
            else:
                # We have already solved for this information
                independent_sets = enumerate_independent_sets(self.hard_constraint_graph.graph)
                # Generate a list of integers corresponding to the independent sets in binary
                # All ones
                k = self.hard_constraint_graph.num_independent_sets - 2
                self.mis_size = 0
                self._diagonal_hamiltonian = np.zeros(self.hard_constraint_graph.num_independent_sets, dtype=float)
                self._diagonal_hamiltonian[-1] = 0
                for i in independent_sets:
                    weight = 0
                    where_rydberg = np.array(i)
                    if len(where_rydberg) > 0:
                        for (i, node1) in enumerate(where_rydberg):
                            for (j, node2) in enumerate(where_rydberg[i + 1:]):
                                weight += self.tails_graph.graph[node1][node2]['weight']
                    self._diagonal_hamiltonian[k] = weight
                    k -= 1

                self._hamiltonian = sparse.csr_matrix(
                    (self._diagonal_hamiltonian,
                     (np.arange(self.hard_constraint_graph.num_independent_sets), np.arange(self.hard_constraint_graph.num_independent_sets))),
                    shape=(self.hard_constraint_graph.num_independent_sets, self.hard_constraint_graph.num_independent_sets))
                self._diagonal_hamiltonian = np.expand_dims(self._diagonal_hamiltonian, axis=1)
        else:
            # Use full Hilbert space
            self._hamiltonian = None

    @property
    def hamiltonian(self):
        if self._hamiltonian is None:
            raise Exception
        return self.energies[0] * self._hamiltonian

    def left_multiply(self, state: State):
        return State(self.energies[0] * self._diagonal_hamiltonian * state, is_ket=state.is_ket,
                     IS_subspace=state.IS_subspace,
                     code=state.code, graph=self.hard_constraint_graph)

    def right_multiply(self, state: State):
        if state.is_ket:
            print('Warning: right multiply functionality currently applies the operator and daggers the state.')
            return self.left_multiply(state).conj().T
        elif not self.IS_subspace:
            temp = np.zeros_like(state)
            # For each physical qubit
            state_shape = state.shape
            for i in range(state.number_logical_qudits):
                if self.code.logical_code:
                    temp = temp + self.code.right_multiply(state, [i], self._operator)
                else:
                    ind = self.code.d ** i
                    out = np.zeros_like(state)
                    out = out.reshape(
                        (-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1), self.code.d, ind),
                        order='F')
                    state = state.reshape(
                        (-1, self.code.d, self.code.d ** (state.number_physical_qudits - 1), self.code.d, ind),
                        order='F')
                    out[:, :, :, self.index, :] = state[:, :, :, self.index, :]
                    state = state.reshape(state_shape, order='F')
                    out = out.reshape(state_shape, order='F')
                    temp = temp + out
            return State(self.energies[0] * temp, is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                         graph=self.hard_constraint_graph)
        else:
            return State(self.energies[0] * state * self._diagonal_hamiltonian.T, is_ket=state.is_ket,
                         IS_subspace=state.IS_subspace,
                         code=state.code, graph=self.hard_constraint_graph)

    def evolve(self, state: State, time):
        r"""
        Use reshape to efficiently implement evolution under :math:`H_B=\\sum_i X_i`
        """
        if not self.IS_subspace:
            # We don't want to modify the original s
            out = state.copy()
            for i in range(state.number_logical_qudits):
                # Note that self._operator is not necessarily involutary
                out = self.code.rotation(out, [i], self.energies[0] * time, self._operator)
            return out
        else:
            if state.is_ket:
                # Handle dimensions
                return State(np.exp(-1j * time * self.energies[0] * self._diagonal_hamiltonian) * state,
                             is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                             graph=self.hard_constraint_graph)
            else:
                exp_hamiltonian = np.exp(-1j * time * self.energies[0] * self._diagonal_hamiltonian)
                return State(exp_hamiltonian * state * exp_hamiltonian.conj().T,
                             is_ket=state.is_ket, IS_subspace=state.IS_subspace, code=state.code,
                             graph=self.hard_constraint_graph)
