# Copyright 2018-2020 Xanadu Quantum Technologies Inc.

# 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.
"""
Unit tests for the :mod:`pennylane` :class:`QubitDevice` class.
"""
import pytest
import numpy as np
from random import random

import pennylane as qml
from pennylane import QubitDevice, DeviceError, QuantumFunctionError
from pennylane.operation import Sample, Variance, Expectation, Probability, State
from pennylane.circuit_graph import CircuitGraph
from pennylane.wires import Wires
from pennylane.tape import QuantumTape
from pennylane.measure import state

mock_qubit_device_paulis = ["PauliX", "PauliY", "PauliZ"]
mock_qubit_device_rotations = ["RX", "RY", "RZ"]


# pylint: disable=abstract-class-instantiated, no-self-use, redefined-outer-name, invalid-name


@pytest.fixture(scope="function")
def mock_qubit_device(monkeypatch):
    """A function to create a mock device that mocks most of the methods except for e.g. probability()"""
    with monkeypatch.context() as m:
        m.setattr(QubitDevice, "__abstractmethods__", frozenset())
        m.setattr(QubitDevice, "_capabilities", mock_qubit_device_capabilities)
        m.setattr(QubitDevice, "operations", ["PauliY", "RX", "Rot"])
        m.setattr(QubitDevice, "observables", ["PauliZ"])
        m.setattr(QubitDevice, "short_name", "MockDevice")
        m.setattr(QubitDevice, "expval", lambda self, *args, **kwargs: 0)
        m.setattr(QubitDevice, "var", lambda self, *args, **kwargs: 0)
        m.setattr(QubitDevice, "sample", lambda self, *args, **kwargs: 0)
        m.setattr(QubitDevice, "apply", lambda self, *args, **kwargs: None)

        def get_qubit_device(wires=1):
            return QubitDevice(wires=wires)

        yield get_qubit_device


@pytest.fixture(scope="function")
def mock_qubit_device_extract_stats(monkeypatch):
    """A function to create a mock device that mocks the methods related to
    statistics (expval, var, sample, probability)"""
    with monkeypatch.context() as m:
        m.setattr(QubitDevice, "__abstractmethods__", frozenset())
        m.setattr(QubitDevice, "_capabilities", mock_qubit_device_capabilities)
        m.setattr(QubitDevice, "operations", ["PauliY", "RX", "Rot"])
        m.setattr(QubitDevice, "observables", ["PauliZ"])
        m.setattr(QubitDevice, "short_name", "MockDevice")
        m.setattr(QubitDevice, "expval", lambda self, *args, **kwargs: 0)
        m.setattr(QubitDevice, "var", lambda self, *args, **kwargs: 0)
        m.setattr(QubitDevice, "sample", lambda self, *args, **kwargs: 0)
        m.setattr(QubitDevice, "state", 0)
        m.setattr(QubitDevice, "density_matrix", lambda self, wires=None: 0)
        m.setattr(QubitDevice, "probability", lambda self, wires=None, *args, **kwargs: 0)
        m.setattr(QubitDevice, "apply", lambda self, x: x)

        def get_qubit_device(wires=1):
            return QubitDevice(wires=wires)

        yield get_qubit_device


@pytest.fixture(scope="function")
def mock_qubit_device_with_original_statistics(monkeypatch):
    """A function to create a mock device that mocks only basis methods and uses the original
    statistics related methods"""
    with monkeypatch.context() as m:
        m.setattr(QubitDevice, "__abstractmethods__", frozenset())
        m.setattr(QubitDevice, "_capabilities", mock_qubit_device_capabilities)
        m.setattr(QubitDevice, "operations", ["PauliY", "RX", "Rot"])
        m.setattr(QubitDevice, "observables", ["PauliZ"])
        m.setattr(QubitDevice, "short_name", "MockDevice")

        def get_qubit_device(wires=1):
            return QubitDevice(wires=wires)

        yield get_qubit_device


mock_qubit_device_capabilities = {
    "measurements": "everything",
    "returns_state": True,
    "noise_models": ["depolarizing", "bitflip"],
}


@pytest.fixture(scope="function")
def mock_qubit_device_with_paulis_and_methods(monkeypatch):
    """A function to create a mock device that supports Paulis in its capabilities"""
    with monkeypatch.context() as m:
        m.setattr(QubitDevice, "__abstractmethods__", frozenset())
        m.setattr(QubitDevice, "_capabilities", mock_qubit_device_capabilities)
        m.setattr(QubitDevice, "operations", mock_qubit_device_paulis)
        m.setattr(QubitDevice, "observables", mock_qubit_device_paulis)
        m.setattr(QubitDevice, "short_name", "MockDevice")
        m.setattr(QubitDevice, "expval", lambda self, *args, **kwargs: 0)
        m.setattr(QubitDevice, "var", lambda self, *args, **kwargs: 0)
        m.setattr(QubitDevice, "sample", lambda self, *args, **kwargs: 0)
        m.setattr(QubitDevice, "apply", lambda self, x, rotations: None)

        def get_qubit_device(wires=1):
            return QubitDevice(wires=wires)

        yield get_qubit_device


@pytest.fixture(scope="function")
def mock_qubit_device_with_paulis_rotations_and_methods(monkeypatch):
    """A function to create a mock device that supports Paulis in its capabilities"""
    with monkeypatch.context() as m:
        m.setattr(QubitDevice, "__abstractmethods__", frozenset())
        m.setattr(QubitDevice, "_capabilities", mock_qubit_device_capabilities)
        m.setattr(QubitDevice, "operations", mock_qubit_device_paulis + mock_qubit_device_rotations)
        m.setattr(QubitDevice, "observables", mock_qubit_device_paulis)
        m.setattr(QubitDevice, "short_name", "MockDevice")
        m.setattr(QubitDevice, "expval", lambda self, *args, **kwargs: 0)
        m.setattr(QubitDevice, "var", lambda self, *args, **kwargs: 0)
        m.setattr(QubitDevice, "sample", lambda self, *args, **kwargs: 0)
        m.setattr(QubitDevice, "apply", lambda self, x: None)

        def get_qubit_device(wires=1):
            return QubitDevice(wires=wires)

        yield get_qubit_device


class TestOperations:
    """Tests the logic related to operations"""

    def test_op_queue_accessed_outside_execution_context(self, mock_qubit_device):
        """Tests that a call to op_queue outside the execution context raises the correct error"""

        with pytest.raises(
            ValueError, match="Cannot access the operation queue outside of the execution context!"
        ):
            dev = mock_qubit_device()
            dev.op_queue

    def test_op_queue_is_filled_during_execution(
        self, mock_qubit_device_with_paulis_and_methods, monkeypatch
    ):
        """Tests that the op_queue is correctly filled when apply is called and that accessing
        op_queue raises no error"""

        with qml.tape.QuantumTape() as tape:
            queue = [qml.PauliX(wires=0), qml.PauliY(wires=1), qml.PauliZ(wires=2)]
            observables = [qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(1))]

        call_history = []

        with monkeypatch.context() as m:
            m.setattr(
                QubitDevice,
                "apply",
                lambda self, x, **kwargs: call_history.extend(x + kwargs.get("rotations", [])),
            )
            m.setattr(QubitDevice, "analytic_probability", lambda *args: None)
            dev = mock_qubit_device_with_paulis_and_methods()
            dev.execute(tape)

        assert call_history == queue

        assert len(call_history) == 3
        assert isinstance(call_history[0], qml.PauliX)
        assert call_history[0].wires == Wires([0])

        assert isinstance(call_history[1], qml.PauliY)
        assert call_history[1].wires == Wires([1])

        assert isinstance(call_history[2], qml.PauliZ)
        assert call_history[2].wires == Wires([2])

    def test_unsupported_operations_raise_error(self, mock_qubit_device_with_paulis_and_methods):
        """Tests that the operations are properly applied and queued"""
        with qml.tape.QuantumTape() as tape:
            queue = [qml.PauliX(wires=0), qml.PauliY(wires=1), qml.Hadamard(wires=2)]
            observables = [qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(1))]

        with pytest.raises(DeviceError, match="Gate Hadamard not supported on device"):
            dev = mock_qubit_device_with_paulis_and_methods()
            dev.execute(tape)

    numeric_queues = [
        [qml.RX(0.3, wires=[0])],
        [
            qml.RX(0.3, wires=[0]),
            qml.RX(0.4, wires=[1]),
            qml.RX(0.5, wires=[2]),
        ],
    ]

    observables = [[qml.PauliZ(0)], [qml.PauliX(0)], [qml.PauliY(0)]]

    @pytest.mark.parametrize("observables", observables)
    @pytest.mark.parametrize("queue", numeric_queues)
    def test_passing_keyword_arguments_to_execute(
        self, mock_qubit_device_with_paulis_rotations_and_methods, monkeypatch, queue, observables
    ):
        """Tests that passing keyword arguments to execute propagates those kwargs to the apply()
        method"""
        with qml.tape.QuantumTape() as tape:
            for op in queue + observables:
                op.queue()

        call_history = {}

        with monkeypatch.context() as m:
            m.setattr(QubitDevice, "apply", lambda self, x, **kwargs: call_history.update(kwargs))
            dev = mock_qubit_device_with_paulis_rotations_and_methods()
            dev.execute(tape, hash=tape.graph.hash)

        len(call_history.items()) == 1
        call_history["hash"] = tape.graph.hash


class TestObservables:
    """Tests the logic related to observables"""

    # pylint: disable=no-self-use, redefined-outer-name

    def test_obs_queue_accessed_outside_execution_context(self, mock_qubit_device):
        """Tests that a call to op_queue outside the execution context raises the correct error"""

        with pytest.raises(
            ValueError,
            match="Cannot access the observable value queue outside of the execution context!",
        ):
            dev = mock_qubit_device()
            dev.obs_queue

    def test_unsupported_observables_raise_error(self, mock_qubit_device_with_paulis_and_methods):
        """Tests that the operations are properly applied and queued"""
        with qml.tape.QuantumTape() as tape:
            queue = [qml.PauliX(wires=0), qml.PauliY(wires=1), qml.PauliZ(wires=2)]

            observables = [
                qml.expval(qml.Hadamard(0)),
                qml.var(qml.PauliZ(1)),
                qml.sample(qml.PauliZ(2)),
            ]

        with pytest.raises(DeviceError, match="Observable Hadamard not supported on device"):
            dev = mock_qubit_device_with_paulis_and_methods()
            dev.execute(tape)

    def test_unsupported_observable_return_type_raise_error(
        self, mock_qubit_device_with_paulis_and_methods, monkeypatch
    ):
        """Check that an error is raised if the return type of an observable is unsupported"""

        with qml.tape.QuantumTape() as tape:
            qml.PauliX(wires=0)
            qml.measure.MeasurementProcess(
                return_type="SomeUnsupportedReturnType", obs=qml.PauliZ(0)
            )

        with monkeypatch.context() as m:
            m.setattr(QubitDevice, "apply", lambda self, x, **kwargs: None)
            with pytest.raises(
                qml.QuantumFunctionError, match="Unsupported return type specified for observable"
            ):
                dev = mock_qubit_device_with_paulis_and_methods()
                dev.execute(tape)


class TestParameters:
    """Test for checking device parameter mappings"""

    def test_parameters_accessed_outside_execution_context(self, mock_qubit_device):
        """Tests that a call to parameters outside the execution context raises the correct error"""

        with pytest.raises(
            ValueError,
            match="Cannot access the free parameter mapping outside of the execution context!",
        ):
            dev = mock_qubit_device()
            dev.parameters


class TestExtractStatistics:
    """Test the statistics method"""

    @pytest.mark.parametrize("returntype", [Expectation, Variance, Sample, Probability, State])
    def test_results_created(self, mock_qubit_device_extract_stats, monkeypatch, returntype):
        """Tests that the statistics method simply builds a results list without any side-effects"""

        class SomeObservable(qml.operation.Observable):
            num_params = 0
            num_wires = 1
            par_domain = "F"
            return_type = returntype

        obs = SomeObservable(wires=0)

        with monkeypatch.context() as m:
            dev = mock_qubit_device_extract_stats()
            results = dev.statistics([obs])

        assert results == [0]

    def test_results_no_state(self, mock_qubit_device_extract_stats, monkeypatch):
        """Tests that the statistics method raises an AttributeError when a State return type is
        requested when QubitDevice does not have a state attribute"""
        with monkeypatch.context():
            dev = mock_qubit_device_extract_stats()
            delattr(dev.__class__, "state")
            with pytest.raises(
                qml.QuantumFunctionError, match="The state is not available in the current"
            ):
                dev.statistics([state()])

    @pytest.mark.parametrize("returntype", [None])
    def test_results_created_empty(self, mock_qubit_device_extract_stats, monkeypatch, returntype):
        """Tests that the statistics method returns an empty list if the return type is None"""

        class SomeObservable(qml.operation.Observable):
            num_params = 0
            num_wires = 1
            par_domain = "F"
            return_type = returntype

        obs = SomeObservable(wires=0)

        with monkeypatch.context() as m:
            dev = mock_qubit_device_extract_stats()
            results = dev.statistics([obs])

        assert results == []

    @pytest.mark.parametrize("returntype", ["not None"])
    def test_error_return_type_none(self, mock_qubit_device_extract_stats, monkeypatch, returntype):
        """Tests that the statistics method raises an error if the return type is not well-defined and is not None"""

        assert returntype not in [Expectation, Variance, Sample, Probability, State, None]

        class SomeObservable(qml.operation.Observable):
            num_params = 0
            num_wires = 1
            par_domain = "F"
            return_type = returntype

        obs = SomeObservable(wires=0)

        with pytest.raises(qml.QuantumFunctionError, match="Unsupported return type"):
            dev = mock_qubit_device_extract_stats()
            dev.statistics([obs])


class TestGenerateSamples:
    """Test the generate_samples method"""

    def test_auxiliary_methods_called_correctly(self, mock_qubit_device, monkeypatch):
        """Tests that the generate_samples method calls on its auxiliary methods correctly"""

        dev = mock_qubit_device()
        number_of_states = 2 ** dev.num_wires

        with monkeypatch.context() as m:
            # Mock the auxiliary methods such that they return the expected values
            m.setattr(QubitDevice, "sample_basis_states", lambda self, wires, b: wires)
            m.setattr(QubitDevice, "states_to_binary", lambda a, b: (a, b))
            m.setattr(QubitDevice, "analytic_probability", lambda *args: None)
            m.setattr(QubitDevice, "shots", 1000)
            dev._samples = dev.generate_samples()

        assert dev._samples == (number_of_states, dev.num_wires)


class TestSampleBasisStates:
    """Test the sample_basis_states method"""

    def test_sampling_with_correct_arguments(self, mock_qubit_device, monkeypatch):
        """Tests that the sample_basis_states method samples with the correct arguments"""

        shots = 1000

        number_of_states = 4
        dev = mock_qubit_device()
        dev.shots = shots
        state_probs = [0.1, 0.2, 0.3, 0.4]

        with monkeypatch.context() as m:
            # Mock the numpy.random.choice method such that it returns the expected values
            m.setattr("numpy.random.choice", lambda x, y, p: (x, y, p))
            res = dev.sample_basis_states(number_of_states, state_probs)

        assert np.array_equal(res[0], np.array([0, 1, 2, 3]))
        assert res[1] == shots
        assert res[2] == state_probs

    def test_raises_deprecation_warning(self, mock_qubit_device, monkeypatch):
        """Test that sampling basis states on a device with shots=None produces a warning."""

        dev = mock_qubit_device()
        number_of_states = 4
        dev.shots = None
        state_probs = [0.1, 0.2, 0.3, 0.4]

        with pytest.raises(
            qml.QuantumFunctionError,
            match="The number of shots has to be explicitly set on the device",
        ):
            dev.sample_basis_states(number_of_states, state_probs)


class TestStatesToBinary:
    """Test the states_to_binary method"""

    def test_correct_conversion_two_states(self, mock_qubit_device):
        """Tests that the sample_basis_states method converts samples to binary correctly"""
        wires = 4
        shots = 10

        number_of_states = 2 ** wires
        basis_states = np.arange(number_of_states)
        samples = np.random.choice(basis_states, shots)

        dev = mock_qubit_device()
        res = dev.states_to_binary(samples, wires)

        format_smt = "{{:0{}b}}".format(wires)
        expected = np.array([[int(x) for x in list(format_smt.format(i))] for i in samples])

        assert np.all(res == expected)

    test_binary_conversion_data = [
        (np.array([2, 3, 2, 0, 0]), np.array([[1, 0], [1, 1], [1, 0], [0, 0], [0, 0]])),
        (np.array([2, 3, 1, 3, 1]), np.array([[1, 0], [1, 1], [0, 1], [1, 1], [0, 1]])),
        (
            np.array([7, 7, 1, 5, 2]),
            np.array([[1, 1, 1], [1, 1, 1], [0, 0, 1], [1, 0, 1], [0, 1, 0]]),
        ),
    ]

    @pytest.mark.parametrize("samples, binary_states", test_binary_conversion_data)
    def test_correct_conversion(self, mock_qubit_device, samples, binary_states, tol):
        """Tests that the states_to_binary method converts samples to binary correctly"""
        dev = mock_qubit_device()
        dev.shots = 5
        wires = binary_states.shape[1]
        res = dev.states_to_binary(samples, wires)
        assert np.allclose(res, binary_states, atol=tol, rtol=0)


class TestExpval:
    """Test the expval method"""

    def test_analytic_expval(self, mock_qubit_device_with_original_statistics, monkeypatch):
        """Tests that expval method when the analytic attribute is True

        Additional QubitDevice methods that are mocked:
        -rotate_basis
        -probability
        """
        obs = qml.PauliX(0)
        probs = [0.5, 0.5]
        dev = mock_qubit_device_with_original_statistics()

        assert dev.shots is None

        call_history = []
        with monkeypatch.context() as m:
            m.setattr(QubitDevice, "probability", lambda self, wires=None: probs)
            res = dev.expval(obs)

        assert res == (obs.eigvals @ probs).real

    def test_non_analytic_expval(self, mock_qubit_device_with_original_statistics, monkeypatch):
        """Tests that expval method when the analytic attribute is False

        Additional QubitDevice methods that are mocked:
        -rotate_basis
        -sample
        -numpy.mean
        """
        obs = qml.PauliX(0)
        dev = mock_qubit_device_with_original_statistics()

        dev.shots = 1000

        call_history = []
        with monkeypatch.context() as m:
            m.setattr(QubitDevice, "sample", lambda self, obs, *args, **kwargs: obs)
            m.setattr("numpy.mean", lambda obs, axis=None: obs)
            res = dev.expval(obs)

        assert res == obs

    def test_no_eigval_error(self, mock_qubit_device_with_original_statistics):
        """Tests that an error is thrown if expval is called with an observable that does not have eigenvalues defined."""
        dev = mock_qubit_device_with_original_statistics()
        with pytest.raises(ValueError, match="Cannot compute analytic expectations"):
            dev.expval(qml.Hamiltonian([1.0], [qml.PauliX(0)]))


class TestVar:
    """Test the var method"""

    def test_analytic_var(self, mock_qubit_device_with_original_statistics, monkeypatch):
        """Tests that var method when the analytic attribute is True

        Additional QubitDevice methods that are mocked:
        -rotate_basis
        -probability
        """
        obs = qml.PauliX(0)
        probs = [0.5, 0.5]
        dev = mock_qubit_device_with_original_statistics()

        assert dev.shots is None

        call_history = []
        with monkeypatch.context() as m:
            m.setattr(QubitDevice, "probability", lambda self, wires=None: probs)
            res = dev.var(obs)

        assert res == (obs.eigvals ** 2) @ probs - (obs.eigvals @ probs).real ** 2

    def test_non_analytic_var(self, mock_qubit_device_with_original_statistics, monkeypatch):
        """Tests that var method when the analytic attribute is False

        Additional QubitDevice methods that are mocked:
        -rotate_basis
        -sample
        -numpy.var
        """
        obs = qml.PauliX(0)
        dev = mock_qubit_device_with_original_statistics()

        dev.shots = 1000

        call_history = []
        with monkeypatch.context() as m:
            m.setattr(QubitDevice, "sample", lambda self, obs, *args, **kwargs: obs)
            m.setattr("numpy.var", lambda obs, axis=None: obs)
            res = dev.var(obs)

        assert res == obs

    def test_no_eigval_error(self, mock_qubit_device_with_original_statistics):
        """Tests that an error is thrown if var is called with an observable that does not have eigenvalues defined."""
        dev = mock_qubit_device_with_original_statistics()
        with pytest.raises(ValueError, match="Cannot compute analytic variance"):
            dev.var(qml.Hamiltonian([1.0], [qml.PauliX(0)]))


class TestSample:
    """Test the sample method"""

    def test_only_ones_minus_ones(
        self, mock_qubit_device_with_original_statistics, monkeypatch, tol
    ):
        """Test that pauli_eigvals_as_samples method only produces -1 and 1 samples"""
        obs = qml.PauliX(0)
        dev = mock_qubit_device_with_original_statistics()
        dev._samples = np.array([[1, 0], [0, 0]])

        with monkeypatch.context() as m:
            res = dev.sample(obs)

        assert np.allclose(res ** 2, 1, atol=tol, rtol=0)

    def test_correct_custom_eigenvalues(
        self, mock_qubit_device_with_original_statistics, monkeypatch, tol
    ):
        """Test that pauli_eigvals_as_samples method only produces samples of eigenvalues"""
        obs = qml.PauliX(0) @ qml.PauliZ(1)
        dev = mock_qubit_device_with_original_statistics(wires=2)
        dev._samples = np.array([[1, 0], [0, 0]])

        with monkeypatch.context() as m:
            res = dev.sample(obs)

        assert np.array_equal(res, np.array([-1, 1]))

    def test_sample_with_no_observable_and_no_wires(
        self, mock_qubit_device_with_original_statistics, tol
    ):
        """Test that when we sample a device without providing an observable or wires then it
        will return the raw samples"""
        obs = qml.measure.sample(op=None, wires=None)
        dev = mock_qubit_device_with_original_statistics(wires=2)
        generated_samples = np.array([[1, 0], [1, 1]])
        dev._samples = generated_samples

        res = dev.sample(obs)
        assert np.array_equal(res, generated_samples)

    def test_sample_with_no_observable_and_with_wires(
        self, mock_qubit_device_with_original_statistics, tol
    ):
        """Test that when we sample a device without providing an observable but we specify
        wires then it returns the generated samples for only those wires"""
        obs = qml.measure.sample(op=None, wires=[1])
        dev = mock_qubit_device_with_original_statistics(wires=2)
        generated_samples = np.array([[1, 0], [1, 1]])
        dev._samples = generated_samples

        wire_samples = np.array([[0], [1]])
        res = dev.sample(obs)

        assert np.array_equal(res, wire_samples)

    def test_no_eigval_error(self, mock_qubit_device_with_original_statistics):
        """Tests that an error is thrown if sample is called with an observable that does not have eigenvalues defined."""
        dev = mock_qubit_device_with_original_statistics()
        dev._samples = np.array([[1, 0], [0, 0]])
        with pytest.raises(ValueError, match="Cannot compute samples"):
            dev.sample(qml.Hamiltonian([1.0], [qml.PauliX(0)]))


class TestEstimateProb:
    """Test the estimate_probability method"""

    @pytest.mark.parametrize(
        "wires, expected", [([0], [0.5, 0.5]), (None, [0.5, 0, 0, 0.5]), ([0, 1], [0.5, 0, 0, 0.5])]
    )
    def test_estimate_probability(
        self, wires, expected, mock_qubit_device_with_original_statistics, monkeypatch
    ):
        """Tests probability method when the analytic attribute is True."""
        dev = mock_qubit_device_with_original_statistics(wires=2)
        samples = np.array([[0, 0], [1, 1], [1, 1], [0, 0]])

        with monkeypatch.context() as m:
            m.setattr(dev, "_samples", samples)
            m.setattr(dev, "shots", 4)
            res = dev.estimate_probability(wires=wires)

        assert np.allclose(res, expected)


class TestMarginalProb:
    """Test the marginal_prob method"""

    @pytest.mark.parametrize(
        "wires, inactive_wires",
        [
            ([0], [1, 2]),
            ([1], [0, 2]),
            ([2], [0, 1]),
            ([0, 1], [2]),
            ([0, 2], [1]),
            ([1, 2], [0]),
            ([0, 1, 2], []),
            (Wires([0]), [1, 2]),
            (Wires([0, 1]), [2]),
            (Wires([0, 1, 2]), []),
        ],
    )
    def test_correct_arguments_for_marginals(
        self, mock_qubit_device_with_original_statistics, mocker, wires, inactive_wires, tol
    ):
        """Test that the correct arguments are passed to the marginal_prob method"""

        # Generate probabilities
        probs = np.array([random() for i in range(2 ** 3)])
        probs /= sum(probs)

        spy = mocker.spy(np, "sum")
        dev = mock_qubit_device_with_original_statistics(wires=3)
        res = dev.marginal_prob(probs, wires=wires)
        array_call = spy.call_args[0][0]
        axis_call = spy.call_args[1]["axis"]

        assert np.allclose(array_call.flatten(), probs, atol=tol, rtol=0)
        assert axis_call == tuple(inactive_wires)

    marginal_test_data = [
        (np.array([0.1, 0.2, 0.3, 0.4]), np.array([0.4, 0.6]), [1]),
        (np.array([0.1, 0.2, 0.3, 0.4]), np.array([0.3, 0.7]), Wires([0])),
        (
            np.array(
                [
                    0.17794671,
                    0.06184147,
                    0.21909549,
                    0.04932204,
                    0.19595214,
                    0.19176834,
                    0.08495311,
                    0.0191207,
                ]
            ),
            np.array([0.3970422, 0.28090525, 0.11116351, 0.21088904]),
            [2, 0],
        ),
    ]

    @pytest.mark.parametrize("probs, marginals, wires", marginal_test_data)
    def test_correct_marginals_returned(
        self, mock_qubit_device_with_original_statistics, probs, marginals, wires, tol
    ):
        """Test that the correct marginals are returned by the marginal_prob method"""
        num_wires = int(np.log2(len(probs)))
        dev = mock_qubit_device_with_original_statistics(num_wires)
        res = dev.marginal_prob(probs, wires=wires)
        assert np.allclose(res, marginals, atol=tol, rtol=0)

    @pytest.mark.parametrize("probs, marginals, wires", marginal_test_data)
    def test_correct_marginals_returned_wires_none(
        self, mock_qubit_device_with_original_statistics, probs, marginals, wires, tol
    ):
        """Test that passing wires=None simply returns the original probability."""
        num_wires = int(np.log2(len(probs)))
        dev = mock_qubit_device_with_original_statistics(wires=num_wires)
        dev.num_wires = num_wires

        res = dev.marginal_prob(probs, wires=None)
        assert np.allclose(res, probs, atol=tol, rtol=0)


class TestActiveWires:
    """Test that the active_wires static method works as required."""

    def test_active_wires_from_queue(self, mock_qubit_device):
        queue = [qml.CNOT(wires=[0, 2]), qml.RX(0.2, wires=0), qml.expval(qml.PauliX(wires=5))]

        dev = mock_qubit_device(wires=6)
        res = dev.active_wires(queue)

        assert res == Wires([0, 2, 5])


class TestCapabilities:
    """Test that a default qubit device defines capabilities that all devices inheriting
    from it will automatically have."""

    def test_defines_correct_capabilities(self):
        """Test that the device defines the right capabilities"""
        capabilities = {
            "model": "qubit",
            "supports_finite_shots": True,
            "supports_tensor_observables": True,
            "returns_probs": True,
        }
        assert capabilities == QubitDevice.capabilities()


class TestExecution:
    """Tests for the execute method"""

    def test_device_executions(self):
        """Test the number of times a qubit device is executed over a QNode's
        lifetime is tracked by `num_executions`"""

        dev_1 = qml.device("default.qubit", wires=2)

        def circuit_1(x, y):
            qml.RX(x, wires=[0])
            qml.RY(y, wires=[1])
            qml.CNOT(wires=[0, 1])
            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))

        node_1 = qml.QNode(circuit_1, dev_1)
        num_evals_1 = 10

        for _ in range(num_evals_1):
            node_1(0.432, 0.12)
        assert dev_1.num_executions == num_evals_1

        # test a second instance of a default qubit device
        dev_2 = qml.device("default.qubit", wires=2)

        def circuit_2(x, y):
            qml.RX(x, wires=[0])
            qml.CNOT(wires=[0, 1])
            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))

        node_2 = qml.QNode(circuit_2, dev_2)
        num_evals_2 = 5

        for _ in range(num_evals_2):
            node_2(0.432, 0.12)
        assert dev_2.num_executions == num_evals_2

        # test a new circuit on an existing instance of a qubit device
        def circuit_3(x, y):
            qml.RY(y, wires=[1])
            qml.CNOT(wires=[0, 1])
            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))

        node_3 = qml.QNode(circuit_3, dev_1)
        num_evals_3 = 7

        for _ in range(num_evals_3):
            node_3(0.432, 0.12)
        assert dev_1.num_executions == num_evals_1 + num_evals_3


class TestBatchExecution:
    """Tests for the batch_execute method."""

    with qml.tape.QuantumTape() as tape1:
        qml.PauliX(wires=0)
        qml.expval(qml.PauliZ(wires=0)), qml.expval(qml.PauliZ(wires=1))

    with qml.tape.JacobianTape() as tape2:
        qml.PauliX(wires=0)
        qml.expval(qml.PauliZ(wires=0))

    @pytest.mark.parametrize("n_tapes", [1, 2, 3])
    def test_calls_to_execute(self, n_tapes, mocker, mock_qubit_device_with_paulis_and_methods):
        """Tests that the device's execute method is called the correct number of times."""

        dev = mock_qubit_device_with_paulis_and_methods(wires=2)
        spy = mocker.spy(QubitDevice, "execute")

        tapes = [self.tape1] * n_tapes
        dev.batch_execute(tapes)

        assert spy.call_count == n_tapes

    @pytest.mark.parametrize("n_tapes", [1, 2, 3])
    def test_calls_to_reset(self, n_tapes, mocker, mock_qubit_device_with_paulis_and_methods):
        """Tests that the device's reset method is called the correct number of times."""

        dev = mock_qubit_device_with_paulis_and_methods(wires=2)

        spy = mocker.spy(QubitDevice, "reset")

        tapes = [self.tape1] * n_tapes
        dev.batch_execute(tapes)

        assert spy.call_count == n_tapes

    def test_result(self, mock_qubit_device_with_paulis_and_methods, tol):
        """Tests that the result has the correct shape and entry types."""

        dev = mock_qubit_device_with_paulis_and_methods(wires=2)

        tapes = [self.tape1, self.tape2]
        res = dev.batch_execute(tapes)

        assert len(res) == 2
        assert np.allclose(res[0], dev.execute(self.tape1), rtol=tol, atol=0)
        assert np.allclose(res[1], dev.execute(self.tape2), rtol=tol, atol=0)

    def test_result_empty_tape(self, mock_qubit_device_with_paulis_and_methods, tol):
        """Tests that the result has the correct shape and entry types for empty tapes."""

        dev = mock_qubit_device_with_paulis_and_methods(wires=2)

        empty_tape = qml.tape.QuantumTape()
        tapes = [empty_tape] * 3
        res = dev.batch_execute(tapes)

        assert len(res) == 3
        assert np.allclose(res[0], dev.execute(empty_tape), rtol=tol, atol=0)


class TestShotList:
    """Tests for passing shots as a list"""

    shot_data = [
        [[1, 2, 3, 10], [(1, 1), (2, 1), (3, 1), (10, 1)], (4,), 16],
        [
            [1, 2, 2, 2, 10, 1, 1, 5, 1, 1, 1],
            [(1, 1), (2, 3), (10, 1), (1, 2), (5, 1), (1, 3)],
            (11,),
            27,
        ],
        [[10, 10, 10], [(10, 3)], (3,), 30],
        [[(10, 3)], [(10, 3)], (3,), 30],
    ]

    @pytest.mark.parametrize("shot_list,shot_vector,expected_shape,total_shots", shot_data)
    def test_single_expval(self, shot_list, shot_vector, expected_shape, total_shots):
        """Test a single expectation value"""
        dev = qml.device("default.qubit", wires=2, shots=shot_list)

        @qml.qnode(dev)
        def circuit(x):
            qml.RX(x, wires=0)
            qml.CNOT(wires=[0, 1])
            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1))

        res = circuit(0.5)

        assert res.shape == expected_shape
        assert circuit.device._shot_vector == shot_vector
        assert circuit.device.shots == total_shots

    shot_data = [
        [[1, 2, 3, 10], [(1, 1), (2, 1), (3, 1), (10, 1)], (4, 2), 16],
        [
            [1, 2, 2, 2, 10, 1, 1, 5, 1, 1, 1],
            [(1, 1), (2, 3), (10, 1), (1, 2), (5, 1), (1, 3)],
            (11, 2),
            27,
        ],
        [[10, 10, 10], [(10, 3)], (3, 2), 30],
        [[(10, 3)], [(10, 3)], (3, 2), 30],
    ]

    @pytest.mark.parametrize("shot_list,shot_vector,expected_shape,total_shots", shot_data)
    def test_multiple_expval(self, shot_list, shot_vector, expected_shape, total_shots):
        """Test multiple expectation values"""
        dev = qml.device("default.qubit", wires=2, shots=shot_list)

        @qml.qnode(dev)
        def circuit(x, y):
            qml.RX(x, wires=0)
            qml.RY(y, wires=0)
            qml.CNOT(wires=[0, 1])
            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0))

        res = circuit(0.5, 0.1)

        assert res.shape == expected_shape
        assert circuit.device._shot_vector == shot_vector
        assert circuit.device.shots == total_shots

        # test gradient works
        res = qml.jacobian(circuit)(0.5, 0.1)
        assert res.shape == (2,) + expected_shape

    shot_data = [
        [[1, 2, 3, 10], [(1, 1), (2, 1), (3, 1), (10, 1)], (4, 4), 16],
        [
            [1, 2, 2, 2, 10, 1, 1, 5, 1, 1, 1],
            [(1, 1), (2, 3), (10, 1), (1, 2), (5, 1), (1, 3)],
            (11, 4),
            27,
        ],
        [[10, 10, 10], [(10, 3)], (3, 4), 30],
        [[(10, 3)], [(10, 3)], (3, 4), 30],
    ]

    @pytest.mark.parametrize("shot_list,shot_vector,expected_shape,total_shots", shot_data)
    def test_probs(self, shot_list, shot_vector, expected_shape, total_shots):
        """Test a probability return"""
        dev = qml.device("default.qubit", wires=2, shots=shot_list)

        @qml.qnode(dev)
        def circuit(x, y):
            qml.RX(x, wires=0)
            qml.RY(y, wires=0)
            qml.CNOT(wires=[0, 1])
            return qml.probs(wires=[0, 1])

        res = circuit(0.5, 0.1)

        assert res.shape == expected_shape
        assert circuit.device._shot_vector == shot_vector
        assert circuit.device.shots == total_shots

        # test gradient works
        res = qml.jacobian(circuit)(0.5, 0.1)

    shot_data = [
        [[1, 2, 3, 10], [(1, 1), (2, 1), (3, 1), (10, 1)], (4, 2, 2), 16],
        [
            [1, 2, 2, 2, 10, 1, 1, 5, 1, 1, 1],
            [(1, 1), (2, 3), (10, 1), (1, 2), (5, 1), (1, 3)],
            (11, 2, 2),
            27,
        ],
        [[10, 10, 10], [(10, 3)], (3, 2, 2), 30],
        [[(10, 3)], [(10, 3)], (3, 2, 2), 30],
    ]

    @pytest.mark.parametrize("shot_list,shot_vector,expected_shape,total_shots", shot_data)
    def test_multiple_probs(self, shot_list, shot_vector, expected_shape, total_shots):
        """Test multiple probability returns"""
        dev = qml.device("default.qubit", wires=2, shots=shot_list)

        @qml.qnode(dev)
        def circuit(x, y):
            qml.RX(x, wires=0)
            qml.RY(y, wires=0)
            qml.CNOT(wires=[0, 1])
            return qml.probs(wires=0), qml.probs(wires=1)

        res = circuit(0.5, 0.1)

        assert res.shape == expected_shape
        assert circuit.device._shot_vector == shot_vector
        assert circuit.device.shots == total_shots

        # test gradient works
        res = qml.jacobian(circuit)(0.5, 0.1)

    shot_data = [
        [[1, 2, 3, 10], [(1, 1), (2, 1), (3, 1), (10, 1)], [(), (2,), (3,), (10,)], 16],
        [
            [1, 2, 2, 2, 10, 1, 1, 5, 1, 1, 1],
            [(1, 1), (2, 3), (10, 1), (1, 2), (5, 1), (1, 3)],
            [(), (2,), (2,), (2,), (10,), (), (), (5,), (), (), ()],
            27,
        ],
        [[10, 10, 10], [(10, 3)], [(10,), (10,), (10,)], 30],
        [[(10, 3)], [(10, 3)], [(10,), (10,), (10,)], 30],
    ]

    @pytest.mark.parametrize("shot_list,shot_vector,expected_shapes,total_shots", shot_data)
    def test_sample(self, shot_list, shot_vector, expected_shapes, total_shots):
        """Test sample returns"""
        dev = qml.device("default.qubit", wires=2, shots=shot_list)

        @qml.qnode(dev)
        def circuit(x, y):
            qml.RX(x, wires=0)
            qml.RY(y, wires=0)
            qml.CNOT(wires=[0, 1])
            return qml.sample(qml.PauliZ(wires=0))

        res = circuit(0.5, 0.1)

        for r, shape in zip(res, expected_shapes):
            assert r.shape == shape

        assert circuit.device._shot_vector == shot_vector
        assert circuit.device.shots == total_shots

    shot_data = [
        [[1, 2, 3, 10], [(1, 1), (2, 1), (3, 1), (10, 1)], [(2,), (2, 2), (3, 2), (10, 2)], 16],
        [
            [1, 2, 2, 2, 10, 1, 1, 5, 1, 1, 1],
            [(1, 1), (2, 3), (10, 1), (1, 2), (5, 1), (1, 3)],
            [(2,), (2, 2), (2, 2), (2, 2), (10, 2), (2,), (2,), (5, 2), (2,), (2,), (2,)],
            27,
        ],
        [[10, 10, 10], [(10, 3)], [(10, 2), (10, 2), (10, 2)], 30],
        [[(10, 3)], [(10, 3)], [(10, 2), (10, 2), (10, 2)], 30],
    ]

    @pytest.mark.parametrize("shot_list,shot_vector,expected_shapes,total_shots", shot_data)
    def test_multiple_sample(self, shot_list, shot_vector, expected_shapes, total_shots):
        """Test sample returns"""
        dev = qml.device("default.qubit", wires=2, shots=shot_list)

        @qml.qnode(dev)
        def circuit(x, y):
            qml.RX(x, wires=0)
            qml.RY(y, wires=0)
            qml.CNOT(wires=[0, 1])
            return qml.sample(qml.PauliZ(wires=0)), qml.sample(qml.PauliZ(wires=1))

        res = circuit(0.5, 0.1)

        for r, shape in zip(res, expected_shapes):
            assert r.shape == shape

        assert circuit.device._shot_vector == shot_vector
        assert circuit.device.shots == total_shots

    def test_invalid_shot_list(self):
        """Test exception raised if the shot list is the wrong type"""
        with pytest.raises(qml.DeviceError, match="Shots must be"):
            qml.device("default.qubit", wires=2, shots=0.5)

        with pytest.raises(ValueError, match="Unknown shot sequence"):
            qml.device("default.qubit", wires=2, shots=["a", "b", "c"])
