import numpy as np
from sympy import Symbol
import unittest

from openqaoa.qaoa_components import PauliOp, Hamiltonian
"""
Unittest based testing of the PauliOp and Hamiltonian classes
"""

class TestingOperators(unittest.TestCase):

    def test_pauli_init(self):
        """
        Test the initialization method of the PauliOp class.

        The test consists in checking correct initializations for representative examples.
        """

        # Define input examples of indices and strings
        input_indices = [(0, 1, 1), (4, 5, 0, 1), (2, 0, 7, 4, 5),
                         (0, 1, 5, 3, 7, 6), (6, 5, 4, 3, 1, 1, 0)]
        input_strings = ['XZX', 'YYYY', 'ZYXZX', 'YYZZXX', 'ZYXYZYX']

        # Create Pauli operators from input examples
        paulis = [PauliOp(string, index)
                  for string, index in zip(input_strings, input_indices)]

        # Correctly ordered indices and strings
        correct_indices = [(0, 1), (0, 1, 4, 5), (0, 2, 4, 5, 7),
                           (0, 1, 3, 5, 6, 7), (0, 1, 3, 4, 5, 6)]
        correct_strings = ['XY', 'YYYY', 'YZZXX', 'YYZZXX', 'XXYXYZ']
        correct_phases = [1j,1,1,1,-1j]

        # Extract indices and strings generated by the Pauli operators
        indices = [pauli.qubit_indices for pauli in paulis]
        strings = [pauli.pauli_str for pauli in paulis]
        phases = [pauli.phase for pauli in paulis]
        
        # Test that indices and strings were generated correctly
        assert indices == correct_indices, f'Pauli indices not generated correctly'
        assert strings == correct_strings, f'Pauli strings not generated correctly'
        assert np.allclose(phases,correct_phases), f'Pauli phases not generated correctly'


        # Exception cases - not valid Pauli operators
        input_exc_indices = [(0, 2, 1), (1,), (2, 2, 2, 4, 5), (0, 1, 5, 3, 7, 6), (6, 5, 4, 3, 2, 1, 0)]
        input_exc_strings = ['XY', 'XZZYZXYZ', 'YXZZ', 'XYZZX', 'Z']

        # Test the exceptions are raised
        for string,index in zip(input_exc_strings, input_exc_indices):
            
            # Attempt construction of Pauli operator 
            with self.assertRaises(AssertionError) as context:
                pauli = PauliOp(string, index)

            # Check exception message
            self.assertEqual(f"Each Pauli operator must have a unique qubit index", str(context.exception))

    def test_sort_pauli_op(self):
        """
        Test of the function that sorts the string attribute and the indices
        of the pauli operator object.

        The test consists in properly ordering a set of string + indices.
        """

        # Define input examples of indices and strings
        input_indices = [(0, 2, 1), (4, 5, 0, 1), (2, 0, 7, 4, 5),
                         (0, 1, 5, 3, 7, 6), (6, 5, 4, 3, 2, 1, 0)]
        input_strings = ['XXZ', 'YYYY', 'ZYXZX', 'YYZZXX', 'ZYXYZYX']

        # Create Pauli operators from input examples
        paulis = [PauliOp(string, index)
                  for string, index in zip(input_strings, input_indices)]

        # Correctly ordered indices and strings
        correct_indices = [(0, 1, 2), (0, 1, 4, 5), (0, 2, 4, 5, 7),
                           (0, 1, 3, 5, 6, 7), (0, 1, 2, 3, 4, 5, 6)]
        correct_strings = ['XZX', 'YYYY', 'YZZXX', 'YYZZXX', 'XYZYXYZ']

        # Extract indices and strings generated by the Pauli operators
        indices = [pauli.qubit_indices for pauli in paulis]
        strings = [pauli.pauli_str for pauli in paulis]

        # Test that indices and strings were correctly ordered when creating the Pauli Operator
        assert indices == correct_indices, f'Pauli indices not sorted correctly'
        assert strings == correct_strings, f'Pauli strings not sorted correctly'


        # Exception cases - not valid Pauli operators
        input_exc_indices = [(0, 2, 1), (1,), (2, 0, 7, 4, 5), (0, 1, 5, 3, 7, 6), (6, 5, 4, 3, 2, 1, 0)]
        input_exc_strings = ['LXY', 'Q', 'ZYXZR', 'FYZZXX', 'ZYXYZYN']

        # Pauli set
        PAULIS_SET = set('XYZI')

        # Test the exceptions are raised
        for string,index in zip(input_exc_strings, input_exc_indices):
            
            # Extract exception from string
            non_pauli_operator = list(set(string) - PAULIS_SET)[0]

            # Attempt construction of Pauli operator 
            with self.assertRaises(ValueError) as context:
                pauli = PauliOp(string, index)

            # Check exception message
            self.assertEqual(f"{non_pauli_operator} is not a valid Pauli. Please choose from the set {PAULIS_SET}", str(context.exception))

    def test_pauli_simplify(self):
        """
        Test of the function that simplifies the expression of Pauli Operator by using
        the multiplication rules that Pauli matrices obey.

        The test consists in simplify a set of Pauli operators.
        """

        # Define input examples of indices and strings
        input_indices = [(0, 1, 1), (0, 0, 0, 2),
                         (0, 1, 3, 3, 4), (1, 1, 1, 1, 1, 1, 1)]
        input_strings = ['ZXY', 'YYYY', 'ZYIZX', 'ZZZZZXZ']

        # Create Pauli operators from input examples
        paulis = [PauliOp(string, index)
                  for string, index in zip(input_strings, input_indices)]

        # Indices, strings and phases from correctly simplified Pauli operators
        correct_indices = [(0, 1), (0, 2), (0, 1, 3, 4), (1,)]
        correct_strings = ['ZZ', 'YY', 'ZYZX', 'X']
        correct_phases = [1j, 1, 1, -1]

        # Extract indices, strings and phases generated by the Pauli operators
        indices = [pauli.qubit_indices for pauli in paulis]
        strings = [pauli.pauli_str for pauli in paulis]
        phases = [pauli.phase for pauli in paulis]

        # Test that indices and strings were correctly ordered when creating the Pauli Operator
        assert indices == correct_indices, f'Indices not simplified correctly'
        assert strings == correct_strings, f'Strings not simplified correctly'
        assert np.allclose(
            phases, correct_phases), f'Phases not simplified correctly'

    def test_pauli_is_trivial(self):
        """
        Test of the function that checks if a Pauli Operator object is trivial,
        i.e. only contains identity operators.

        The tests consists in checking the triviality of a set of examples.
        """

        # Define input examples of indices and strings
        input_indices = [(0, 1, 2), (0, 1, 4, 5), (0, 2, 4, 5, 7),
                         (0, 1, 3, 5, 6, 7), (0, 1, 2, 3, 4, 5, 6)]
        input_strings = ['III', 'IYYY', 'YZIXX', 'IIIIII', 'XIIIIII']

        # Create Pauli operators from input examples
        paulis = [PauliOp(string, index)
                  for string, index in zip(input_strings, input_indices)]

        # Correct list containing whether the Pauli Operator is trivial or not
        correct_triviality = [True, False, False, True, False]

        # Compute triviality of the Pauli operators
        triviality = [pauli._is_trivial for pauli in paulis]

        # Test that indices and strings were correctly ordered when creating the Pauli Operator
        assert np.allclose(
            triviality, correct_triviality), f'Triviality classification is incorrect'

    def test_pauli_matrix(self):
        """
        Test of the function that computes the matrix representation of the Pauli operator.

        The test consists in computing the matrix representation of set of examples.
        """

        # Define input examples of indices and strings
        input_indices = [(0, 1), (0, 1), (0, 1), (0, 1), (0, 1, 2)]
        input_strings = ['IX', 'XY', 'XZ', 'ZY', 'IIZ']

        # Create Pauli operators from input examples
        paulis = [PauliOp(string, index)
                  for string, index in zip(input_strings, input_indices)]

        # Correct pauli representations
        correct_matrix_representations = [np.array(([0, 1, 0, 0], [1, 0, 0, 0], [0, 0, 0, 1], [0, 0, 1, 0]), dtype=complex),
                                          np.array(
                                              ([0, 0, 0, -1j], [0, 0, 1j, 0], [0, -1j, 0, 0], [1j, 0, 0, 0]), dtype=complex),
                                          np.array(
                                              ([0, 0, 1, 0], [0, 0, 0, -1], [1, 0, 0, 0], [0, -1, 0, 0]), dtype=complex),
                                          np.array(([0, -1j, 0, 0], [1j, 0, 0, 0], [0, 0, 0, 1j],
                                                    [0, 0, -1j, 0]), dtype=complex),
                                          np.array(([1, 0, 0, 0, 0, 0, 0, 0], [0, -1, 0, 0, 0, 0, 0, 0], [0, 0, 1, 0, 0, 0, 0, 0], [0, 0, 0, -1, 0, 0, 0, 0],
                                                    [0, 0, 0, 0, 1, 0, 0, 0], [0, 0, 0, 0, 0, -1, 0, 0], [0, 0, 0, 0, 0, 0, 1, 0], [0, 0, 0, 0, 0, 0, 0, -1]), dtype=complex)]

        # Compute matrix representations
        matrix_representations = [pauli.matrix for pauli in paulis]

        # Compare each matrix representation
        matrix_comparison = [(correct_matrix_representations[i] == matrix_representations[i]).all(
        ) for i in range(len(matrix_representations))]

        # Test that the computed representations are correct
        assert np.all(matrix_comparison), f'Matrix representation is incorrect'

    def test_pauli_len(self):
        """
        Test the function that computes the length of the Pauli operator.

        The tests consists in computing the length for several examples.
        """

        # Define input examples of indices and strings
        input_indices = [(0, 1, 2), (0, 1, 4, 5), (0, 2, 4, 5, 7),
                         (0, 1, 3, 5, 6, 7), (0, 1, 2, 3, 4, 5, 6)]
        input_strings = ['III', 'IYYY', 'YZIXX', 'IIIIII', 'XIIIIII']

        # Create Pauli operators from input examples
        paulis = [PauliOp(string, index)
                  for string, index in zip(input_strings, input_indices)]

        # Correct Pauli operator lengths
        correct_lengths = [3, 4, 5, 6, 7]

        # Compute Pauli operator lengths
        lengths = [pauli.__len__() for pauli in paulis]

        # Test that computed lengths are correct
        assert np.allclose(
            lengths, correct_lengths), f'Computed lengths are incorrect'

    def test_pauli_eq(self):
        """
        Test the function that checks whether two Pauli operators are equivalent.

        The test consists in checking equivalence between different Pauli operators for several examples.
        """

        # Define input examples of indices and strings
        input_indices = [(0, 1, 2), (0, 1, 2), (0, 1, 4, 5), (0, 2, 4, 5, 7),
                         (0, 1, 3, 5, 6, 7), (0, 1, 3, 5, 6, 7), (0, 1, 2, 3, 4, 5, 6)]
        input_strings = ['XZX', 'XZX', 'YYYY',
                         'YZZXX', 'YYZZXX', 'YYZZXX', 'XYZYXYZ']

        # Create Pauli operators from input examples
        paulis = [PauliOp(string, index)
                  for string, index in zip(input_strings, input_indices)]

        # Check equivalence between contiguous Pauli operators in the list
        equivalences = [paulis[i].__eq__(paulis[i+1])
                        for i in range(len(paulis)-1)]

        # Correct equivalences
        correct_equivalences = [True, False, False, False, True, False]

        # Test the checks have been performes correctly
        assert np.allclose(
            equivalences, correct_equivalences), f'Equivalences have not been computed correctly'

    def test_pauli_copy(self):
        """
        Test the function that computes a copy of a Pauli operator.

        The test consists in creating copies for a set of examples.
        """

        # Define input examples of indices and strings
        input_indices = [(0, 1, 2), (0, 1, 4, 5), (0, 2, 4, 5, 7),
                         (0, 1, 3, 5, 6, 7), (0, 1, 2, 3, 4, 5, 6)]
        input_strings = ['XZX', 'YYYY', 'YZZXX', 'YYZZXX', 'XYZYXYZ']

        # Create Pauli operators from input examples
        paulis = [PauliOp(string, index)
                  for string, index in zip(input_strings, input_indices)]

        # Generate copies of the Pauli operators
        copy_paulis = [pauli.__copy__() for pauli in paulis]

        # Test the copies have been correctly performes

        check_copy = [paulis[j].__eq__(copy_paulis[j])
                      for j in range(len(paulis))]

        assert sum(check_copy) == len(
            check_copy), f'Copies were not performed correctly'

    def test_pauli_str(self):
        """
        Test the function that generates a string representation of the Pauli operator.

        The test consists in computing the string representation of a set of examples.
        """

        # Define input examples of indices and strings
        input_indices = [(0, 1, 2), (0, 1, 4, 5), (0, 2, 4, 5, 7),
                         (0, 1, 3, 5, 6, 7), (0, 1, 2, 3, 4, 5, 6)]
        input_strings = ['XZX', 'YYYY', 'YZZXX', 'YYZZXX', 'XYZYXYZ']

        # Create Pauli operators from input examples
        paulis = [PauliOp(string, index)
                  for string, index in zip(input_strings, input_indices)]

        # Correct string representations
        correct_str_rep = ['X_{0}Z_{1}X_{2}', 'Y_{0}Y_{1}Y_{4}Y_{5}', 'Y_{0}Z_{2}Z_{4}X_{5}X_{7}',
                           'Y_{0}Y_{1}Z_{3}Z_{5}X_{6}X_{7}', 'X_{0}Y_{1}Z_{2}Y_{3}X_{4}Y_{5}Z_{6}']

        # Generate string rtepresentations
        str_rep = [pauli.__str__() for pauli in paulis]

        # Test that string representations were correctly generated
        assert str_rep == correct_str_rep, f'String representations were not correctly generated'

    def test_pauli_repr(self):
        """
        Test the function that generates a string representation of the Pauli operator object.

        The test consists in computing the string representation of a set of examples.
        """

        # Define input examples of indices and strings
        input_indices = [(0, 1, 2), (0, 1, 4, 5), (0, 2, 4, 5, 7),
                         (0, 1, 3, 5, 6, 7), (0, 1, 2, 3, 4, 5, 6)]
        input_strings = ['XZX', 'YYYY', 'YZZXX', 'YYZZXX', 'XYZYXYZ']

        # Create Pauli operators from input examples
        paulis = [PauliOp(string, index)
                  for string, index in zip(input_strings, input_indices)]

        # Correct string representations
        correct_str_rep = ['PauliOp(XZX,(0, 1, 2))', 'PauliOp(YYYY,(0, 1, 4, 5))', 'PauliOp(YZZXX,(0, 2, 4, 5, 7))',
                           'PauliOp(YYZZXX,(0, 1, 3, 5, 6, 7))', 'PauliOp(XYZYXYZ,(0, 1, 2, 3, 4, 5, 6))']

        # Generate string rtepresentations
        str_rep = [pauli.__repr__() for pauli in paulis]

        # Test that string representations were correctly generated
        assert str_rep == correct_str_rep, f'String representations for Pauli operator object were not correctly generated'

    def test_pauli_mul(self):
        """
        Test the functions that multiply two Pauli operators.

        The test consists in multipliying a set of example Pauli operators.
        """

        # Define input examples of indices and strings
        input_indices = [(0, 1, 2), (0, 1, 4, 5), (0, 2, 4, 5, 7),
                         (0, 1, 3, 5, 6, 7), (0, 1, 2, 3, 4, 5, 6)]
        input_strings = ['XZX', 'YYYY', 'YZZXX', 'YYZZXX', 'XYZYXYX']

        # Create Pauli operators from input examples
        paulis = [PauliOp(string, index)
                  for string, index in zip(input_strings, input_indices)]

        # Compute multiplication between contiguous Pauli operators in the list
        pauli_mult = [paulis[j].__mul__(paulis[j+1])
                      for j in range(len(paulis)-1)]
        indices = [pauli.qubit_indices for pauli in pauli_mult]
        strings = [pauli.pauli_str for pauli in pauli_mult]
        phases = [pauli.phase for pauli in pauli_mult]

        # Correct matrix multiplications
        correct_indices = [(0, 1, 2, 4, 5), (1, 2, 4, 5, 7),
                           (1, 2, 3, 4, 5, 6), (0, 2, 3, 4, 5, 7)]
        correct_strings = ['ZXXYY', 'YZXZX', 'YZZZYX', 'ZZXXXX']
        correct_phases = [1, 1, -1j, 1j]

        # Test that matrix multiplications are correct
        assert indices == correct_indices, f'Indices from matrix multiplications are not correct'
        assert strings == correct_strings, f'Strings from matrix multiplications are not correct'
        assert np.allclose(
            phases, correct_phases), f'Phases from matrix multiplications are not correct'

    def test_pauli_XYZI(self):
        """
        Test the functions that define each Pauli matrix.

        The test consists in checking the properties of each Pauli matrix when defined as a Pauli
        operator object.
        """

        # Define a qubit index to define the Pauli operator object
        index = 1

        # Define the Pauli matrices
        X = PauliOp.X(index)
        Y = PauliOp.Y(index)
        Z = PauliOp.Z(index)
        I = PauliOp.I(index)
        paulis = [X, Y, Z, I]

        # Extract attributes
        strings = [pauli.pauli_str for pauli in paulis]
        indices = [pauli.qubit_indices for pauli in paulis]
        phases = [pauli.phase for pauli in paulis]
        string_rep = [pauli.__str__() for pauli in paulis]
        string_obj_rep = [pauli.__repr__() for pauli in paulis]
        matrix_rep = [pauli.matrix for pauli in paulis]

        # Correct attributes
        correct_strings = ['X', 'Y', 'Z', 'I']
        correct_indices = [(1,), (1,), (1,), (1,)]
        correct_phases = [1, 1, 1, 1]
        correct_string_rep = ['X_{1}', 'Y_{1}', 'Z_{1}', 'I_{1}']
        correct_string_obj_rep = [
            'PauliOp(X,(1,))', 'PauliOp(Y,(1,))', 'PauliOp(Z,(1,))', 'PauliOp(I,(1,))']
        correct_matrix_rep = [np.array(([0, 1], [1, 0]), dtype=complex), np.array(([0, -1j], [1j, 0]), dtype=complex),
                              np.array(([1, 0], [0, -1]), dtype=complex), np.array(([1, 0], [0, 1]), dtype=complex)]

        # Check indices, strings, phases and matrix representations
        assert strings == correct_strings, f'Strings are not correctly defined for Pauli matrices'
        assert indices == correct_indices, f'Indices are not correctly defined for Pauli matrices'
        assert np.allclose(
            phases, correct_phases), f'Phases are not correctly defined for Pauli matrices'
        assert string_rep == correct_string_rep, f'String representations are not correctly defined for Pauli matrices'
        assert string_obj_rep == correct_string_obj_rep, f'String object representations are not correctly defined for Pauli matrices'
        assert np.all([(matrix_rep[i] == correct_matrix_rep[i]).all() for i in range(len(
            matrix_rep))]), f'Matrix representations are not correctly defined for Pauli matrices'

    def test_hamiltonian_init(self):
        """
        Test initialization method of the Hamiltonian class.
        
        The tests consists in checking correct initializations for representative examples.
        """

        # Number of qubits in the register
        n_qubits = 3

        ## First example

        # Define Pauli terms using shifted indices
        idx_ref = 3
        indices_ref = list(range(idx_ref,n_qubits+idx_ref))
        trivial_indices = [(i,) for i in indices_ref]
        singlet_indices = [(i,) for i in indices_ref]
        pair_indices = [(i, j) for j in indices_ref for i in range(idx_ref,j)]

        trivial_terms = [PauliOp('I', indices) for indices in trivial_indices]
        linear_terms = [PauliOp('Z', indices) for indices in singlet_indices]
        quadratic_terms = [PauliOp('ZZ', indices) for indices in pair_indices]

        terms = trivial_terms + linear_terms + quadratic_terms

        # Define coefficients and constant
        trivial_coeffs = [0.5 for _ in range(len(trivial_indices))]
        linear_coeffs = [-1 for _ in range(len(singlet_indices))]
        quadratic_coeffs = [0.5 for _ in range(len(pair_indices))]
        coefficients = trivial_coeffs + linear_coeffs + quadratic_coeffs

        constant = 2

        # Define Hamiltonian
        hamiltonian = Hamiltonian(terms, coefficients, constant)

        # Correct attributes
        correct_indices = list(range(n_qubits))
        correct_singlet_indices = [(i,) for i in correct_indices]
        correct_pair_indices = [(i, j) for j in correct_indices for i in range(j)]
        correct_linear_terms = [PauliOp('Z', indices) for indices in correct_singlet_indices]
        correct_quadratic_terms = [PauliOp('ZZ', indices) for indices in correct_pair_indices]
        correct_linear_coeffs = [-1 for _ in range(len(correct_singlet_indices))]
        correct_quadratic_coeffs = [0.5 for _ in range(len(correct_pair_indices))]

        correct_terms = correct_linear_terms + correct_quadratic_terms
        correct_coefficients = correct_linear_coeffs + correct_quadratic_coeffs
        correct_constant = constant + sum(trivial_coeffs)

        # Test Hamiltonian was defined correctly
        assert hamiltonian.terms  == correct_terms, f'Terms in the Hamiltonian were not constructed correctly'
        assert hamiltonian.coeffs == correct_coefficients, f'Coefficients in the Hamiltonian were not constructed correctly'
        assert hamiltonian.constant == correct_constant, f'Constant in the Hamiltonian was not correctly constructed'

        ## Second example

        # Define Pauli terms
        trivial_indices = [(i,) for i in range(n_qubits)]
        pair_indices = [(i, j) for j in range(n_qubits) for i in range(j)]

        trivial_terms = [PauliOp('I', indices) for indices in trivial_indices]
        terms_strings = ['XX', 'YY']
        quadratic_terms = [PauliOp(string, indices)
                 for indices in pair_indices for string in terms_strings]

        terms = trivial_terms + quadratic_terms

        # Define coefficients and constant
        trivial_coeffs = [-1 for _ in range(len(trivial_terms))]
        quadratic_coeffs = [0.5 for _ in range(len(quadratic_terms))]
        coefficients = trivial_coeffs + quadratic_coeffs

        constant = 2
        
        # Define Hamiltonian
        hamiltonian = Hamiltonian(terms, coefficients, constant)

        # Correct attributes
        correct_terms = quadratic_terms
        correct_coefficients = quadratic_coeffs
        correct_constant = constant + sum(trivial_coeffs)

        # Test Hamiltonian was defined correctly
        assert hamiltonian.terms  == correct_terms, f'Terms in the Hamiltonian were not constructed correctly'
        assert hamiltonian.coeffs == correct_coefficients, f'Coefficients in the Hamiltonian were not constructed correctly'
        assert hamiltonian.constant == correct_constant, f'Constant in the Hamiltonian was not correctly constructed'

        ## Exception case - number of terms does not match number of coefficients
        
        # Number of qubits
        n_qubits = 5

        # Define terms and coefficients
        pair_indices = [(i, j) for j in range(n_qubits) for i in range(j)]
        terms = [PauliOp('ZZ', indices) for indices in pair_indices]
        coefficients = [1 for _ in range(len(terms)+1)]

        # Attempt construction of Hamiltonian
        with self.assertRaises(AssertionError) as context:
            hamiltonian = Hamiltonian(terms, coefficients, constant = 0)
            
        # Check exception message
        self.assertEqual("Number of Pauli terms in Hamiltonian should be same as number of coefficients", str(context.exception))

        ## Exception case - non Pauli object term present in list of terms
        
        # Number of qubits
        n_qubits = 5

        # Define terms and coefficients
        pair_indices = [(i, j) for j in range(n_qubits) for i in range(j)]
        error_term = (0,5)
        terms = [PauliOp('ZZ', indices) for indices in pair_indices] + [error_term]
        coefficients = [1 for _ in range(len(terms))]
        
        # Attempt construction of Hamiltonian
        with self.assertRaises(TypeError) as context:
            hamiltonian = Hamiltonian(terms, coefficients, constant = 0)

        # Check exception message
        self.assertEqual(f"Pauli terms should be of type PauliOp and not {type(error_term)}", str(context.exception))
        
    def test_hamiltonian_qureg(self):
        """
        Tests the function that generated the register in the Hamiltonian object.

        The test consists in generating registers from a set of examples.
        """

        # Number of qubits in the register
        n_qubits = 5

        # Define Pauli terms
        terms_indices = [(i, j) for j in range(n_qubits) for i in range(j)]
        terms_strings = ['XX', 'YY']
        terms = [PauliOp(string, indices)
                 for indices in terms_indices for string in terms_strings]

        # Define coefficients and constant
        coefficients = [1 for _ in range(len(terms))]
        constant = 2

        # Define Hamiltonian
        hamiltonian = Hamiltonian(terms, coefficients, constant)

        # Extract physical register
        register = hamiltonian.qureg

        # Correct register
        correct_register = [0, 1, 2, 3, 4]

        # Test if the extracted register is correct
        assert np.allclose(
            register, correct_register), f'Register has been incorrectly constructed'

    def test_hamiltonian_divide_into_singlets_pairs(self):
        """
        Tests the function that divides that Hamiltonian into singlets and pairs.

        The tests consists in splittig terms into singlets and pairs from a given Hamiltonian.
        """

        # Number of qubits in the register
        n_qubits = 5

        # Define Pauli terms
        singlet_indices = [(i,) for i in range(n_qubits)]
        pair_indices = [(i, j) for j in range(n_qubits) for i in range(j)]
        correct_linear_terms = [PauliOp('Z', indices)
                                for indices in singlet_indices]
        correct_quadratic_terms = [
            PauliOp('ZZ', indices) for indices in pair_indices]

        constant_term = [PauliOp('I', (0,))]

        terms = correct_linear_terms + correct_quadratic_terms + constant_term

        # Define coefficients and constant
        correct_linear_coeffs = [-1 for _ in range(len(singlet_indices))]
        correct_quadratic_coeffs = [1 for _ in range(len(pair_indices))]
        constant_coeff = [1]
        coefficients = correct_linear_coeffs + correct_quadratic_coeffs + constant_coeff

        constant = 1
        correct_constant = 2

        # Define Hamiltonian
        hamiltonian = Hamiltonian(terms, coefficients, constant)

        # Extract singlets and pairs attributes
        linear_terms = hamiltonian.qubits_singles
        quadratic_terms = hamiltonian.qubits_pairs
        linear_coeffs = hamiltonian.single_qubit_coeffs
        quadratic_coeffs = hamiltonian.pair_qubit_coeffs
        constant = hamiltonian.constant

        # Test if the hamiltonian terms and coefficients have been correctly separated
        assert linear_terms == correct_linear_terms, f'Linear terms have not been correctly generated'
        assert np.allclose(
            linear_coeffs, correct_linear_coeffs), f'Linear coefficients have not been correctly generated'
        assert quadratic_terms == correct_quadratic_terms, f'Quadratic terms have not been correctly generated'
        assert np.allclose(
            quadratic_coeffs, correct_quadratic_coeffs), f'Quadratic coefficients have not been correctly generated'
        assert np.allclose(
            constant, correct_constant), f'Constant coefficient have not been correctly generated'

        # Exception case - generating terms beyond quadratic order

        # Define hamiltonian attributes
        terms_exc = [PauliOp('Z',(0,)), PauliOp('ZZ',(0,1)), PauliOp('ZZ',(0,2)),
                    PauliOp('ZZ',(1,2)), PauliOp('ZZZ',(1,2,3))]
        coefficients_exc = [1,0.5,0.5,0.5,3]
        constant_exc = 2

        # Test the exception is raised - Attempt construction of Hamiltonian
        with self.assertRaises(NotImplementedError) as context:
            hamiltonian_exc = Hamiltonian(terms_exc,coefficients_exc,constant_exc)

        # Check exception message
        self.assertEqual("Hamiltonian only supports Linear and Quadratic terms", str(context.exception))
        
    def test_hamiltonian_str(self):
        """
        Tests the function that constructs a string representation of the Hamiltonian object.

        The test consists of constructing the string representation of a given Hamiltonian.
        """

        # Number of qubits in the register
        n_qubits = 3

        # Define Pauli terms
        singlet_indices = [(i,) for i in range(n_qubits)]
        pair_indices = [(i, j) for j in range(n_qubits) for i in range(j)]
        linear_terms = [PauliOp('Z', indices) for indices in singlet_indices]
        quadratic_terms = [PauliOp('ZZ', indices) for indices in pair_indices]

        terms = linear_terms + quadratic_terms

        # Define coefficients and constant
        linear_coeffs = [-1 for _ in range(len(singlet_indices))]
        quadratic_coeffs = [0.5 for _ in range(len(pair_indices))]
        coefficients = linear_coeffs + quadratic_coeffs

        constant = 2

        # Define Hamiltonian
        hamiltonian = Hamiltonian(terms, coefficients, constant)

        # Extract string representation
        string_rep = hamiltonian.__str__()

        # Correct string representation
        correct_string_rep = '-1*Z_{0} + -1*Z_{1} + -1*Z_{2} + 0.5*Z_{0}Z_{1} + 0.5*Z_{0}Z_{2} + 0.5*Z_{1}Z_{2} + 2'

        # Test the obtained string representation
        assert string_rep == correct_string_rep, f'Hamiltonian string representation was incorrectly constructed'

    def test_hamiltonian_len(self):
        """
        Tests the function that computes the number of terms in the Hamiltonian.

        The test consists in checking the number of terms for a fully connected Hamiltonian
        with linear terms for an increasing number of qubits
        """

        # Initialize lists containing correct and computed lengths
        correct_lengths = []
        lengths = []

        for n_qubits in range(3, 10):

            # Compute correct lengths
            correct_lengths.append(n_qubits + n_qubits*(n_qubits-1)/2)

            # Define Pauli terms
            singlet_indices = [(i,) for i in range(n_qubits)]
            pair_indices = [(i, j) for j in range(n_qubits) for i in range(j)]
            linear_terms = [PauliOp('Z', indices)
                            for indices in singlet_indices]
            quadratic_terms = [PauliOp('ZZ', indices)
                               for indices in pair_indices]

            terms = linear_terms + quadratic_terms

            # Define coefficients and constant
            linear_coeffs = [-1 for _ in range(len(singlet_indices))]
            quadratic_coeffs = [0.5 for _ in range(len(pair_indices))]
            coefficients = linear_coeffs + quadratic_coeffs

            constant = 2

            # Define Hamiltonian
            hamiltonian = Hamiltonian(terms, coefficients, constant)

            # Compute length
            length = hamiltonian.__len__()
            lengths.append(length)

        # Test that computed lengths are correct
        assert np.allclose(
            lengths, correct_lengths), f'Hamiltonian lengths are incorrectly computed'

    def test_hamiltonian_expression(self):
        """
        Tests the function that computes the symbolic expression of the Hamiltonian object.

        The test consists in computing the symbolic expression for a specific example.
        """

        # Number of qubits in the register
        n_qubits = 3

        # Define Pauli terms
        singlet_indices = [(i,) for i in range(n_qubits)]
        pair_indices = [(i, j) for j in range(n_qubits) for i in range(j)]
        linear_terms = [PauliOp('Z', indices) for indices in singlet_indices]
        quadratic_terms = [PauliOp('ZZ', indices) for indices in pair_indices]

        terms = linear_terms + quadratic_terms

        # Define coefficients and constant
        linear_coeffs = [-1 for _ in range(len(singlet_indices))]
        quadratic_coeffs = [0.5 for _ in range(len(pair_indices))]
        coefficients = linear_coeffs + quadratic_coeffs

        constant = 2

        # Define Hamiltonian
        hamiltonian = Hamiltonian(terms, coefficients, constant)

        # Extract symbolic representation
        sym_rep = hamiltonian.expression

        # Correct symbolic representation
        correct_sym_rep = Symbol(str(constant))

        for ind, coeff in enumerate(linear_coeffs):
            correct_sym_rep += Symbol(str(coeff)+'Z'+'_{'+str(ind)+'}')

        for ind, coeff in enumerate(quadratic_coeffs):
            i, j = pair_indices[ind]
            correct_sym_rep += Symbol(str(coeff) +
                                      'Z'+'_{'+str(i)+'}'+'Z'+'_{'+str(j)+'}')

        # Test if string representation was correctly generated
        assert sym_rep == correct_sym_rep, f'Symbolic representation was incorrectly generated'

    def test_hamiltonian_add(self):
        """
        Test the function that adds two Hamiltonian objects.

        The test consists in adding two given Hamiltonians with a repeated term.
        """

        # Number of qubits in the register
        n_qubits = 4

        # Define the first Hamiltonian
        pair_indices = [(i, j) for j in range(n_qubits-1) for i in range(j)]
        input_terms = [PauliOp('ZZ', indices) for indices in pair_indices]

        quadratic_coeffs = [0.5 for _ in range(len(pair_indices))]
        input_coefficients = quadratic_coeffs
        input_constant = 2

        hamiltonian = Hamiltonian(
            input_terms, input_coefficients, input_constant)

        # Define the other Hamiltonian
        singlet_indices = [(i,) for i in range(n_qubits)]
        pair_indices = [(i, j) for j in range(n_qubits) for i in range(j)]
        linear_terms = [PauliOp('Z', indices) for indices in singlet_indices]
        quadratic_terms = [PauliOp('ZZ', indices) for indices in pair_indices]
        other_input_terms = linear_terms + quadratic_terms

        linear_coeffs = [-1 for _ in range(len(singlet_indices))]
        quadratic_coeffs = [0.5 for _ in range(
            int((n_qubits-1)*(n_qubits-2)/2))] + [1 for _ in range(n_qubits-1)]
        other_input_coefficients = linear_coeffs + quadratic_coeffs

        other_input_constant = 0

        other_hamiltonian = Hamiltonian(
            other_input_terms, other_input_coefficients, other_input_constant)

        # Add hamiltonians
        hamiltonian.__add__(other_hamiltonian)

        # Exctract terms, coefficients and constants from added Hamiltonians
        terms = hamiltonian.terms
        coefficients = hamiltonian.coeffs
        constant = hamiltonian.constant

        # Create Hamiltonian dictionary - To be removed upon implementation of _simplify()
        hamiltonian_dict = {term.qubit_indices: coeff for term,
                            coeff in zip(hamiltonian.terms, hamiltonian.coeffs)}

        # Correct terms, coefficients and constants
        correct_terms = input_terms + other_input_terms
        correct_coefficients = input_coefficients + other_input_coefficients
        correct_constant = input_constant + other_input_constant

        # Create Hamiltonian dictionary for correct Hamiltonian - To be removed upon implementation of _simplify()
        correct_hamiltonian_dict = {term.qubit_indices: coeff for term, coeff in zip(
            input_terms, input_coefficients)}

        # Add terms and coefficients from the other hamiltonian
        for term, coeff in zip(other_input_terms, other_input_coefficients):

            # If term was absent add it together with coefficient
            if correct_hamiltonian_dict.get(term.qubit_indices) is None:
                correct_hamiltonian_dict.update({term.qubit_indices: coeff})

            # If term was present add coefficient to current term
            else:
                correct_hamiltonian_dict[term.qubit_indices] += coeff

        # Test if Hamiltonians were added correctly
        assert hamiltonian_dict == correct_hamiltonian_dict, f'Hamiltonian adder did not yield correct terms and/or coefficients'
        assert np.allclose(constant, correct_constant), f'Hamiltonian adder did not yield correct constant'

    def test_hamiltonian_squared(self):
        """
        Tests the function that computes the squared Hamiltonian, also a Hamiltonian object.

        The test consists in computing the Hamiltonian squared for two given small examples.
        """

        # First example

        # Number of qubits in the register
        n_qubits = 2

        # Define Pauli terms
        singlet_indices = [(i,) for i in range(n_qubits)]
        pair_indices = [(i, j) for j in range(n_qubits) for i in range(j)]
        linear_terms = [PauliOp('Z', indices) for indices in singlet_indices]
        quadratic_terms = [PauliOp('ZZ', indices) for indices in pair_indices]
        terms = linear_terms + quadratic_terms

        # Define coefficients and constant
        linear_coeffs = [-1 for _ in range(len(singlet_indices))]
        quadratic_coeffs = [0.5 for _ in range(len(pair_indices))]
        coefficients = linear_coeffs + quadratic_coeffs
        constant = 2

        # Define Hamiltonian
        hamiltonian = Hamiltonian(terms, coefficients, constant)

        # Compute squared hamiltonian and extract attributes
        hamiltonian_sq = hamiltonian.hamiltonian_squared

        terms_sq = hamiltonian_sq.terms
        coefficients_sq = hamiltonian_sq.coeffs
        constant_sq = hamiltonian_sq.constant

        # Correct Hamiltonian squared attributes
        correct_terms_sq = [PauliOp('ZZ', (0, 1)), PauliOp('Z', (1,)),
                            PauliOp('ZZ', (0, 1)), PauliOp('Z', (0,)),
                            PauliOp('Z', (1,)),  PauliOp('Z', (0,)),
                            PauliOp('Z', (0,)), PauliOp('Z', (1,)),
                            PauliOp('ZZ', (0, 1))]

        correct_coefficients_sq = [1, -0.5, 1, -0.5, -0.5, -0.5, -4, -4, 2]

        correct_constant_sq = 4 + n_qubits * \
            (-1)**2 + n_qubits*(n_qubits-1)/2*(0.5)**2

        # Test if Hamiltonian squared was correctly computed
        assert terms_sq == correct_terms_sq, f'Hamiltonian squared did not yield correct terms'
        assert np.allclose(
            coefficients_sq, correct_coefficients_sq), f'Hamiltonian squared did not yield correct coefficients'
        assert np.allclose(
            constant_sq, correct_constant_sq), f'Hamiltonian squared did not yield correct constant'

        # Second example

        # Number of qubits in the register
        n_qubits = 2

        # Define Pauli terms
        pair_indices = [(i, j) for j in range(n_qubits) for i in range(j)]
        terms_strings = ['XX', 'YY']
        terms = [PauliOp(string, indices)
                 for indices in pair_indices for string in terms_strings]

        # Define coefficients and constant
        coefficients = [1 for _ in range(len(terms))]
        constant = 2

        # Define Hamiltonian
        hamiltonian = Hamiltonian(terms, coefficients, constant)

        # Compute squared hamiltonian and extract attributes
        hamiltonian_sq = hamiltonian.hamiltonian_squared

        terms_sq = hamiltonian_sq.terms
        coefficients_sq = hamiltonian_sq.coeffs
        constant_sq = hamiltonian_sq.constant

        # Correct Hamiltonian squared attributes
        correct_terms_sq = [PauliOp('ZZ', (0, 1)), PauliOp(
            'ZZ', (0, 1)), PauliOp('XX', (0, 1)), PauliOp('YY', (0, 1))]
        correct_coefficients_sq = [-1, -1, 4, 4]
        correct_constant_sq = 6

        # Test if Hamiltonian squared was correctly computed
        assert terms_sq == correct_terms_sq, f'Hamiltonian squared did not yield correct terms'
        assert np.allclose(
            coefficients_sq, correct_coefficients_sq), f'Hamiltonian squared did not yield correct coefficients'
        assert np.allclose(
            constant_sq, correct_constant_sq), f'Hamiltonian squared did not yield correct constant'

    def test_classical_hamiltonian(self):
        """
        Tests the function that generates a classical Hamiltonian, i.e. a Hamiltonian composed only
        of Z Pauli operators.

        The tests consists in generating a Hamiltonian from a given set of graph edges and weights.
        """

        # Number of nodes in the graph, i.e. qubits in the register
        n_nodes = 5

        # Graph properties
        input_edges = [(i,) for i in range(n_nodes)] + [(i, j)
                                                        for j in range(n_nodes) for i in range(j)]
        input_weights = [-1 for _ in range(len(input_edges))]
        input_constant = 2

        # Generate classical Hamiltonian and extract attributes
        cl_hamiltonian = Hamiltonian.classical_hamiltonian(
            input_edges, input_weights, input_constant)

        terms = cl_hamiltonian.terms
        terms_indices = [term.qubit_indices for term in terms]
        coefficients = cl_hamiltonian.coeffs
        constant = cl_hamiltonian.constant

        # Test that the generated Hamiltonian match the properties of the input graph
        assert terms_indices == input_edges, f'Classical Hamiltonian terms did not match input graph edges'
        assert np.allclose(
            coefficients, input_weights), f'Hamiltonian squared did not yield correct coefficients'
        assert np.allclose(
            constant, input_constant), f'Hamiltonian squared did not yield correct constant'

        # Exception case - inserting coefficients that are not interegers or floats

        # Define hamiltonian attributes
        terms_exc = [(0,1),(0,2),(1,2)]
        constant_exc = 2

        # Define different types of coefficients containing exceptions
        coefficients_exc_types = [[-1j,1,-1j],[2,3,0.1*1j]]
        
        # Check for every coefficient type
        for coefficients_exc in coefficients_exc_types:
            
            # Test the exception is raised - Attempt construction of Hamiltonian
            with self.assertRaises(ValueError) as context:
                hamiltonian_exc = Hamiltonian.classical_hamiltonian(terms_exc,coefficients_exc,constant_exc)
            
        # Check exception message
        self.assertEqual("Classical Hamiltonians only support Integer or Float coefficients", str(context.exception))
        
        # Exception case - generating terms beyond quadratic order

        # Define hamiltonian attributes
        terms_exc = [(0,),(0,1),(0,2),(1,2),(1,2,3)]
        coefficients_exc = [1,0.5,0.5,0.5,3]
        constant_exc = 2

        # Test the exception is raised - Attempt construction of Hamiltonian
        with self.assertRaises(ValueError) as context:
            hamiltonian_exc = Hamiltonian.classical_hamiltonian(terms_exc,coefficients_exc,constant_exc)

        # Check exception message
        self.assertEqual("Hamiltonian only supports Linear and Quadratic terms", str(context.exception))

if __name__ == "__main__":
    unittest.main()