# This code is part of Qiskit.
#
# (C) Copyright IBM 2018, 2020.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

" Test MatrixExpectation"

import unittest
from test.aqua import QiskitAquaTestCase

import itertools
import numpy as np

from qiskit.aqua import QuantumInstance
from qiskit.aqua.operators import (X, Y, Z, I, CX, H, S,
                                   ListOp, Zero, One, Plus, Minus, StateFn,
                                   MatrixExpectation, CircuitSampler)
from qiskit import BasicAer


# pylint: disable=invalid-name

class TestMatrixExpectation(QiskitAquaTestCase):
    """Pauli Change of Basis Expectation tests."""

    def setUp(self) -> None:
        super().setUp()
        self.seed = 97
        backend = BasicAer.get_backend('statevector_simulator')
        q_instance = QuantumInstance(backend, seed_simulator=self.seed, seed_transpiler=self.seed)
        self.sampler = CircuitSampler(q_instance, attach_results=True)
        self.expect = MatrixExpectation()

    def test_pauli_expect_pair(self):
        """ pauli expect pair test """
        op = (Z ^ Z)
        # wf = (Pl^Pl) + (Ze^Ze)
        wf = CX @ (H ^ I) @ Zero

        converted_meas = self.expect.convert(~StateFn(op) @ wf)
        self.assertAlmostEqual(converted_meas.eval(), 0, delta=.1)
        sampled = self.sampler.convert(converted_meas)
        self.assertAlmostEqual(sampled.eval(), 0, delta=.1)

    def test_pauli_expect_single(self):
        """ pauli expect single test """
        paulis = [Z, X, Y, I]
        states = [Zero, One, Plus, Minus, S @ Plus, S @ Minus]
        for pauli, state in itertools.product(paulis, states):
            converted_meas = self.expect.convert(~StateFn(pauli) @ state)
            matmulmean = state.adjoint().to_matrix() @ pauli.to_matrix() @ state.to_matrix()
            self.assertAlmostEqual(converted_meas.eval(), matmulmean, delta=.1)

            sampled = self.sampler.convert(converted_meas)
            self.assertAlmostEqual(sampled.eval(), matmulmean, delta=.1)

    def test_pauli_expect_op_vector(self):
        """ pauli expect op vector test """
        paulis_op = ListOp([X, Y, Z, I])
        converted_meas = self.expect.convert(~StateFn(paulis_op))

        plus_mean = (converted_meas @ Plus)
        np.testing.assert_array_almost_equal(plus_mean.eval(), [1, 0, 0, 1], decimal=1)
        sampled_plus = self.sampler.convert(plus_mean)
        np.testing.assert_array_almost_equal(sampled_plus.eval(), [1, 0, 0, 1], decimal=1)

        minus_mean = (converted_meas @ Minus)
        np.testing.assert_array_almost_equal(minus_mean.eval(), [-1, 0, 0, 1], decimal=1)
        sampled_minus = self.sampler.convert(minus_mean)
        np.testing.assert_array_almost_equal(sampled_minus.eval(), [-1, 0, 0, 1], decimal=1)

        zero_mean = (converted_meas @ Zero)
        np.testing.assert_array_almost_equal(zero_mean.eval(), [0, 0, 1, 1], decimal=1)
        sampled_zero = self.sampler.convert(zero_mean)
        np.testing.assert_array_almost_equal(sampled_zero.eval(), [0, 0, 1, 1], decimal=1)

        sum_zero = (Plus + Minus) * (.5 ** .5)
        sum_zero_mean = (converted_meas @ sum_zero)
        np.testing.assert_array_almost_equal(sum_zero_mean.eval(), [0, 0, 1, 1], decimal=1)
        sampled_zero = self.sampler.convert(sum_zero)
        np.testing.assert_array_almost_equal((converted_meas @ sampled_zero).eval(), [0, 0, 1, 1],
                                             decimal=1)

        for i, op in enumerate(paulis_op.oplist):
            mat_op = op.to_matrix()
            np.testing.assert_array_almost_equal(zero_mean.eval()[i],
                                                 Zero.adjoint().to_matrix() @
                                                 mat_op @ Zero.to_matrix(),
                                                 decimal=1)
            np.testing.assert_array_almost_equal(plus_mean.eval()[i],
                                                 Plus.adjoint().to_matrix() @
                                                 mat_op @ Plus.to_matrix(),
                                                 decimal=1)
            np.testing.assert_array_almost_equal(minus_mean.eval()[i],
                                                 Minus.adjoint().to_matrix() @
                                                 mat_op @ Minus.to_matrix(),
                                                 decimal=1)

    def test_pauli_expect_state_vector(self):
        """ pauli expect state vector test """
        states_op = ListOp([One, Zero, Plus, Minus])

        paulis_op = X
        converted_meas = self.expect.convert(~StateFn(paulis_op) @ states_op)
        np.testing.assert_array_almost_equal(converted_meas.eval(), [0, 0, 1, -1], decimal=1)

        sampled = self.sampler.convert(converted_meas)
        np.testing.assert_array_almost_equal(sampled.eval(), [0, 0, 1, -1], decimal=1)

        # Small test to see if execution results are accessible
        for composed_op in sampled:
            self.assertIn('statevector', composed_op[1].execution_results)

    def test_pauli_expect_op_vector_state_vector(self):
        """ pauli expect op vector state vector test """
        paulis_op = ListOp([X, Y, Z, I])
        states_op = ListOp([One, Zero, Plus, Minus])

        valids = [[+0, 0, 1, -1],
                  [+0, 0, 0, 0],
                  [-1, 1, 0, -0],
                  [+1, 1, 1, 1]]
        converted_meas = self.expect.convert(~StateFn(paulis_op))
        np.testing.assert_array_almost_equal((converted_meas @ states_op).eval(), valids, decimal=1)

        sampled = self.sampler.convert(states_op)
        np.testing.assert_array_almost_equal((converted_meas @ sampled).eval(), valids, decimal=1)

    def test_multi_representation_ops(self):
        """ Test observables with mixed representations """
        mixed_ops = ListOp([X.to_matrix_op(),
                            H,
                            H + I,
                            X])
        converted_meas = self.expect.convert(~StateFn(mixed_ops))

        plus_mean = (converted_meas @ Plus)
        sampled_plus = self.sampler.convert(plus_mean)
        np.testing.assert_array_almost_equal(sampled_plus.eval(),
                                             [1, .5**.5, (1 + .5**.5), 1],
                                             decimal=1)


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