#   Copyright 2017 The OpenFermion Developers
#   Modifications copyright 2022 Zapata Computing, Inc. for compatibility reasons.
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
"""Tests  jordan_wigner.py."""
import itertools
import os
import unittest

import numpy
import sympy

from orquestra.quantum.openfermion.chem import MolecularData
from orquestra.quantum.openfermion.config import DATA_DIRECTORY
from orquestra.quantum.openfermion.hamiltonians.special_operators import number_operator
from orquestra.quantum.openfermion.ops.operators import FermionOperator, QubitOperator
from orquestra.quantum.openfermion.ops.representations import InteractionOperator
from orquestra.quantum.openfermion.testing.testing_utils import (
    random_interaction_operator,
)
from orquestra.quantum.openfermion.transforms.opconversions import (
    get_fermion_operator,
    normal_ordered,
    reverse_jordan_wigner,
)
from orquestra.quantum.openfermion.transforms.opconversions.jordan_wigner import (
    _jordan_wigner_interaction_op,
    jordan_wigner,
    jordan_wigner_one_body,
    jordan_wigner_two_body,
)
from orquestra.quantum.openfermion.transforms.repconversions import (
    get_interaction_operator,
)
from orquestra.quantum.openfermion.utils import hermitian_conjugated


class JordanWignerTransformTest(unittest.TestCase):
    def setUp(self):
        self.n_qubits = 5

    def test_bad_input(self):
        with self.assertRaises(TypeError):
            jordan_wigner(3)

    def test_transm_raise3(self):
        raising = jordan_wigner(FermionOperator(((3, 1),)))
        self.assertEqual(len(raising.terms), 2)

        correct_operators_x = ((0, "Z"), (1, "Z"), (2, "Z"), (3, "X"))
        correct_operators_y = ((0, "Z"), (1, "Z"), (2, "Z"), (3, "Y"))
        qtermx = QubitOperator(correct_operators_x, 0.5)
        qtermy = QubitOperator(correct_operators_y, -0.5j)

        self.assertEqual(raising.terms[correct_operators_x], 0.5)
        self.assertEqual(raising.terms[correct_operators_y], -0.5j)
        self.assertTrue(raising == qtermx + qtermy)

    def test_transm_hopping_sympy(self):
        coeff = sympy.Symbol("x")
        hopping = FermionOperator(((3, 1), (2, 0))) * coeff
        hopping += FermionOperator(((2, 1), (3, 0))) * coeff

        jw_hopping = jordan_wigner(hopping)
        print(jw_hopping)
        correct_operators_xx = ((2, "X"), (3, "X"))
        correct_operators_yy = ((2, "Y"), (3, "Y"))
        qtermxx = QubitOperator(correct_operators_xx, 0.5 * coeff)
        qtermyy = QubitOperator(correct_operators_yy, 0.5 * coeff)

        self.assertEqual(jw_hopping.terms[correct_operators_xx], 0.5 * coeff)
        self.assertEqual(jw_hopping.terms[correct_operators_yy], 0.5 * coeff)
        self.assertTrue(jw_hopping == qtermxx + qtermyy)

    def test_transm_raise3_sympy(self):
        coeff = sympy.Symbol("x")
        raising = jordan_wigner(coeff * FermionOperator(((3, 1),)))
        self.assertEqual(len(raising.terms), 2)

        correct_operators_x = ((0, "Z"), (1, "Z"), (2, "Z"), (3, "X"))
        correct_operators_y = ((0, "Z"), (1, "Z"), (2, "Z"), (3, "Y"))
        qtermx = QubitOperator(correct_operators_x, 0.5 * coeff)
        qtermy = QubitOperator(correct_operators_y, -0.5j * coeff)

        self.assertEqual(raising.terms[correct_operators_x], 0.5 * coeff)
        self.assertEqual(raising.terms[correct_operators_y], -0.5j * coeff)
        self.assertTrue(raising == qtermx + qtermy)

    def test_transm_raise1(self):
        raising = jordan_wigner(FermionOperator(((1, 1),)))

        correct_operators_x = ((0, "Z"), (1, "X"))
        correct_operators_y = ((0, "Z"), (1, "Y"))
        qtermx = QubitOperator(correct_operators_x, 0.5)
        qtermy = QubitOperator(correct_operators_y, -0.5j)

        self.assertEqual(raising.terms[correct_operators_x], 0.5)
        self.assertEqual(raising.terms[correct_operators_y], -0.5j)
        self.assertTrue(raising == qtermx + qtermy)

    def test_transm_lower3(self):
        lowering = jordan_wigner(FermionOperator(((3, 0),)))

        correct_operators_x = ((0, "Z"), (1, "Z"), (2, "Z"), (3, "X"))
        correct_operators_y = ((0, "Z"), (1, "Z"), (2, "Z"), (3, "Y"))
        qtermx = QubitOperator(correct_operators_x, 0.5)
        qtermy = QubitOperator(correct_operators_y, 0.5j)

        self.assertEqual(lowering.terms[correct_operators_x], 0.5)
        self.assertEqual(lowering.terms[correct_operators_y], 0.5j)
        self.assertTrue(lowering == qtermx + qtermy)

    def test_transm_lower2(self):
        lowering = jordan_wigner(FermionOperator(((2, 0),)))

        correct_operators_x = ((0, "Z"), (1, "Z"), (2, "X"))
        correct_operators_y = ((0, "Z"), (1, "Z"), (2, "Y"))
        qtermx = QubitOperator(correct_operators_x, 0.5)
        qtermy = QubitOperator(correct_operators_y, 0.5j)

        self.assertEqual(lowering.terms[correct_operators_x], 0.5)
        self.assertEqual(lowering.terms[correct_operators_y], 0.5j)
        self.assertTrue(lowering == qtermx + qtermy)

    def test_transm_lower1(self):
        lowering = jordan_wigner(FermionOperator(((1, 0),)))

        correct_operators_x = ((0, "Z"), (1, "X"))
        correct_operators_y = ((0, "Z"), (1, "Y"))
        qtermx = QubitOperator(correct_operators_x, 0.5)
        qtermy = QubitOperator(correct_operators_y, 0.5j)

        self.assertEqual(lowering.terms[correct_operators_x], 0.5)
        self.assertEqual(lowering.terms[correct_operators_y], 0.5j)
        self.assertTrue(lowering == qtermx + qtermy)

    def test_transm_lower0(self):
        lowering = jordan_wigner(FermionOperator(((0, 0),)))

        correct_operators_x = ((0, "X"),)
        correct_operators_y = ((0, "Y"),)
        qtermx = QubitOperator(correct_operators_x, 0.5)
        qtermy = QubitOperator(correct_operators_y, 0.5j)

        self.assertEqual(lowering.terms[correct_operators_x], 0.5)
        self.assertEqual(lowering.terms[correct_operators_y], 0.5j)
        self.assertTrue(lowering == qtermx + qtermy)

    def test_transm_raise3lower0(self):
        # recall that creation gets -1j on Y and annihilation gets +1j on Y.
        term = jordan_wigner(FermionOperator(((3, 1), (0, 0))))
        self.assertEqual(
            term.terms[((0, "X"), (1, "Z"), (2, "Z"), (3, "Y"))], 0.25 * 1 * -1j
        )
        self.assertEqual(
            term.terms[((0, "Y"), (1, "Z"), (2, "Z"), (3, "Y"))], 0.25 * 1j * -1j
        )
        self.assertEqual(
            term.terms[((0, "Y"), (1, "Z"), (2, "Z"), (3, "X"))], 0.25 * 1j * 1
        )
        self.assertEqual(
            term.terms[((0, "X"), (1, "Z"), (2, "Z"), (3, "X"))], 0.25 * 1 * 1
        )

    def test_transm_number(self):
        n = number_operator(self.n_qubits, 3)
        n_jw = jordan_wigner(n)
        self.assertEqual(n_jw.terms[((3, "Z"),)], -0.5)
        self.assertEqual(n_jw.terms[()], 0.5)
        self.assertEqual(len(n_jw.terms), 2)

    def test_ccr_offsite_even_ca(self):
        c2 = FermionOperator(((2, 1),))
        a4 = FermionOperator(((4, 0),))

        self.assertTrue(normal_ordered(c2 * a4) == normal_ordered(-a4 * c2))
        self.assertTrue(jordan_wigner(c2 * a4) == jordan_wigner(-a4 * c2))

    def test_ccr_offsite_odd_ca(self):
        c1 = FermionOperator(((1, 1),))
        a4 = FermionOperator(((4, 0),))
        self.assertTrue(normal_ordered(c1 * a4) == normal_ordered(-a4 * c1))

        self.assertTrue(jordan_wigner(c1 * a4) == jordan_wigner(-a4 * c1))

    def test_ccr_offsite_even_cc(self):
        c2 = FermionOperator(((2, 1),))
        c4 = FermionOperator(((4, 1),))
        self.assertTrue(normal_ordered(c2 * c4) == normal_ordered(-c4 * c2))

        self.assertTrue(jordan_wigner(c2 * c4) == jordan_wigner(-c4 * c2))

    def test_ccr_offsite_odd_cc(self):
        c1 = FermionOperator(((1, 1),))
        c4 = FermionOperator(((4, 1),))
        self.assertTrue(normal_ordered(c1 * c4) == normal_ordered(-c4 * c1))

        self.assertTrue(jordan_wigner(c1 * c4) == jordan_wigner(-c4 * c1))

    def test_ccr_offsite_even_aa(self):
        a2 = FermionOperator(((2, 0),))
        a4 = FermionOperator(((4, 0),))
        self.assertTrue(normal_ordered(a2 * a4) == normal_ordered(-a4 * a2))

        self.assertTrue(jordan_wigner(a2 * a4) == jordan_wigner(-a4 * a2))

    def test_ccr_offsite_odd_aa(self):
        a1 = FermionOperator(((1, 0),))
        a4 = FermionOperator(((4, 0),))
        self.assertTrue(normal_ordered(a1 * a4) == normal_ordered(-a4 * a1))

        self.assertTrue(jordan_wigner(a1 * a4) == jordan_wigner(-a4 * a1))

    def test_ccr_onsite(self):
        c1 = FermionOperator(((1, 1),))
        a1 = hermitian_conjugated(c1)
        self.assertTrue(
            normal_ordered(c1 * a1) == FermionOperator(()) - normal_ordered(a1 * c1)
        )
        self.assertTrue(
            jordan_wigner(c1 * a1) == QubitOperator(()) - jordan_wigner(a1 * c1)
        )

    def test_jordan_wigner_transm_op(self):
        n = number_operator(self.n_qubits)
        n_jw = jordan_wigner(n)
        self.assertEqual(self.n_qubits + 1, len(n_jw.terms))
        self.assertEqual(self.n_qubits / 2.0, n_jw.terms[()])
        for qubit in range(self.n_qubits):
            operators = ((qubit, "Z"),)
            self.assertEqual(n_jw.terms[operators], -0.5)


class InteractionOperatorsJWTest(unittest.TestCase):
    def setUp(self):
        self.n_qubits = 5
        self.constant = 0.0
        self.one_body = numpy.zeros((self.n_qubits, self.n_qubits), float)
        self.two_body = numpy.zeros(
            (self.n_qubits, self.n_qubits, self.n_qubits, self.n_qubits), float
        )
        self.interaction_operator = InteractionOperator(
            self.constant, self.one_body, self.two_body
        )

    def test_consistency(self):
        """Test consistency with JW for FermionOperators."""
        # Random interaction operator
        n_qubits = 5
        iop = random_interaction_operator(n_qubits, real=False)
        op1 = jordan_wigner(iop)
        op2 = jordan_wigner(get_fermion_operator(iop))

        self.assertEqual(op1, op2)

        # Interaction operator from molecule
        geometry = [("Li", (0.0, 0.0, 0.0)), ("H", (0.0, 0.0, 1.45))]
        basis = "sto-3g"
        multiplicity = 1

        filename = os.path.join(DATA_DIRECTORY, "H1-Li1_sto-3g_singlet_1.45")
        molecule = MolecularData(geometry, basis, multiplicity, filename=filename)
        molecule.load()

        iop = molecule.get_molecular_hamiltonian()
        op1 = jordan_wigner(iop)
        op2 = jordan_wigner(get_fermion_operator(iop))

        self.assertEqual(op1, op2)

    def testjordan_wigner_one_body(self):
        # Make sure it agrees with jordan_wigner(FermionTerm).
        for p in range(self.n_qubits):
            for q in range(self.n_qubits):
                coefficient = numpy.random.randn()
                if p != q:
                    coefficient += 1j * numpy.random.randn()
                # Get test qubit operator.
                test_operator = jordan_wigner_one_body(p, q, coefficient)

                # Get correct qubit operator.
                fermion_term = FermionOperator(((p, 1), (q, 0)), coefficient)
                if p != q:
                    fermion_term += FermionOperator(
                        ((q, 1), (p, 0)), coefficient.conjugate()
                    )
                correct_op = jordan_wigner(fermion_term)

                self.assertTrue(test_operator == correct_op)

    def testjordan_wigner_two_body(self):
        # Make sure it agrees with jordan_wigner(FermionTerm).
        for p, q, r, s in itertools.product(range(self.n_qubits), repeat=4):
            coefficient = numpy.random.randn()
            if set([p, q]) != set([r, s]):
                coefficient += 1j * numpy.random.randn()

            # Get test qubit operator.
            test_operator = jordan_wigner_two_body(p, q, r, s, coefficient)

            # Get correct qubit operator.
            fermion_term = FermionOperator(
                ((p, 1), (q, 1), (r, 0), (s, 0)), coefficient
            )
            if set([p, q]) != set([r, s]):
                fermion_term += FermionOperator(
                    ((s, 1), (r, 1), (q, 0), (p, 0)), coefficient.conjugate()
                )
            correct_op = jordan_wigner(fermion_term)

            self.assertTrue(
                test_operator == correct_op, str(test_operator - correct_op)
            )

    def test_jordan_wigner_twobody_interaction_op_allunique(self):
        test_op = FermionOperator("1^ 2^ 3 4")
        test_op += hermitian_conjugated(test_op)

        retransformed_test_op = reverse_jordan_wigner(
            jordan_wigner(get_interaction_operator(test_op))
        )

        self.assertTrue(
            normal_ordered(retransformed_test_op) == normal_ordered(test_op)
        )

    def test_jordan_wigner_twobody_interaction_op_reversal_symmetric(self):
        test_op = FermionOperator("1^ 2^ 2 1")
        test_op += hermitian_conjugated(test_op)
        self.assertTrue(
            jordan_wigner(test_op) == jordan_wigner(get_interaction_operator(test_op))
        )

    def test_jordan_wigner_interaction_op_too_few_n_qubits(self):
        with self.assertRaises(ValueError):
            _jordan_wigner_interaction_op(self.interaction_operator, self.n_qubits - 2)

    def test_jordan_wigner_interaction_op_with_zero_term(self):
        test_op = FermionOperator("1^ 2^ 3 4")
        test_op += hermitian_conjugated(test_op)

        interaction_op = get_interaction_operator(test_op)
        interaction_op.constant = 0.0

        retransformed_test_op = reverse_jordan_wigner(jordan_wigner(interaction_op))

        self.assertEqual(normal_ordered(retransformed_test_op), normal_ordered(test_op))


class GetInteractionOperatorTest(unittest.TestCase):
    def setUp(self):
        self.n_qubits = 5
        self.constant = 0.0
        self.one_body = numpy.zeros((self.n_qubits, self.n_qubits), float)
        self.two_body = numpy.zeros(
            (self.n_qubits, self.n_qubits, self.n_qubits, self.n_qubits), float
        )

    def test_get_interaction_operator_identity(self):
        interaction_operator = InteractionOperator(-2j, self.one_body, self.two_body)
        qubit_operator = jordan_wigner(interaction_operator)
        self.assertTrue(qubit_operator == -2j * QubitOperator(()))
        self.assertEqual(
            interaction_operator,
            get_interaction_operator(
                reverse_jordan_wigner(qubit_operator), self.n_qubits
            ),
        )

    def test_get_interaction_operator_one_body(self):
        interaction_operator = get_interaction_operator(
            FermionOperator("2^ 2"), self.n_qubits
        )
        one_body = numpy.zeros((self.n_qubits, self.n_qubits), float)
        one_body[2, 2] = 1.0
        self.assertEqual(
            interaction_operator, InteractionOperator(0.0, one_body, self.two_body)
        )

    def test_get_interaction_operator_one_body_twoterm(self):
        interaction_operator = get_interaction_operator(
            FermionOperator("2^ 3", -2j) + FermionOperator("3^ 2", 3j), self.n_qubits
        )
        one_body = numpy.zeros((self.n_qubits, self.n_qubits), complex)
        one_body[2, 3] = -2j
        one_body[3, 2] = 3j
        self.assertEqual(
            interaction_operator, InteractionOperator(0.0, one_body, self.two_body)
        )

    def test_get_interaction_operator_two_body(self):
        interaction_operator = get_interaction_operator(
            FermionOperator("2^ 2 3^ 4"), self.n_qubits
        )
        two_body = numpy.zeros(
            (self.n_qubits, self.n_qubits, self.n_qubits, self.n_qubits), float
        )
        two_body[3, 2, 4, 2] = -1.0
        self.assertEqual(
            interaction_operator, InteractionOperator(0.0, self.one_body, two_body)
        )

    def test_get_interaction_operator_two_body_distinct(self):
        interaction_operator = get_interaction_operator(
            FermionOperator("0^ 1^ 2 3"), self.n_qubits
        )
        two_body = numpy.zeros(
            (self.n_qubits, self.n_qubits, self.n_qubits, self.n_qubits), float
        )
        two_body[1, 0, 3, 2] = 1.0
        self.assertEqual(
            interaction_operator, InteractionOperator(0.0, self.one_body, two_body)
        )
