# This code is part of Qiskit.
#
# (C) Copyright IBM 2018, 2019.
#
# 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.
"""
Integration Tests for Parameterized Qobj execution, testing qasm_simulator,
statevector_simulator, and expectation value snapshots.
"""

import unittest
from math import pi
import numpy as np

from test.terra import common

from qiskit.compiler import transpile
from qiskit.circuit import QuantumCircuit, Parameter
from test.terra.reference.ref_save_expval import (
    save_expval_circuits,
    save_expval_counts,
    save_expval_labels,
    save_expval_pre_meas_values,
    save_expval_circuit_parameterized,
    save_expval_final_statevecs,
)
from qiskit_aer.library import SaveStatevector
from qiskit_aer import AerSimulator, AerError


class TestParameterizedCircuit(common.QiskitAerTestCase):
    """Parameterized circuit extension tests"""

    BACKEND_OPTS = {"seed_simulator": 2113}

    def test_run_path(self):
        """Test parameterized circuit path via backed.run()"""
        shots = 1000
        backend = AerSimulator()
        circuit = QuantumCircuit(2)
        theta = Parameter("theta")
        circuit.rx(theta, 0)
        circuit.cx(0, 1)
        circuit.measure_all()
        parameter_binds = [{theta: [0, pi, 2 * pi]}]
        res = backend.run(circuit, shots=shots, parameter_binds=parameter_binds).result()
        counts = res.get_counts()
        self.assertEqual(counts, [{"00": shots}, {"11": shots}, {"00": shots}])

    def test_run_path_already_bound_parameter_expression(self):
        """Test parameterizations with a parameter expression that's already bound."""
        shots = 1000
        backend = AerSimulator()
        circuit = QuantumCircuit(2)
        tmp = Parameter("x")
        theta = Parameter("theta")
        expr = tmp - tmp
        bound_expr = expr.bind({tmp: 1})
        circuit.rx(theta, 0)
        circuit.rx(bound_expr, 0)
        circuit.cx(0, 1)
        circuit.measure_all()
        parameter_binds = [{theta: [0, pi, 2 * pi]}]
        res = backend.run(circuit, shots=shots, parameter_binds=parameter_binds).result()
        counts = res.get_counts()
        self.assertEqual(counts, [{"00": shots}, {"11": shots}, {"00": shots}])

    def test_run_path_already_transpiled_parameter_expression(self):
        """Test parameterizations with a transpiled parameter expression."""
        shots = 1000
        backend = AerSimulator()
        circuit = QuantumCircuit(1)
        theta = Parameter("theta")
        circuit.rx(theta, 0)
        circuit.measure_all()
        parameter_binds = [{theta: [0, pi, 2 * pi]}]
        tqc = transpile(circuit, basis_gates=["u3"])
        res = backend.run(tqc, shots=shots, parameter_binds=parameter_binds).result()
        counts = res.get_counts()
        self.assertEqual(counts, [{"0": shots}, {"1": shots}, {"0": shots}])

    def test_run_path_with_expressions(self):
        """Test parameterized circuit path via backed.run()"""
        shots = 1000
        backend = AerSimulator()
        circuit = QuantumCircuit(2)
        theta = Parameter("theta")
        theta_squared = theta * theta
        circuit.rx(theta, 0)
        circuit.cx(0, 1)
        circuit.rz(theta_squared, 1)
        circuit.measure_all()
        parameter_binds = [{theta: [0, pi, 2 * pi]}]
        res = backend.run(circuit, shots=shots, parameter_binds=parameter_binds).result()
        counts = res.get_counts()
        self.assertEqual(counts, [{"00": shots}, {"11": shots}, {"00": shots}])

    def test_run_path_with_expressions_multiple_params_per_instruction(self):
        """Test parameterized circuit path via backed.run()"""
        shots = 1000
        backend = AerSimulator()
        circuit = QuantumCircuit(2)
        theta = Parameter("theta")
        theta_squared = theta * theta
        circuit.rx(theta, 0)
        circuit.cx(0, 1)
        circuit.rz(theta_squared, 1)
        circuit.u(theta, theta_squared, theta, 1)
        circuit.measure_all()
        parameter_binds = [{theta: [0, pi, 2 * pi]}]
        res = backend.run(circuit, shots=shots, parameter_binds=parameter_binds).result()
        counts = res.get_counts()
        self.assertEqual(counts, [{"00": shots}, {"01": shots}, {"00": shots}])

    def test_run_path_with_more_params_than_expressions(self):
        """Test parameterized circuit path via backed.run()"""
        shots = 2000
        backend = AerSimulator()
        circuit = QuantumCircuit(2)
        theta = Parameter("theta")
        theta_squared = theta * theta
        phi = Parameter("phi")
        circuit.rx(theta, 0)
        circuit.cx(0, 1)
        circuit.rz(theta_squared, 1)
        circuit.ry(phi, 1)
        circuit.measure_all()
        parameter_binds = [{theta: [0, pi, 2 * pi], phi: [0, 1, pi]}]
        res = backend.run(
            circuit, shots=shots, parameter_binds=parameter_binds, **self.BACKEND_OPTS
        ).result()
        counts = res.get_counts()
        for index, expected in enumerate(
            [{"00": shots}, {"01": 0.25 * shots, "11": 0.75 * shots}, {"10": shots}]
        ):
            self.assertDictAlmostEqual(counts[index], expected, delta=0.05 * shots)

    def test_run_path_multiple_circuits(self):
        """Test parameterized circuit path via backed.run()"""
        shots = 1000
        backend = AerSimulator()
        circuit = QuantumCircuit(2)
        theta = Parameter("theta")
        circuit.rx(theta, 0)
        circuit.cx(0, 1)
        circuit.measure_all()
        parameter_binds = [{theta: [0, pi, 2 * pi]}] * 3
        res = backend.run([circuit] * 3, shots=shots, parameter_binds=parameter_binds).result()
        counts = res.get_counts()
        self.assertEqual(counts, [{"00": shots}, {"11": shots}, {"00": shots}] * 3)

    def test_run_path_multiple_different_circuits(self):
        """Test parameterized circuit path via backed.run()"""
        shots = 1000
        backend = AerSimulator()

        circuit1 = QuantumCircuit(2)
        theta1 = Parameter("theta1")
        circuit1.rx(theta1, 0)
        circuit1.cx(0, 1)
        circuit1.measure_all()

        circuit2 = QuantumCircuit(2)
        theta2 = Parameter("theta2")
        circuit2.rx(theta2, 0)
        circuit2.cx(0, 1)
        circuit2.measure_all()

        circuit3 = QuantumCircuit(2)
        theta3_1 = Parameter("theta3_1")
        theta3_2 = Parameter("theta3_2")
        circuit3.rx(theta3_1, 0)
        circuit3.rx(theta3_2, 0)
        circuit3.cx(0, 1)
        circuit3.measure_all()

        parameter_binds = [
            {theta1: [0, pi, 2 * pi]},
            {theta2: [0, pi, 2 * pi]},
            {theta3_1: [0, pi / 2, pi], theta3_2: [0, pi / 2, pi]},
        ]
        res = backend.run(
            [circuit1, circuit2, circuit3], shots=shots, parameter_binds=parameter_binds
        ).result()
        counts = res.get_counts()
        self.assertEqual(counts, [{"00": shots}, {"11": shots}, {"00": shots}] * 3)

    def test_run_path_with_expressions_multiple_circuits(self):
        """Test parameterized circuit path via backed.run()"""
        shots = 1000
        backend = AerSimulator()
        circuit = QuantumCircuit(2)
        theta = Parameter("theta")
        theta_squared = theta * theta
        circuit.rx(theta, 0)
        circuit.cx(0, 1)
        circuit.rz(theta_squared, 1)
        circuit.measure_all()
        parameter_binds = [{theta: [0, pi, 2 * pi]}] * 3
        res = backend.run([circuit] * 3, shots=shots, parameter_binds=parameter_binds).result()
        counts = res.get_counts()
        self.assertEqual(counts, [{"00": shots}, {"11": shots}, {"00": shots}] * 3)

    def test_run_path_with_expressions_multiple_params_per_instruction(self):
        """Test parameterized circuit path via backed.run()"""
        shots = 1000
        backend = AerSimulator()
        circuit = QuantumCircuit(2)
        theta = Parameter("theta")
        theta_squared = theta * theta
        circuit.rx(theta, 0)
        circuit.cx(0, 1)
        circuit.rz(theta_squared, 1)
        circuit.u(theta, theta_squared, theta, 1)
        circuit.measure_all()
        parameter_binds = [{theta: [0, pi, 2 * pi]}] * 3
        res = backend.run([circuit] * 3, shots=shots, parameter_binds=parameter_binds).result()
        counts = res.get_counts()
        self.assertEqual(counts, [{"00": shots}, {"01": shots}, {"00": shots}] * 3)

    def test_run_path_with_more_params_than_expressions_multiple_circuits(self):
        """Test parameterized circuit path via backed.run()"""
        shots = 2000
        backend = AerSimulator()
        circuit = QuantumCircuit(2)
        theta = Parameter("theta")
        theta_squared = theta * theta
        phi = Parameter("phi")
        circuit.rx(theta, 0)
        circuit.cx(0, 1)
        circuit.rz(theta_squared, 1)
        circuit.ry(phi, 1)
        circuit.measure_all()
        parameter_binds = [{theta: [0, pi, 2 * pi], phi: [0, 1, pi]}] * 3
        res = backend.run(
            [circuit] * 3, shots=shots, parameter_binds=parameter_binds, **self.BACKEND_OPTS
        ).result()
        counts = res.get_counts()
        for index, expected in enumerate(
            [{"00": shots}, {"01": 0.25 * shots, "11": 0.75 * shots}, {"10": shots}] * 3
        ):
            self.assertDictAlmostEqual(counts[index], expected, delta=0.05 * shots)

    def test_run_path_multiple_circuits_mismatch_length(self):
        """Test parameterized circuit path via backed.run()"""
        shots = 1000
        backend = AerSimulator()
        circuit = QuantumCircuit(2)
        theta = Parameter("theta")
        circuit.rx(theta, 0)
        circuit.cx(0, 1)
        circuit.measure_all()
        parameter_binds = [{theta: [0, pi, 2 * pi]}]
        with self.assertRaises(AerError):
            backend.run([circuit] * 3, shots=shots, parameter_binds=[parameter_binds]).result()

    def test_run_path_with_truncation(self):
        """Test parameterized circuits with truncation"""
        backend = AerSimulator(method="statevector")
        theta = Parameter("theta")
        circuit = QuantumCircuit(5, 2)
        for q in range(5):
            circuit.ry(theta, q)
        circuit.cx(0, 1)
        circuit.cx(1, 2)
        for q in range(5):
            circuit.ry(theta, q)
        circuit.cx(0, 1)
        circuit.cx(1, 2)
        circuit.append(SaveStatevector(3, label="sv", pershot=False, conditional=False), range(3))

        param_map = {theta: [0.1 * i for i in range(3)]}
        param_sets = [{theta: 0.1 * i} for i in range(3)]

        resolved_circuits = [circuit.assign_parameters(param_set) for param_set in param_sets]

        result = backend.run(circuit, parameter_binds=[param_map]).result()
        self.assertSuccess(result)

        result_without_parameters = backend.run(resolved_circuits).result()
        self.assertSuccess(result_without_parameters)

        for actual_result in result.results:
            metadata = actual_result.metadata
            self.assertEqual(metadata["active_input_qubits"], [q for q in range(3)])
        for i in range(3):
            self.assertEqual(result.data(i)["sv"], result_without_parameters.data(i)["sv"])

    def test_different_seed(self):
        """Test parameterized circuits have different seeds"""
        shots = 1000
        backend = AerSimulator()
        circuit = QuantumCircuit(2)
        theta = Parameter("theta")
        circuit.rx(theta, 0)
        circuit.cx(0, 1)
        circuit.measure_all()
        parameter_binds = [{theta: [0, pi, 2 * pi]}]
        res = backend.run(circuit, shots=shots, parameter_binds=parameter_binds).result()
        seed_simulator_list = [result.seed_simulator for result in res.results]
        self.assertEqual(len(seed_simulator_list), len(np.unique(seed_simulator_list)))

        res2 = backend.run(
            circuit,
            shots=shots,
            parameter_binds=parameter_binds,
            seed_simulator=seed_simulator_list[0],
        ).result()
        self.assertEqual(seed_simulator_list, [result.seed_simulator for result in res2.results])

    def test_run_empty(self):
        """Test parameterized circuit with empty dict path via backed.run()"""
        shots = 1000
        backend = AerSimulator()
        circuit = QuantumCircuit(2)
        theta = Parameter("theta")
        circuit.rx(theta, 0)
        circuit.cx(0, 1)
        circuit.measure_all()
        parameter_binds = [{}]
        with self.assertRaises(AerError):
            res = backend.run(circuit, shots=shots, parameter_binds=parameter_binds).result()

    def test_parameters_with_barrier(self):
        """Test parameterized circuit path with barrier"""
        backend = AerSimulator()
        circuit = QuantumCircuit(3)
        theta = Parameter("theta")
        phi = Parameter("phi")
        circuit.rx(theta, 0)
        circuit.rx(theta, 1)
        circuit.rx(theta, 2)
        circuit.barrier()
        circuit.rx(phi, 0)
        circuit.rx(phi, 1)
        circuit.rx(phi, 2)
        circuit.barrier()
        circuit.measure_all()

        parameter_binds = [{theta: [pi / 2], phi: [pi / 2]}]
        res = backend.run([circuit], shots=1024, parameter_binds=parameter_binds).result()

        self.assertSuccess(res)
        self.assertEqual(res.get_counts(), {"111": 1024})

    def test_parameters_with_conditional(self):
        """Test parameterized circuit path with conditional"""
        backend = AerSimulator()
        circuit = QuantumCircuit(3, 3)
        theta = Parameter("theta")
        phi = Parameter("phi")
        with circuit.if_test((1, False)):
            circuit.rx(theta, 0)
        with circuit.if_test((2, False)):
            circuit.rx(theta, 1)
        with circuit.if_test((0, False)):
            circuit.rx(theta, 2)
        circuit.rx(phi, 0)
        circuit.rx(phi, 1)
        circuit.rx(phi, 2)
        circuit.measure_all()

        parameter_binds = [{theta: [pi / 2], phi: [pi / 2]}]
        res = backend.run([circuit], shots=1024, parameter_binds=parameter_binds).result()

        self.assertSuccess(res)
        self.assertEqual(res.get_counts(), {"111 000": 1024})

    def test_check_parameter_binds_exist(self):
        """Test parameter_binds exists to simulate parameterized circuits"""

        shots = 1000
        backend = AerSimulator()
        circuit = QuantumCircuit(2)
        theta = Parameter("theta")
        circuit.rx(theta, 0)
        circuit.cx(0, 1)
        circuit.measure_all()
        with self.assertRaises(AerError):
            res = backend.run(circuit, shots=shots).result()

    def test_global_phase_parameters(self):
        """Test parameterized global phase"""
        backend = AerSimulator(method="extended_stabilizer", basis_gates=["h", "x", "u1"])

        theta = Parameter("theta")
        circ = QuantumCircuit(2)
        circ.ry(theta, 0)
        circ.ry(theta, 1)
        circ.measure_all()

        circ = transpile(circ, backend)

        parameter_binds = [{theta: [1, 2, 3]}]
        res = backend.run(
            [circ], shots=10, parameter_binds=parameter_binds, seed_simulator=100
        ).result()

        self.assertSuccess(res)

        circs = []
        for v in [1, 2, 3]:
            circs.append(circ.assign_parameters({theta: v}))

        expected = backend.run(circs, shots=10, seed_simulator=100).result()

        self.assertEqual(res.get_counts(), expected.get_counts())


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