# 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.
"""Integration tests for using the Torch interface with a QNode"""
import pytest
import numpy as np

torch = pytest.importorskip("torch", minversion="1.3")
from torch.autograd.functional import hessian, jacobian

import pennylane as qml
from pennylane import qnode, QNode
from pennylane.tape import JacobianTape

qubit_device_and_diff_method = [
    ["default.qubit", "finite-diff", "backward"],
    ["default.qubit", "parameter-shift", "backward"],
    ["default.qubit", "backprop", "forward"],
    ["default.qubit", "adjoint", "forward"],
    ["default.qubit", "adjoint", "backward"],
]

@pytest.mark.parametrize("dev_name,diff_method,mode", qubit_device_and_diff_method)
class TestQNode:
    """Test that using the QNode with Torch integrates with the PennyLane stack"""

    def test_execution_with_interface(self, dev_name, diff_method, mode):
        """Test execution works with the interface"""
        if diff_method == "backprop":
            pytest.skip("Test does not support backprop")

        dev = qml.device(dev_name, wires=1)

        @qnode(dev, diff_method=diff_method, mode=mode, interface="torch")
        def circuit(a):
            qml.RY(a, wires=0)
            qml.RX(0.2, wires=0)
            return qml.expval(qml.PauliZ(0))

        a = torch.tensor(0.1, requires_grad=True)
        res = circuit(a)

        assert circuit.interface == "torch"

        # with the interface, the tape returns torch tensors

        assert isinstance(res, torch.Tensor)
        assert res.shape == tuple()

        # the tape is able to deduce trainable parameters
        assert circuit.qtape.trainable_params == [0]

        # gradients should work
        res.backward()
        grad = a.grad
        assert isinstance(grad, torch.Tensor)
        assert grad.shape == tuple()

    def test_interface_swap(self, dev_name, diff_method, mode, tol):
        """Test that the Torch interface can be applied to a QNode
        with a pre-existing interface"""
        dev = qml.device(dev_name, wires=1)

        @qnode(dev, diff_method=diff_method, interface="autograd")
        def circuit(a):
            qml.RY(a, wires=0)
            qml.RX(0.2, wires=0)
            return qml.expval(qml.PauliZ(0))

        from pennylane import numpy as anp

        a = anp.array(0.1, requires_grad=True)

        res1 = circuit(a)
        grad_fn = qml.grad(circuit)
        grad1 = grad_fn(a)

        # switch to Torch interface
        circuit.interface = "torch"

        a = torch.tensor(0.1, dtype=torch.float64, requires_grad=True)

        res2 = circuit(a)
        res2.backward()
        grad2 = a.grad
        assert np.allclose(res1, res2.detach().numpy(), atol=tol, rtol=0)
        assert np.allclose(grad1, grad2, atol=tol, rtol=0)

    def test_drawing(self, dev_name, diff_method, mode):
        """Test circuit drawing when using the torch interface"""

        x = torch.tensor(0.1, requires_grad=True)
        y = torch.tensor([0.2, 0.3], requires_grad=True)
        z = torch.tensor(0.4, requires_grad=True)

        dev = qml.device(dev_name, wires=2)

        @qnode(dev, interface="torch", diff_method=diff_method, mode=mode)
        def circuit(p1, p2=y, **kwargs):
            qml.RX(p1, wires=0)
            qml.RY(p2[0] * p2[1], wires=1)
            qml.RX(kwargs["p3"], wires=0)
            qml.CNOT(wires=[0, 1])
            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))

        circuit(p1=x, p3=z)

        result = qml.draw(circuit)(p1=x, p3=z)
        expected = "0: ──RX(0.10)──RX(0.40)─╭C─┤  <Z>\n" "1: ──RY(0.06)───────────╰X─┤  <Z>"

        assert result == expected

    def test_jacobian(self, dev_name, diff_method, mode, mocker, tol):
        """Test jacobian calculation"""
        if diff_method == "parameter-shift":
            spy = mocker.spy(qml.gradients.param_shift, "transform_fn")
        elif diff_method == "finite-diff":
            spy = mocker.spy(qml.gradients.finite_diff, "transform_fn")

        a_val = 0.1
        b_val = 0.2

        a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True)
        b = torch.tensor(b_val, dtype=torch.float64, requires_grad=True)

        dev = qml.device(dev_name, wires=2)

        @qnode(dev, diff_method=diff_method, mode=mode, interface="torch")
        def circuit(a, b):
            qml.RY(a, wires=0)
            qml.RX(b, wires=1)
            qml.CNOT(wires=[0, 1])
            return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))]

        res = circuit(a, b)

        assert circuit.qtape.trainable_params == [0, 1]

        assert isinstance(res, torch.Tensor)
        assert res.shape == (2,)

        expected = [np.cos(a_val), -np.cos(a_val) * np.sin(b_val)]
        assert np.allclose(res.detach().numpy(), expected, atol=tol, rtol=0)

        loss = torch.sum(res)

        loss.backward()
        expected = [
            -np.sin(a_val) + np.sin(a_val) * np.sin(b_val),
            -np.cos(a_val) * np.cos(b_val),
        ]
        assert np.allclose(a.grad, expected[0], atol=tol, rtol=0)
        assert np.allclose(b.grad, expected[1], atol=tol, rtol=0)

        if diff_method in ("parameter-shift", "finite-diff"):
            spy.assert_called()

    @pytest.mark.xfail
    def test_jacobian_dtype(self, dev_name, diff_method, mode, tol):
        """Test calculating the jacobian with a different datatype"""
        if diff_method == "backprop":
            pytest.skip("Test does not support backprop")

        a = torch.tensor(0.1, dtype=torch.float32, requires_grad=True)
        b = torch.tensor(0.2, dtype=torch.float32, requires_grad=True)

        dev = qml.device(dev_name, wires=2)

        @qnode(dev, interface="torch", diff_method=diff_method)
        def circuit(a, b):
            qml.RY(a, wires=0)
            qml.RX(b, wires=1)
            qml.CNOT(wires=[0, 1])
            return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))]

        res = circuit(a, b)

        assert circuit.interface == "torch"
        assert circuit.qtape.trainable_params == [0, 1]

        assert isinstance(res, torch.Tensor)
        assert res.shape == (2,)
        assert res.dtype is torch.float32

        loss = torch.sum(res)
        loss.backward()
        assert a.grad.dtype is torch.float32
        assert b.grad.dtype is torch.float32

    def test_jacobian_options(self, dev_name, diff_method, mode, mocker, tol):
        """Test setting jacobian options"""
        if diff_method != "finite-diff":
            pytest.skip("Test only works with finite-diff")

        spy = mocker.spy(qml.gradients.finite_diff, "transform_fn")

        a = torch.tensor([0.1, 0.2], requires_grad=True)

        dev = qml.device(dev_name, wires=1)

        @qnode(dev, diff_method=diff_method, mode=mode, interface="torch", h=1e-8, approx_order=2)
        def circuit(a):
            qml.RY(a[0], wires=0)
            qml.RX(a[1], wires=0)
            return qml.expval(qml.PauliZ(0))

        res = circuit(a)
        res.backward()

        for args in spy.call_args_list:
            assert args[1]["approx_order"] == 2
            assert args[1]["h"] == 1e-8

    def test_changing_trainability(self, dev_name, diff_method, mode, mocker, tol):
        """Test that changing the trainability of parameters changes the
        number of differentiation requests made"""
        if diff_method != "parameter-shift":
            pytest.skip("Test only supports parameter-shift")

        a_val = 0.1
        b_val = 0.2

        a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True)
        b = torch.tensor(b_val, dtype=torch.float64, requires_grad=True)

        dev = qml.device(dev_name, wires=2)

        @qnode(dev, interface="torch", diff_method=diff_method)
        def circuit(a, b):
            qml.RY(a, wires=0)
            qml.RX(b, wires=1)
            qml.CNOT(wires=[0, 1])
            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))

        res = circuit(a, b)

        # the tape has reported both gate arguments as trainable
        assert circuit.qtape.trainable_params == [0, 1]

        expected = [np.cos(a_val), -np.cos(a_val) * np.sin(b_val)]
        assert np.allclose(res.detach().numpy(), expected, atol=tol, rtol=0)

        spy = mocker.spy(qml.gradients.param_shift, "transform_fn")

        loss = torch.sum(res)
        loss.backward()

        expected = [
            -np.sin(a_val) + np.sin(a_val) * np.sin(b_val),
            -np.cos(a_val) * np.cos(b_val),
        ]
        assert np.allclose([a.grad, b.grad], expected, atol=tol, rtol=0)

        # The parameter-shift rule has been called for each argument
        assert len(spy.spy_return[0]) == 4

        # make the second QNode argument a constant
        a_val = 0.54
        b_val = 0.8

        a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True)
        b = torch.tensor(b_val, dtype=torch.float64, requires_grad=False)

        res = circuit(a, b)

        # the tape has reported only the first argument as trainable
        assert circuit.qtape.trainable_params == [0]

        expected = [np.cos(a_val), -np.cos(a_val) * np.sin(b_val)]
        assert np.allclose(res.detach().numpy(), expected, atol=tol, rtol=0)

        spy.call_args_list = []
        loss = torch.sum(res)
        loss.backward()
        expected = -np.sin(a_val) + np.sin(a_val) * np.sin(b_val)
        assert np.allclose(a.grad, expected, atol=tol, rtol=0)

        # the gradient transform has only been called once
        assert len(spy.call_args_list) == 1

    def test_classical_processing(self, dev_name, diff_method, mode, tol):
        """Test classical processing within the quantum tape"""
        a = torch.tensor(0.1, dtype=torch.float64, requires_grad=True)
        b = torch.tensor(0.2, dtype=torch.float64, requires_grad=False)
        c = torch.tensor(0.3, dtype=torch.float64, requires_grad=True)

        dev = qml.device(dev_name, wires=1)

        @qnode(dev, diff_method=diff_method, mode=mode, interface="torch")
        def circuit(a, b, c):
            qml.RY(a * c, wires=0)
            qml.RZ(b, wires=0)
            qml.RX(c + c**2 + torch.sin(a), wires=0)
            return qml.expval(qml.PauliZ(0))

        res = circuit(a, b, c)

        if diff_method == "finite-diff":
            assert circuit.qtape.trainable_params == [0, 2]
            assert circuit.qtape.get_parameters() == [a * c, c + c**2 + torch.sin(a)]

        res.backward()

        assert isinstance(a.grad, torch.Tensor)
        assert b.grad is None
        assert isinstance(c.grad, torch.Tensor)

    def test_no_trainable_parameters(self, dev_name, diff_method, mode, tol):
        """Test evaluation and Jacobian if there are no trainable parameters"""
        dev = qml.device(dev_name, wires=2)

        @qnode(dev, diff_method=diff_method, mode=mode, interface="torch")
        def circuit(a, b):
            qml.RY(a, wires=0)
            qml.RX(b, wires=0)
            qml.CNOT(wires=[0, 1])
            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))

        a = 0.1
        b = torch.tensor(0.2, dtype=torch.float64, requires_grad=False)

        res = circuit(a, b)

        if diff_method == "finite-diff":
            assert circuit.qtape.trainable_params == []

        assert res.shape == (2,)
        assert isinstance(res, torch.Tensor)

        with pytest.raises(
            RuntimeError,
            match="element 0 of tensors does not require grad and does not have a grad_fn",
        ):
            res.backward()

    @pytest.mark.parametrize(
        "U",
        [
            torch.tensor([[0, 1], [1, 0]], requires_grad=False),
            np.array([[0, 1], [1, 0]]),
        ],
    )
    def test_matrix_parameter(self, dev_name, diff_method, mode, U, tol):
        """Test that the Torch interface works correctly
        with a matrix parameter"""
        a_val = 0.1
        a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True)

        dev = qml.device(dev_name, wires=2)

        @qnode(dev, diff_method=diff_method, mode=mode, interface="torch")
        def circuit(U, a):
            qml.QubitUnitary(U, wires=0)
            qml.RY(a, wires=0)
            return qml.expval(qml.PauliZ(0))

        res = circuit(U, a)

        if diff_method == "finite-diff":
            assert circuit.qtape.trainable_params == [1]

        assert np.allclose(res.detach(), -np.cos(a_val), atol=tol, rtol=0)

        res.backward()
        assert np.allclose(a.grad, np.sin(a_val), atol=tol, rtol=0)

    def test_differentiable_expand(self, dev_name, diff_method, mode, tol):
        """Test that operation and nested tapes expansion
        is differentiable"""

        class U3(qml.U3):
            def expand(self):
                theta, phi, lam = self.data
                wires = self.wires

                with JacobianTape() as tape:
                    qml.Rot(lam, theta, -lam, wires=wires)
                    qml.PhaseShift(phi + lam, wires=wires)

                return tape

        dev = qml.device(dev_name, wires=1)
        a = np.array(0.1)
        p_val = [0.1, 0.2, 0.3]
        p = torch.tensor(p_val, dtype=torch.float64, requires_grad=True)

        @qnode(dev, diff_method=diff_method, mode=mode, interface="torch")
        def circuit(a, p):
            qml.RX(a, wires=0)
            U3(p[0], p[1], p[2], wires=0)
            return qml.expval(qml.PauliX(0))

        res = circuit(a, p)

        assert circuit.qtape.trainable_params == [1, 2, 3]

        expected = np.cos(a) * np.cos(p_val[1]) * np.sin(p_val[0]) + np.sin(a) * (
            np.cos(p_val[2]) * np.sin(p_val[1])
            + np.cos(p_val[0]) * np.cos(p_val[1]) * np.sin(p_val[2])
        )
        assert np.allclose(res.detach().numpy(), expected, atol=tol, rtol=0)

        res.backward()
        expected = np.array(
            [
                np.cos(p_val[1])
                * (np.cos(a) * np.cos(p_val[0]) - np.sin(a) * np.sin(p_val[0]) * np.sin(p_val[2])),
                np.cos(p_val[1]) * np.cos(p_val[2]) * np.sin(a)
                - np.sin(p_val[1])
                * (np.cos(a) * np.sin(p_val[0]) + np.cos(p_val[0]) * np.sin(a) * np.sin(p_val[2])),
                np.sin(a)
                * (
                    np.cos(p_val[0]) * np.cos(p_val[1]) * np.cos(p_val[2])
                    - np.sin(p_val[1]) * np.sin(p_val[2])
                ),
            ]
        )
        assert np.allclose(p.grad, expected, atol=tol, rtol=0)

class TestShotsIntegration:
    """Test that the QNode correctly changes shot value, and
    differentiates it."""

    def test_changing_shots(self, mocker, tol):
        """Test that changing shots works on execution"""
        dev = qml.device("default.qubit", wires=2, shots=None)
        a, b = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64)

        @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift)
        def circuit(a, b):
            qml.RY(a, wires=0)
            qml.RX(b, wires=1)
            qml.CNOT(wires=[0, 1])
            return qml.expval(qml.PauliY(1))

        spy = mocker.spy(dev, "sample")

        # execute with device default shots (None)
        res = circuit(a, b)
        assert torch.allclose(res, -torch.cos(a) * torch.sin(b), atol=tol, rtol=0)
        spy.assert_not_called()

        # execute with shots=100
        res = circuit(a, b, shots=100)
        spy.assert_called()
        assert spy.spy_return.shape == (100,)

        # device state has been unaffected
        assert dev.shots is None
        spy = mocker.spy(dev, "sample")
        res = circuit(a, b)
        assert torch.allclose(res, -torch.cos(a) * torch.sin(b), atol=tol, rtol=0)
        spy.assert_not_called()

    def test_gradient_integration(self, tol):
        """Test that temporarily setting the shots works
        for gradient computations"""
        dev = qml.device("default.qubit", wires=2, shots=None)
        a, b = torch.tensor([0.543, -0.654], requires_grad=True)

        @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift)
        def cost_fn(a, b):
            qml.RY(a, wires=0)
            qml.RX(b, wires=1)
            qml.CNOT(wires=[0, 1])
            return qml.expval(qml.PauliY(1))

        res = jacobian(lambda a, b: cost_fn(a, b, shots=[10000, 10000, 10000]), (a, b))
        res = qml.math.transpose(torch.stack(res))
        assert dev.shots is None
        assert len(res) == 3

        expected = torch.tensor([torch.sin(a) * torch.sin(b), -torch.cos(a) * torch.cos(b)])
        assert torch.allclose(torch.mean(res, axis=0), expected, atol=0.1, rtol=0)

    def test_multiple_gradient_integration(self, tol):
        """Test that temporarily setting the shots works
        for gradient computations, even if the QNode has been re-evaluated
        with a different number of shots in the meantime."""
        dev = qml.device("default.qubit", wires=2, shots=None)
        weights = torch.tensor([0.543, -0.654], requires_grad=True)
        a, b = weights

        @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift)
        def circuit(a, b):
            qml.RY(a, wires=0)
            qml.RX(b, wires=1)
            qml.CNOT(wires=[0, 1])
            return qml.expval(qml.PauliY(1))

        res1 = circuit(*weights)
        assert qml.math.shape(res1) == tuple()

        res2 = circuit(*weights, shots=[(1, 1000)])
        assert qml.math.shape(res2) == (1000,)

        res1.backward()

        expected = torch.tensor([torch.sin(a) * torch.sin(b), -torch.cos(a) * torch.cos(b)])
        assert torch.allclose(weights.grad, expected, atol=tol, rtol=0)

    def test_update_diff_method(self, mocker, tol):
        """Test that temporarily setting the shots updates the diff method"""
        dev = qml.device("default.qubit", wires=2, shots=100)
        a, b = torch.tensor([0.543, -0.654], requires_grad=True)

        spy = mocker.spy(qml, "execute")

        @qnode(dev, interface="torch")
        def cost_fn(a, b):
            qml.RY(a, wires=0)
            qml.RX(b, wires=1)
            qml.CNOT(wires=[0, 1])
            return qml.expval(qml.PauliY(1))

        # since we are using finite shots, parameter-shift will
        # be chosen
        assert cost_fn.gradient_fn is qml.gradients.param_shift

        cost_fn(a, b)
        assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift

        # if we set the shots to None, backprop can now be used
        cost_fn(a, b, shots=None)
        assert spy.call_args[1]["gradient_fn"] == "backprop"

        # original QNode settings are unaffected
        assert cost_fn.gradient_fn is qml.gradients.param_shift
        cost_fn(a, b)
        assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift

class TestAdjoint:
    """Specific integration tests for the adjoint method"""

    def test_reuse_state(self, mocker):
        """Tests that the Torch interface reuses the device state for adjoint differentiation"""
        dev = qml.device("default.qubit", wires=2)

        @qnode(dev, diff_method="adjoint", interface="torch")
        def circ(x):
            qml.RX(x[0], wires=0)
            qml.RY(x[1], wires=1)
            qml.CNOT(wires=(0, 1))
            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliX(1))

        expected_grad = lambda x: torch.tensor([-torch.sin(x[0]), torch.cos(x[1])])

        spy = mocker.spy(dev, "adjoint_jacobian")

        x1 = torch.tensor([0.1, 0.2], requires_grad=True)
        res1 = circ(x1)
        res1.backward(torch.Tensor([1, 1]))

        assert np.allclose(x1.grad, expected_grad(x1))
        assert circ.device.num_executions == 1
        spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY)

    def test_resuse_state_multiple_evals(self, mocker, tol):
        """Tests that the Torch interface reuses the device state for adjoint differentiation,
        even where there are intermediate evaluations."""
        dev = qml.device("default.qubit", wires=2)

        x_val = 0.543
        y_val = -0.654
        x = torch.tensor(x_val, requires_grad=True)
        y = torch.tensor(y_val, requires_grad=True)

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

        spy = mocker.spy(dev, "adjoint_jacobian")

        res1 = circuit(x, y)
        assert np.allclose(res1.detach(), np.cos(x_val), atol=tol, rtol=0)

        # intermediate evaluation with different values
        res2 = circuit(torch.tan(x), torch.cosh(y))

        # the adjoint method will continue to compute the correct derivative
        res1.backward()
        assert np.allclose(x.grad.detach(), -np.sin(x_val), atol=tol, rtol=0)
        assert dev.num_executions == 2
        spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY)

@pytest.mark.parametrize("dev_name,diff_method,mode", qubit_device_and_diff_method)
class TestQubitIntegration:
    """Tests that ensure various qubit circuits integrate correctly"""

    def test_probability_differentiation(self, dev_name, diff_method, mode, tol):
        """Tests correct output shape and evaluation for a tape
        with prob and expval outputs"""

        if diff_method == "adjoint":
            pytest.skip("The adjoint method does not currently support returning probabilities")

        dev = qml.device(dev_name, wires=2)
        x_val = 0.543
        y_val = -0.654
        x = torch.tensor(x_val, requires_grad=True, dtype=torch.float64)
        y = torch.tensor(y_val, requires_grad=True, dtype=torch.float64)

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

        res = circuit(x, y)

        expected = np.array(
            [
                [np.cos(x_val / 2) ** 2, np.sin(x_val / 2) ** 2],
                [
                    (1 + np.cos(x_val) * np.cos(y_val)) / 2,
                    (1 - np.cos(x_val) * np.cos(y_val)) / 2,
                ],
            ]
        )

        if diff_method == "backprop":
            # TODO: check why this differs from other interfaces
            # https://github.com/PennyLaneAI/pennylane/issues/1607
            expected = expected.flatten()

        assert np.allclose(res.detach().numpy(), expected, atol=tol, rtol=0)

        loss = torch.sum(res)
        loss.backward()
        expected = np.array(
            [
                -np.sin(x_val) / 2
                + np.sin(x_val) / 2
                - np.sin(x_val) * np.cos(y_val) / 2
                + np.cos(y_val) * np.sin(x_val) / 2,
                -np.cos(x_val) * np.sin(y_val) / 2 + np.cos(x_val) * np.sin(y_val) / 2,
            ]
        )
        assert np.allclose(x.grad, expected[0], atol=tol, rtol=0)
        assert np.allclose(y.grad, expected[1], atol=tol, rtol=0)

    def test_ragged_differentiation(self, dev_name, diff_method, mode, monkeypatch, tol):
        """Tests correct output shape and evaluation for a tape
        with prob and expval outputs"""
        if diff_method == "adjoint":
            pytest.skip("The adjoint method does not currently support returning probabilities")

        dev = qml.device(dev_name, wires=2)
        x_val = 0.543
        y_val = -0.654
        x = torch.tensor(x_val, requires_grad=True, dtype=torch.float64)
        y = torch.tensor(y_val, requires_grad=True, dtype=torch.float64)

        @qnode(dev, diff_method=diff_method, mode=mode, interface="torch")
        def circuit(x, y):
            qml.RX(x, wires=[0])
            qml.RY(y, wires=[1])
            qml.CNOT(wires=[0, 1])
            return [qml.expval(qml.PauliZ(0)), qml.probs(wires=[1])]

        res = circuit(x, y)

        expected = np.array(
            [
                np.cos(x_val),
                (1 + np.cos(x_val) * np.cos(y_val)) / 2,
                (1 - np.cos(x_val) * np.cos(y_val)) / 2,
            ]
        )
        assert np.allclose(res.detach().numpy(), expected, atol=tol, rtol=0)

        loss = torch.sum(res)
        loss.backward()
        expected = np.array(
            [
                -np.sin(x_val)
                + -np.sin(x_val) * np.cos(y_val) / 2
                + np.cos(y_val) * np.sin(x_val) / 2,
                -np.cos(x_val) * np.sin(y_val) / 2 + np.cos(x_val) * np.sin(y_val) / 2,
            ]
        )
        assert np.allclose(x.grad, expected[0], atol=tol, rtol=0)
        assert np.allclose(y.grad, expected[1], atol=tol, rtol=0)

    def test_chained_qnodes(self, dev_name, diff_method, mode):
        """Test that the gradient of chained QNodes works without error"""
        dev = qml.device(dev_name, wires=2)

        @qnode(dev, interface="torch", diff_method=diff_method, mode=mode)
        def circuit1(weights):
            qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1])
            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))

        @qnode(dev, interface="torch", diff_method=diff_method, mode=mode)
        def circuit2(data, weights):
            qml.templates.AngleEmbedding(data, wires=[0, 1])
            qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1])
            return qml.expval(qml.PauliX(0))

        def cost(weights):
            w1, w2 = weights
            c1 = circuit1(w1)
            c2 = circuit2(c1, w2)
            return torch.sum(c2) ** 2

        w1 = np.random.random(qml.templates.StronglyEntanglingLayers.shape(3, 2))
        w2 = np.random.random(qml.templates.StronglyEntanglingLayers.shape(4, 2))

        w1 = torch.tensor(w1, requires_grad=True)
        w2 = torch.tensor(w2, requires_grad=True)

        weights = [w1, w2]

        loss = cost(weights)
        loss.backward()

    def test_hessian(self, dev_name, diff_method, mode, tol):
        """Test hessian calculation of a scalar valued QNode"""
        if diff_method not in {"parameter-shift", "backprop"}:
            pytest.skip("Test only supports parameter-shift or backprop")

        dev = qml.device(dev_name, wires=1)

        @qnode(dev, diff_method=diff_method, mode=mode, max_diff=2, interface="torch")
        def circuit(x):
            qml.RY(x[0], wires=0)
            qml.RX(x[1], wires=0)
            return qml.expval(qml.PauliZ(0))

        x = torch.tensor([1.0, 2.0], requires_grad=True)
        res = circuit(x)

        res.backward()
        g = x.grad

        hess = hessian(circuit, x)
        a, b = x.detach().numpy()

        expected_res = np.cos(a) * np.cos(b)
        assert np.allclose(res.detach(), expected_res, atol=tol, rtol=0)

        expected_g = [-np.sin(a) * np.cos(b), -np.cos(a) * np.sin(b)]
        assert np.allclose(g.detach(), expected_g, atol=tol, rtol=0)

        expected_hess = [
            [-np.cos(a) * np.cos(b), np.sin(a) * np.sin(b)],
            [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)],
        ]
        assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0)

    def test_hessian_vector_valued(self, dev_name, diff_method, mode, tol):
        """Test hessian calculation of a vector valued QNode"""
        if diff_method not in {"parameter-shift", "backprop"}:
            pytest.skip("Test only supports parameter-shift or backprop")

        dev = qml.device(dev_name, wires=1)

        @qnode(dev, diff_method=diff_method, mode=mode, max_diff=2, interface="torch")
        def circuit(x):
            qml.RY(x[0], wires=0)
            qml.RX(x[1], wires=0)
            return qml.probs(wires=0)

        x = torch.tensor([1.0, 2.0], requires_grad=True)
        res = circuit(x)
        jac_fn = lambda x: jacobian(circuit, x, create_graph=True)

        g = jac_fn(x)
        hess = jacobian(jac_fn, x)
        a, b = x.detach().numpy()

        expected_res = [
            0.5 + 0.5 * np.cos(a) * np.cos(b),
            0.5 - 0.5 * np.cos(a) * np.cos(b),
        ]
        assert np.allclose(res.detach(), expected_res, atol=tol, rtol=0)

        expected_g = [
            [-0.5 * np.sin(a) * np.cos(b), -0.5 * np.cos(a) * np.sin(b)],
            [0.5 * np.sin(a) * np.cos(b), 0.5 * np.cos(a) * np.sin(b)],
        ]
        assert np.allclose(g.detach(), expected_g, atol=tol, rtol=0)

        expected_hess = [
            [
                [-0.5 * np.cos(a) * np.cos(b), 0.5 * np.sin(a) * np.sin(b)],
                [0.5 * np.sin(a) * np.sin(b), -0.5 * np.cos(a) * np.cos(b)],
            ],
            [
                [0.5 * np.cos(a) * np.cos(b), -0.5 * np.sin(a) * np.sin(b)],
                [-0.5 * np.sin(a) * np.sin(b), 0.5 * np.cos(a) * np.cos(b)],
            ],
        ]
        assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0)

    def test_hessian_ragged(self, dev_name, diff_method, mode, tol):
        """Test hessian calculation of a ragged QNode"""
        if diff_method not in {"parameter-shift", "backprop"}:
            pytest.skip("Test only supports parameter-shift or backprop")

        dev = qml.device(dev_name, wires=2)

        @qnode(dev, diff_method=diff_method, mode=mode, max_diff=2, interface="torch")
        def circuit(x):
            qml.RY(x[0], wires=0)
            qml.RX(x[1], wires=0)
            qml.RY(x[0], wires=1)
            qml.RX(x[1], wires=1)
            return qml.expval(qml.PauliZ(0)), qml.probs(wires=1)

        x = torch.tensor([1.0, 2.0], requires_grad=True)
        res = circuit(x)
        jac_fn = lambda x: jacobian(circuit, x, create_graph=True)

        g = jac_fn(x)
        hess = jacobian(jac_fn, x)
        a, b = x.detach().numpy()

        expected_res = [
            np.cos(a) * np.cos(b),
            0.5 + 0.5 * np.cos(a) * np.cos(b),
            0.5 - 0.5 * np.cos(a) * np.cos(b),
        ]
        assert np.allclose(res.detach(), expected_res, atol=tol, rtol=0)

        expected_g = [
            [-np.sin(a) * np.cos(b), -np.cos(a) * np.sin(b)],
            [-0.5 * np.sin(a) * np.cos(b), -0.5 * np.cos(a) * np.sin(b)],
            [0.5 * np.sin(a) * np.cos(b), 0.5 * np.cos(a) * np.sin(b)],
        ]
        assert np.allclose(g.detach(), expected_g, atol=tol, rtol=0)

        expected_hess = [
            [
                [-np.cos(a) * np.cos(b), np.sin(a) * np.sin(b)],
                [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)],
            ],
            [
                [-0.5 * np.cos(a) * np.cos(b), 0.5 * np.sin(a) * np.sin(b)],
                [0.5 * np.sin(a) * np.sin(b), -0.5 * np.cos(a) * np.cos(b)],
            ],
            [
                [0.5 * np.cos(a) * np.cos(b), -0.5 * np.sin(a) * np.sin(b)],
                [-0.5 * np.sin(a) * np.sin(b), 0.5 * np.cos(a) * np.cos(b)],
            ],
        ]
        assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0)

    def test_hessian_vector_valued_postprocessing(self, dev_name, diff_method, mode, tol):
        """Test hessian calculation of a vector valued QNode with post-processing"""
        if diff_method not in {"parameter-shift", "backprop"}:
            pytest.skip("Test only supports parameter-shift or backprop")

        dev = qml.device(dev_name, wires=1)

        @qnode(dev, diff_method=diff_method, mode=mode, max_diff=2, interface="torch")
        def circuit(x):
            qml.RX(x[0], wires=0)
            qml.RY(x[1], wires=0)
            return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(0))]

        x = torch.tensor([0.76, -0.87], requires_grad=True, dtype=torch.float64)

        def cost_fn(x):
            return x @ circuit(x)

        a, b = x.detach().numpy()

        res = cost_fn(x)
        expected_res = np.array([a, b]) @ [np.cos(a) * np.cos(b), np.cos(a) * np.cos(b)]
        assert np.allclose(res.detach(), expected_res, atol=tol, rtol=0)

        res.backward()

        g = x.grad
        expected_g = [
            np.cos(b) * (np.cos(a) - (a + b) * np.sin(a)),
            np.cos(a) * (np.cos(b) - (a + b) * np.sin(b)),
        ]
        assert np.allclose(g.detach(), expected_g, atol=tol, rtol=0)

        hess = hessian(cost_fn, x)
        expected_hess = [
            [
                -(np.cos(b) * ((a + b) * np.cos(a) + 2 * np.sin(a))),
                -(np.cos(b) * np.sin(a)) + (-np.cos(a) + (a + b) * np.sin(a)) * np.sin(b),
            ],
            [
                -(np.cos(b) * np.sin(a)) + (-np.cos(a) + (a + b) * np.sin(a)) * np.sin(b),
                -(np.cos(a) * ((a + b) * np.cos(b) + 2 * np.sin(b))),
            ],
        ]

        assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0)

    def test_state(self, dev_name, diff_method, mode, tol):
        """Test that the state can be returned and differentiated"""
        if diff_method == "adjoint":
            pytest.skip("Adjoint does not support states")

        dev = qml.device(dev_name, wires=2)

        x = torch.tensor(0.543, requires_grad=True)
        y = torch.tensor(-0.654, requires_grad=True)

        @qnode(dev, diff_method=diff_method, interface="torch", mode=mode)
        def circuit(x, y):
            qml.RX(x, wires=[0])
            qml.RY(y, wires=[1])
            qml.CNOT(wires=[0, 1])
            return qml.state()

        def cost_fn(x, y):
            res = circuit(x, y)
            assert res.dtype is torch.complex128
            probs = torch.abs(res) ** 2
            return probs[0] + probs[2]

        res = cost_fn(x, y)

        if diff_method not in {"backprop"}:
            pytest.skip("Test only supports backprop")

        res.backward()
        res = torch.tensor([x.grad, y.grad])
        expected = torch.tensor(
            [-torch.sin(x) * torch.cos(y) / 2, -torch.cos(x) * torch.sin(y) / 2]
        )
        assert torch.allclose(res, expected, atol=tol, rtol=0)

    def test_projector(self, dev_name, diff_method, mode, tol):
        """Test that the variance of a projector is correctly returned"""
        if diff_method == "adjoint":
            pytest.skip("Adjoint does not support projectors")

        dev = qml.device(dev_name, wires=2)
        P = torch.tensor([1], requires_grad=False)

        x, y = 0.765, -0.654
        weights = torch.tensor([x, y], requires_grad=True, dtype=torch.float64)

        @qnode(dev, diff_method=diff_method, interface="torch", mode=mode)
        def circuit(x, y):
            qml.RX(x, wires=0)
            qml.RY(y, wires=1)
            qml.CNOT(wires=[0, 1])
            return qml.var(qml.Projector(P, wires=0) @ qml.PauliX(1))

        res = circuit(*weights)
        expected = 0.25 * np.sin(x / 2) ** 2 * (3 + np.cos(2 * y) + 2 * np.cos(x) * np.sin(y) ** 2)
        assert np.allclose(res.detach(), expected, atol=tol, rtol=0)

        res.backward()
        expected = np.array(
            [
                [
                    0.5 * np.sin(x) * (np.cos(x / 2) ** 2 + np.cos(2 * y) * np.sin(x / 2) ** 2),
                    -2 * np.cos(y) * np.sin(x / 2) ** 4 * np.sin(y),
                ]
            ]
        )
        assert np.allclose(weights.grad.detach(), expected, atol=tol, rtol=0)

@pytest.mark.parametrize(
    "diff_method,kwargs",
    [["finite-diff", {}], ("parameter-shift", {}), ("parameter-shift", {"force_order2": True})],
)
class TestCV:
    """Tests for CV integration"""

    def test_first_order_observable(self, diff_method, kwargs, tol):
        """Test variance of a first order CV observable"""
        dev = qml.device("default.gaussian", wires=1)

        r = torch.tensor(0.543, dtype=torch.float64, requires_grad=True)
        phi = torch.tensor(-0.654, dtype=torch.float64, requires_grad=True)

        @qnode(dev, interface="torch", diff_method=diff_method, **kwargs)
        def circuit(r, phi):
            qml.Squeezing(r, 0, wires=0)
            qml.Rotation(phi, wires=0)
            return qml.var(qml.X(0))

        res = circuit(r, phi)
        expected = torch.exp(2 * r) * torch.sin(phi) ** 2 + torch.exp(-2 * r) * torch.cos(phi) ** 2
        assert torch.allclose(res, expected, atol=tol, rtol=0)

        # circuit jacobians
        res.backward()
        res = torch.tensor([r.grad, phi.grad])
        expected = torch.tensor(
            [
                [
                    2 * torch.exp(2 * r) * torch.sin(phi) ** 2
                    - 2 * torch.exp(-2 * r) * torch.cos(phi) ** 2,
                    2 * torch.sinh(2 * r) * torch.sin(2 * phi),
                ]
            ]
        )
        assert torch.allclose(res, expected, atol=tol, rtol=0)

    def test_second_order_observable(self, diff_method, kwargs, tol):
        """Test variance of a second order CV expectation value"""
        dev = qml.device("default.gaussian", wires=1)

        n = torch.tensor(0.12, dtype=torch.float64, requires_grad=True)
        a = torch.tensor(0.765, dtype=torch.float64, requires_grad=True)

        @qnode(dev, interface="torch", diff_method=diff_method, **kwargs)
        def circuit(n, a):
            qml.ThermalState(n, wires=0)
            qml.Displacement(a, 0, wires=0)
            return qml.var(qml.NumberOperator(0))

        res = circuit(n, a)
        expected = n**2 + n + torch.abs(a) ** 2 * (1 + 2 * n)
        assert torch.allclose(res, expected, atol=tol, rtol=0)

        # circuit jacobians
        res.backward()
        res = torch.tensor([n.grad, a.grad])
        expected = torch.tensor([[2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]])
        assert torch.allclose(res, expected, atol=tol, rtol=0)

@pytest.mark.parametrize("dev_name,diff_method,mode", qubit_device_and_diff_method)
class TestTapeExpansion:
    """Test that tape expansion within the QNode integrates correctly
    with the Torch interface"""

    def test_gradient_expansion(self, dev_name, diff_method, mode, mocker):
        """Test that a *supported* operation with no gradient recipe is
        expanded for both parameter-shift and finite-differences, but not for execution."""
        if diff_method not in ("parameter-shift", "finite-diff"):
            pytest.skip("Only supports gradient transforms")

        dev = qml.device(dev_name, wires=1)

        class PhaseShift(qml.PhaseShift):
            grad_method = None

            def expand(self):
                with qml.tape.QuantumTape() as tape:
                    qml.RY(3 * self.data[0], wires=self.wires)
                return tape

        @qnode(dev, diff_method=diff_method, mode=mode, max_diff=2, interface="torch")
        def circuit(x):
            qml.Hadamard(wires=0)
            PhaseShift(x, wires=0)
            return qml.expval(qml.PauliX(0))

        spy = mocker.spy(circuit.device, "batch_execute")
        x = torch.tensor(0.5, requires_grad=True)

        loss = circuit(x)

        tape = spy.call_args[0][0][0]

        spy = mocker.spy(circuit.gradient_fn, "transform_fn")
        loss.backward()
        res = x.grad

        input_tape = spy.call_args[0][0]
        assert len(input_tape.operations) == 2
        assert input_tape.operations[1].name == "RY"
        assert input_tape.operations[1].data[0] == 3 * x

        shifted_tape1, shifted_tape2 = spy.spy_return[0]

        assert len(shifted_tape1.operations) == 2
        assert shifted_tape1.operations[1].name == "RY"

        assert len(shifted_tape2.operations) == 2
        assert shifted_tape2.operations[1].name == "RY"

        assert torch.allclose(res, -3 * torch.sin(3 * x))

        if diff_method == "parameter-shift":
            # test second order derivatives
            res = torch.autograd.functional.hessian(circuit, x)
            assert torch.allclose(res, -9 * torch.cos(3 * x))

    @pytest.mark.parametrize("max_diff", [1, 2])
    def test_gradient_expansion_trainable_only(self, dev_name, diff_method, mode, max_diff, mocker):
        """Test that a *supported* operation with no gradient recipe is only
        expanded for parameter-shift and finite-differences when it is trainable."""
        if diff_method not in ("parameter-shift", "finite-diff"):
            pytest.skip("Only supports gradient transforms")

        dev = qml.device(dev_name, wires=1)

        class PhaseShift(qml.PhaseShift):
            grad_method = None

            def expand(self):
                with qml.tape.QuantumTape() as tape:
                    qml.RY(3 * self.data[0], wires=self.wires)
                return tape

        @qnode(dev, diff_method=diff_method, mode=mode, max_diff=max_diff, interface="torch")
        def circuit(x, y):
            qml.Hadamard(wires=0)
            PhaseShift(x, wires=0)
            PhaseShift(2 * y, wires=0)
            return qml.expval(qml.PauliX(0))

        spy = mocker.spy(circuit.device, "batch_execute")
        x = torch.tensor(0.5, requires_grad=True)
        y = torch.tensor(0.7, requires_grad=False)

        loss = circuit(x, y)

        spy = mocker.spy(circuit.gradient_fn, "transform_fn")
        loss.backward()

        input_tape = spy.call_args[0][0]
        assert len(input_tape.operations) == 3
        assert input_tape.operations[1].name == "RY"
        assert input_tape.operations[1].data[0] == 3 * x
        assert input_tape.operations[2].name == "PhaseShift"
        assert input_tape.operations[2].grad_method is None

    @pytest.mark.parametrize("max_diff", [1, 2])
    def test_hamiltonian_expansion_analytic(self, dev_name, diff_method, mode, max_diff):
        """Test that if there
        are non-commuting groups and the number of shots is None
        the first and second order gradients are correctly evaluated"""
        if diff_method == "adjoint":
            pytest.skip("The adjoint method does not yet support Hamiltonians")

        dev = qml.device(dev_name, wires=3, shots=None)
        obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)]

        @qnode(dev, diff_method=diff_method, mode=mode, max_diff=max_diff, interface="torch")
        def circuit(data, weights, coeffs):
            weights = torch.reshape(weights, [1, -1])
            qml.templates.AngleEmbedding(data, wires=[0, 1])
            qml.templates.BasicEntanglerLayers(weights, wires=[0, 1])
            return qml.expval(qml.Hamiltonian(coeffs, obs))

        d = torch.tensor([0.1, 0.2], requires_grad=False, dtype=torch.float64)
        w = torch.tensor([0.654, -0.734], requires_grad=True, dtype=torch.float64)
        c = torch.tensor([-0.6543, 0.24, 0.54], requires_grad=True, dtype=torch.float64)

        # test output
        res = circuit(d, w, c)

        expected = c[2] * torch.cos(d[1] + w[1]) - c[1] * torch.sin(d[0] + w[0]) * torch.sin(
            d[1] + w[1]
        )
        assert torch.allclose(res, expected)

        # test gradients
        res.backward()
        grad = (w.grad, c.grad)

        expected_w = torch.tensor(
            [
                -c[1] * torch.cos(d[0] + w[0]) * torch.sin(d[1] + w[1]),
                -c[1] * torch.cos(d[1] + w[1]) * torch.sin(d[0] + w[0])
                - c[2] * torch.sin(d[1] + w[1]),
            ]
        )
        expected_c = torch.tensor(
            [0, -torch.sin(d[0] + w[0]) * torch.sin(d[1] + w[1]), torch.cos(d[1] + w[1])]
        )
        assert torch.allclose(grad[0], expected_w)
        assert torch.allclose(grad[1], expected_c)

        # test second-order derivatives
        if diff_method in ("parameter-shift", "backprop") and max_diff == 2:
            hessians = torch.autograd.functional.hessian(circuit, (d, w, c))

            grad2_c = hessians[2][2]
            assert torch.allclose(grad2_c, torch.zeros([3, 3], dtype=torch.float64))

            grad2_w_c = hessians[1][2]
            expected = torch.tensor(
                [
                    [0, -torch.cos(d[0] + w[0]) * torch.sin(d[1] + w[1]), 0],
                    [
                        0,
                        -torch.cos(d[1] + w[1]) * torch.sin(d[0] + w[0]),
                        -torch.sin(d[1] + w[1]),
                    ],
                ]
            )
            assert torch.allclose(grad2_w_c, expected)

    @pytest.mark.parametrize("max_diff", [1, 2])
    def test_hamiltonian_expansion_finite_shots(
        self, dev_name, diff_method, mode, max_diff, mocker
    ):
        """Test that the Hamiltonian is expanded if there
        are non-commuting groups and the number of shots is finite
        and the first and second order gradients are correctly evaluated"""
        if diff_method in ("adjoint", "backprop", "finite-diff"):
            pytest.skip("The adjoint and backprop methods do not yet support sampling")

        dev = qml.device(dev_name, wires=3, shots=50000)
        spy = mocker.spy(qml.transforms, "hamiltonian_expand")
        obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)]

        @qnode(dev, diff_method=diff_method, mode=mode, max_diff=max_diff, interface="torch")
        def circuit(data, weights, coeffs):
            weights = torch.reshape(weights, [1, -1])
            qml.templates.AngleEmbedding(data, wires=[0, 1])
            qml.templates.BasicEntanglerLayers(weights, wires=[0, 1])
            H = qml.Hamiltonian(coeffs, obs)
            H.compute_grouping()
            return qml.expval(H)

        d = torch.tensor([0.1, 0.2], requires_grad=False, dtype=torch.float64)
        w = torch.tensor([0.654, -0.734], requires_grad=True, dtype=torch.float64)
        c = torch.tensor([-0.6543, 0.24, 0.54], requires_grad=True, dtype=torch.float64)

        # test output
        res = circuit(d, w, c)

        expected = c[2] * torch.cos(d[1] + w[1]) - c[1] * torch.sin(d[0] + w[0]) * torch.sin(
            d[1] + w[1]
        )
        assert torch.allclose(res, expected, atol=0.1)
        spy.assert_called()

        # test gradients
        res.backward()
        grad = (w.grad, c.grad)

        expected_w = torch.tensor(
            [
                -c[1] * torch.cos(d[0] + w[0]) * torch.sin(d[1] + w[1]),
                -c[1] * torch.cos(d[1] + w[1]) * torch.sin(d[0] + w[0])
                - c[2] * torch.sin(d[1] + w[1]),
            ]
        )
        expected_c = torch.tensor(
            [0, -torch.sin(d[0] + w[0]) * torch.sin(d[1] + w[1]), torch.cos(d[1] + w[1])]
        )
        assert torch.allclose(grad[0], expected_w, atol=0.1)
        assert torch.allclose(grad[1], expected_c, atol=0.1)

        # test second-order derivatives
        if diff_method == "parameter-shift" and max_diff == 2:
            hessians = torch.autograd.functional.hessian(circuit, (d, w, c))

            grad2_c = hessians[2][2]
            assert torch.allclose(grad2_c, torch.zeros([3, 3], dtype=torch.float64), atol=0.1)

            grad2_w_c = hessians[1][2]
            expected = torch.tensor(
                [
                    [0, -torch.cos(d[0] + w[0]) * torch.sin(d[1] + w[1]), 0],
                    [
                        0,
                        -torch.cos(d[1] + w[1]) * torch.sin(d[0] + w[0]),
                        -torch.sin(d[1] + w[1]),
                    ],
                ]
            )
            assert torch.allclose(grad2_w_c, expected, atol=0.1)

class TestSample:
    """Tests for the sample integration"""

    def test_sample_dimension(self):
        """Test sampling works as expected"""
        dev = qml.device("default.qubit", wires=2, shots=10)

        @qnode(dev, diff_method="parameter-shift", interface="torch")
        def circuit():
            qml.Hadamard(wires=[0])
            qml.CNOT(wires=[0, 1])
            return [qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1))]

        res = circuit()

        assert res.shape == (2, 10)
        assert isinstance(res, torch.Tensor)

    def test_sampling_expval(self):
        """Test sampling works as expected if combined with expectation values"""
        dev = qml.device("default.qubit", wires=2, shots=10)

        @qnode(dev, diff_method="parameter-shift", interface="torch")
        def circuit():
            qml.Hadamard(wires=[0])
            qml.CNOT(wires=[0, 1])
            return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1))

        res = circuit()

        assert len(res) == 2
        assert isinstance(res, tuple)
        assert res[0].shape == (10,)
        assert isinstance(res[0], torch.Tensor)
        assert isinstance(res[1], torch.Tensor)

    def test_sample_combination(self, tol):
        """Test the output of combining expval, var and sample"""
        n_sample = 10

        dev = qml.device("default.qubit", wires=3, shots=n_sample)

        @qnode(dev, diff_method="parameter-shift", interface="torch")
        def circuit():
            qml.RX(0.54, wires=0)

            return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)), qml.var(qml.PauliY(2))

        result = circuit()

        assert len(result) == 3
        assert np.array_equal(result[0].shape, (n_sample,))
        assert isinstance(result[1], torch.Tensor)
        assert isinstance(result[2], torch.Tensor)
        assert result[0].dtype is torch.int64

    def test_single_wire_sample(self, tol):
        """Test the return type and shape of sampling a single wire"""
        n_sample = 10

        dev = qml.device("default.qubit", wires=1, shots=n_sample)

        @qnode(dev, diff_method="parameter-shift", interface="torch")
        def circuit():
            qml.RX(0.54, wires=0)

            return qml.sample(qml.PauliZ(0))

        result = circuit()

        assert isinstance(result, torch.Tensor)
        assert np.array_equal(result.shape, (n_sample,))

    def test_multi_wire_sample_regular_shape(self, tol):
        """Test the return type and shape of sampling multiple wires
        where a rectangular array is expected"""
        n_sample = 10

        dev = qml.device("default.qubit", wires=3, shots=n_sample)

        @qnode(dev, diff_method="parameter-shift", interface="torch")
        def circuit():
            return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2))

        result = circuit()

        # If all the dimensions are equal the result will end up to be a proper rectangular array
        assert isinstance(result, torch.Tensor)
        assert np.array_equal(result.shape, (3, n_sample))
        assert result.dtype == torch.int64


# MULTI-INDEX DATAFRAMES IN PANDAS (AKA: "INDEX-HIGHER KEY" )
# Santiago Garcia Arango, June 2020

import numpy as np
import pandas as pd

# ---------------MULTI-INDEX DATAFRAME PANDAS--------------------
# With Pandas, we can achieve a multi-level indexing layout, where...
# ... the DataFrame has N levels of indexing. This is really usefull...
# ... when working with multiple parameters and sub-parameters of info.

# First index level (that will contain the second one)...
outside = ["G1", "G1", "G1", "G2", "G2", "G2", "G3", "G3", "G3"]  # List 1

# Second index level (that are "inside" the First index level)...
inside = [1, 2, 3, 1, 2, 3, 1, 2, 3]  # List 2

# We create the "interconnection" of these two index levels...
# Remark: this must be done in a list of tuples (to work properly)
h_index = list(zip(outside, inside))  # List of both lists (in tuples)
h_index = pd.MultiIndex.from_tuples(h_index)  # Create MultiIndex DF
print("h_index =\n", h_index, "\n")

# Now, we can go ahead and create the MultiIndex Dataframe...
data_frame = pd.DataFrame(np.random.randn(9, 3), h_index, ['A', 'B', 'C'])
print("data_frame =\n", data_frame, "\n")

# Extra/optional step (naming the index-levels)...
data_frame.index.names = ["Groups", "Numbers"]
print("data_frame (with index names) =\n", data_frame, "\n")

# -------------SEARCHING FOR MULTI-INDEX DATAFRAMES----------------
# Searching first index level...
#  1) First way of doing this (access corresponding index):
print("data_frame.loc['G1']=\n", data_frame.loc['G1'], "\n")
#  2) Second way of doing this (access cross_section of index):
print("data_frame.xs('G1')=\n", data_frame.xs('G1'), "\n")

# Searching second index level...
print("data_frame.loc['G1'].loc[1]=\n", data_frame.loc['G1'].loc[1], "\n")

# Searching general condition for a "N level" index...
# Note: will search in Any 1st level, only focusing on 2nd level condition
print("data_frame.xs(2, level='Numbers') =")
print(data_frame.xs(2, level='Numbers'))


import time
import pandas as pd
import numpy as np

CITY_DATA = { 'chicago': 'chicago.csv',
              'new york city': 'new_york_city.csv',
              'washington': 'washington.csv' }

def get_filters():
    """
    Asks user to specify a city, month, and day to analyze.

    Returns:
        (str) city - name of the city to analyze
        (str) month - name of the month to filter by, or "all" to apply no month filter
        (str) day - name of the day of week to filter by, or "all" to apply no day filter
    """
    print('\nHello! Let\'s explore some US bikeshare data!\n')

    # get user input for city (chicago, new york city, washington). HINT: Use a while loop to handle invalid inputs
    while True:
        city = str(input('Please enter a city to analyze (Chicago, New Yor City, Washington): ')).lower()
        # checking, if the correct selection is entered
        if city in ('chicago', 'new york city', 'washington'):
            break
        else:
            print('Wrong entry. Please enter the name of the cities.')

    # get user input for month (all, january, february, ... , june)
    while True:
        month = str(input('Enter a month to filter for (all, january, february, ... , june): ')).lower()
        # checking, if the correct selection is entered
        if month in ('all', 'january', 'february', 'march', 'april', 'may', 'june'):
            break
        else:
            print('Wrong entry. Please enter the name of a month.')

    # get user input for day of week (all, monday, tuesday, ... sunday)
    while True:
        day = str(input('Enter a day to filter for (all, monday, tuesday, ... sunday): ')).lower()
        # checking, if the correct selection is entered
        if day in ('all', 'monday', 'tuesday', 'wednesday', 'thursday', 'friday', 'saturday', 'sunday'):
            break
        else:
            print('Wrong entry. Please enter the name of a day.')

    print('-'*40)
    return city, month, day

def load_data(city, month, day):
    """
    Loads data for the specified city and filters by month and day if applicable.

    Args:
        (str) city - name of the city to analyze
        (str) month - name of the month to filter by, or "all" to apply no month filter
        (str) day - name of the day of week to filter by, or "all" to apply no day filter
    Returns:
        df - Pandas DataFrame containing city data filtered by month and day
    """
    # load data file into a dataframe
    df = pd.read_csv(CITY_DATA[city])

    # convert the Start Time column to datetime
    df['Start Time'] = pd.to_datetime(df['Start Time'])

    # extract month and day of week from Start Time to create new columns
    df['month'] = df['Start Time'].dt.month_name
    df['day'] = df['Start Time'].dt.weekday_name

    # filter by month if applicable
    if month != 'all':
        # use the index of the months list to get the corresponding int
        months = ['january', 'february', 'march', 'april', 'may', 'june']
        month = months.index(month) + 1

        # filter by month to create the new dataframe
        df = df[df['month'] == month]

    # filter by day of week if applicable
    if day != 'all':
        # filter by day of week to create the new dataframe
        df = df[df['day'] == day.title()]

    return df

def time_stats(df):
    """Displays statistics on the most frequent times of travel."""

    print('\nCalculating The Most Frequent Times of Travel...\n')
    start_time = time.time()

    # display the most common month
    print('The most common month is: ', df['month'].mode()[0])

    # display the most common day of week
    print('The most common day of the week is: ', df['day'].mode()[0])

    # display the most common start hour
    df['hour'] = df['Start Time'].dt.hour
    print('The most common start hour is: ', df['hour'].mode()[0])

    print("\nThis took %s seconds." % round((time.time() - start_time), 2))
    print('-'*40)

def station_stats(df):
    """Displays statistics on the most popular stations and trip."""

    print('\nCalculating The Most Popular Stations and Trip...\n')
    start_time = time.time()

    # display most commonly used start station
    print('Most commonly used start station: ', df['Start Station'].mode()[0])

    # display most commonly used end station
    print('Most commonly used end station: ', df['End Station'].mode()[0])

    # display most frequent combination of start station and end station trip
    df['Start to End Station'] = df['Start Station'] + ' to ' + df['End Station']
    print('Most frequent combination of start station and end station trip: ', df['Start to End Station'].mode()[0])

    print("\nThis took %s seconds." % round((time.time() - start_time), 2))
    print('-'*40)

def trip_duration_stats(df):
    """Displays statistics on the total and average trip duration."""

    print('\nCalculating Trip Duration...\n')
    start_time = time.time()

    df['Start Time'] = pd.to_datetime(df['Start Time'])
    df['End Time'] = pd.to_datetime(df['End Time'])
    df['Travel Time'] = df['End Time'] - df['Start Time']

    # display total travel time
    print('Total travel time: ', df['Travel Time'].sum())

    # display mean travel time
    print('Mean travel time: ', df['Travel Time'].mean())

    print("\nThis took %s seconds." % round((time.time() - start_time), 2))
    print('-'*40)

def user_stats(df):
    """Displays statistics on bikeshare users."""

    print('\nCalculating User Stats...\n')
    start_time = time.time()

    # Display counts of user types
    if 'User Type' in df.columns:
        print('Counts of user types:')
        print(df['User Type'].value_counts())
    else:
        print('\nNo user tpye data collected.\n')

    # Display counts of gender
    if 'Gender' in df.columns:
        print('\nCounts of gender:')
        print(df['Gender'].value_counts())
    else:
        print('\nNo gender data collected.\n')

    # Display earliest, most recent, and most common year of birth
    if 'Birth Year' in df.columns:
        print('\nEarliest year of birth: ', int(df['Birth Year'].min()))
        print('\nMost recent year of birth: ', int(df['Birth Year'].max()))
        print('\nMost common year of birth: ', int(df['Birth Year'].mode()[0]))
    else:
        print('\nNo birth year collected.\n')

    print("\nThis took %s seconds." % round((time.time() - start_time), 2))
    print('-'*40)

def raw_data(df):
    """Displays the raw data, if user enters 'y' or doesn't display, if user enters 'n'"""

    counter = 0

    while True:
        show_data = str(input('\nWould you like to see the raw data (yes/no)?\n')).lower()
        if show_data in ('yes'):
            print(df[:][counter:counter+5])
            counter += 5
        elif show_data in ('no'):
            break
        else:
            print('Wrong entry.')

def main():
    while True:
        city, month, day = get_filters()
        df = load_data(city, month, day)
        time_stats(df)
        station_stats(df)
        trip_duration_stats(df)
        user_stats(df)
        raw_data(df)
        restart = input('\nWould you like to restart? Enter yes or no.\n')
        if restart.lower() != 'yes':
            break

if __name__ == "__main__":
	main()


"""
Utility functions for working with DataFrames
"""

import pandas as pd
import numpy as np

TEST_DF = pd.DataFrame([1, 2, 3], [4, np.NaN, 2])

# create utility functions for assignment

def nulls(df):
    """Take a DataFrame and outputs a new df of NaN sums by column."""
    return pd.DataFrame(df.isna().sum())

def display_settings(rows, columns):
    """Change the df display settings in a notebook."""
    pd.set_option('display.max_rows', rows)
    pd.set_option('display.max_columns', columns)
    print('Display options change complete.')


"""
Data loader for torch.
"""
# Author: Arturs Berzins <berzins@cats.rwth-aachen.de>
# License: BSD 3 clause

import torch
import torch.utils.data

def create_loader(features_np, targets_np, batch_size, shuffle):
    # https://stackoverflow.com/questions/41924453/pytorch-how-to-use-dataloaders-for-custom-datasets

    # Convert from numpy array to torch tensor
    features = torch.from_numpy(features_np).float()
    targets = torch.from_numpy(targets_np).float()

    # Create dataset from torch tensors
    dataset = torch.utils.data.TensorDataset(features, targets)

    # Create data loader
    loader = torch.utils.data.DataLoader(dataset, batch_size = batch_size, shuffle=shuffle, num_workers=0)
    return loader


import streamlit as st
import numpy as np
import pandas as pd
import plost

st.set_page_config(page_title='Plost', page_icon=':tomato:')

"""
# 🍅 Plost

A deceptively simple plotting library for [Streamlit](https://github.com/streamlit/streamlit).

_“Because you've been writing plots wrong all this time”_

Below you'll find documentation and live examples showing how to use Plost. Of course,
the first step is:

```
pip install streamlit
pip install plost
```

---

## Basics

Plost makes it easy to build common plots using the
[Vega-Lite](https://vega.github.io/vega-lite/)
library but without having to delve into Vega-Lite specs (unless you're doing
something tricky), and without having to melt your DataFrame from long format to wide
format (the bane of most Vega-Lite plots!)

For example, let's say you have a "long-format" table like this:

| time | stock_name | stock_value |
|------|------------|-------------|
| ...  | stock1     | 1           |
| ...  | stock2     | 2           |
| ...  | stock1     | 100         |
| ...  | stock2     | 200         |

Then you can draw a line chart by simply calling `line_chart()` with some
column names:

```python
import plost

plost.line_chart(
  my_dataframe,
  x='time',  # The name of the column to use for the x axis.
  y='stock_value',  # The name of the column to use for the data itself.
  color='stock_name', # The name of the column to use for the line colors.
)
```

Simple enough! But what if you instead have a "wide-format" table like this, which is
super common in reality:

| time | stock1 | stock2 |
|------|--------|--------|
| ...  | 1      | 100    |
| ...  | 2      | 200    |

Normally you'd have to `melt()` the table with Pandas first or create a complex
Vega-Lite layered plot. But with Plost, you can just specify what you're trying
to accomplish and it will melt the data internally for you:

```python
import plost

plost.line_chart(
  my_dataframe,
  x='time',
  y=('stock1', 'stock2'),  # 👈 This is magic!
)
```

Ok, now let's add a mini-map to make panning/zooming even easier:

```python
import plost

plost.line_chart(
  my_dataframe,
  x='time',
  y=('stock1', 'stock2'),
  pan_zoom='minimap',  # 👈 This is magic!
)
```

But we're just scratching the surface. Basically the idea is that Plost allows
you to make beautiful Vega-Lite-driven charts for your most common needs, without
having to learn about the powerful yet complex language behind Vega-Lite.
"""

@st.cache
def get_datasets():
    N = 50
    rand = pd.DataFrame()
    rand['a'] = np.arange(N)
    rand['b'] = np.random.rand(N)
    rand['c'] = np.random.rand(N)

    N = 500
    events = pd.DataFrame()
    events['time_delta_s'] = np.random.randn(N)
    events['servers'] = np.random.choice(['server 1', 'server 2', 'server 3'], N)

    N = 500
    randn = pd.DataFrame(
        np.random.randn(N, 4),
        columns=['a', 'b', 'c', 'd'],
    )

    stocks = pd.DataFrame(dict(
        company=['goog', 'fb', 'ms', 'amazon'],
        q2=[4, 6, 8, 2],
        q3=[2, 5, 2, 6],
    ))

    N = 200
    pageviews = pd.DataFrame()
    pageviews['pagenum'] = [f'page-{i:03d}' for i in range(N)]
    pageviews['pageviews'] = np.random.randint(0, 1000, N)

    return dict(
        rand=rand,
        randn=randn,
        events=events,
        pageviews=pageviews,
        stocks=stocks,
        seattle_weather=pd.read_csv('./data/seattle-weather.csv', parse_dates=['date']),
        sp500=pd.read_csv('./data/sp500.csv', parse_dates=['date']),
    )

datasets = get_datasets()

"""
---

## Datasets used for these examples

Let's say you have some datasets like these:
"""

dataset_name = st.selectbox("Datasets", datasets)
st.write(datasets[dataset_name])

"Where the columns have the following types:"

datasets[dataset_name].dtypes.to_dict(),

"""
Now let's take this data and go _plost_ some _plosts_!

---

## The basics
"""

"### line_chart()"

with st.expander('Documentation'):
    st.write(plost.line_chart)

""

with st.echo():
    plost.line_chart(
        data=datasets['seattle_weather'],
        x='date',
        y='temp_max')

""

with st.echo():
    plost.line_chart(
        data=datasets['seattle_weather'],
        x='date',
        y=('temp_max', 'temp_min'))

"---"

"### area_chart()"

with st.expander('Documentation'):
    st.write(plost.area_chart)
""

with st.echo():
    plost.area_chart(
        data=datasets['rand'],
        x='a',
        y=('b', 'c'))

""

with st.echo():
    plost.area_chart(
        data=datasets['rand'],
        x='a',
        y=('b', 'c'),
        opacity=0.5,
        stack=False)

""

with st.echo():
    plost.area_chart(
        data=datasets['rand'],
        x='a',
        y=('b', 'c'),
        stack='normalize')

"---"

"### bar_chart()"

with st.expander('Documentation'):
    st.write(plost.bar_chart)
""

with st.echo():
    plost.bar_chart(
        data=datasets['stocks'],
        bar='company',
        value='q2')

""

with st.echo():
    plost.bar_chart(
        data=datasets['stocks'],
        bar='company',
        value='q2',
        direction='horizontal')

""

with st.echo():
    plost.bar_chart(
        data=datasets['stocks'],
        bar='company',
        value=['q2', 'q3'],
    )

""

with st.echo():
    plost.bar_chart(
        data=datasets['stocks'],
        bar='company',
        value=['q2', 'q3'],
        stack='normalize')

""

with st.echo():
    plost.bar_chart(
        data=datasets['stocks'],
        bar='company',
        value=['q2', 'q3'],
        group=True)
""

with st.echo():
    plost.bar_chart(
        data=datasets['stocks'],
        bar='company',
        value=['q2', 'q3'],
        group='value',
        color='company',
        legend=None,
    )

"---"

"### pie_chart()"

with st.expander('Documentation'):
    st.write(plost.pie_chart)
""

with st.echo():
    plost.pie_chart(
        data=datasets['stocks'],
        theta='q2',
        color='company')

"---"

"### donut_chart()"

with st.expander('Documentation'):
    st.write(plost.donut_chart)
""

with st.echo():
    plost.donut_chart(
        data=datasets['stocks'],
        theta='q2',
        color='company')

"---"

"### scatter_chart()"

with st.expander('Documentation'):
    st.write(plost.scatter_chart)
""

with st.echo():
    plost.scatter_chart(
        data=datasets['randn'],
        x='a',
        y='b',
        size='c',
        opacity='b',
        height=500)

""

with st.echo():
    plost.scatter_chart(
        data=datasets['randn'],
        x='a',
        y=['b', 'c'],
        size='d',
        height=500)

"---"

"### event_chart()"

with st.expander('Documentation'):
    st.write(plost.event_chart)
""

with st.echo():
    plost.event_chart(
        data=datasets['events'],
        x='time_delta_s',
        y='servers')

""

with st.echo():
    plost.event_chart(
        data=datasets['events'],
        x='time_delta_s',
        y='servers',
        color='servers',
        legend=None)

"""
---

## Histograms
"""

"### hist()"

with st.expander('Documentation'):
    st.write(plost.hist)
""

with st.echo():
    plost.hist(
        data=datasets['randn'],
        x='a',
        aggregate='count')

""

with st.echo():
    plost.hist(
        data=datasets['seattle_weather'],
        x='date',
        y='temp_max',
        aggregate='median')

"---"

"### time_hist()"

with st.expander('Documentation'):
    st.write(plost.time_hist)
""

with st.echo():
    plost.time_hist(
        data=datasets['seattle_weather'],
        date='date',
        x_unit='week',
        y_unit='day',
        color='temp_max',
        aggregate='median',
        legend=None,
    )

"---"

"### xy_hist()"

with st.expander('Documentation'):
    st.write(plost.xy_hist)
""

with st.echo():
    plost.xy_hist(
        data=datasets['randn'],
        x='a',
        y='b',
    )

"---"

with st.echo():
    plost.xy_hist(
        data=datasets['randn'],
        x='a',
        y='b',
        x_bin=dict(maxbins=20),
        y_bin=dict(maxbins=20),
        height=400,
    )

"---"

"""
Woah, double histogram :rainbow:
"""

"### scatter_hist()"

with st.expander('Documentation'):
    st.write(plost.scatter_hist)
""

with st.echo():
    plost.scatter_hist(
        data=datasets['randn'],
        x='a',
        y='b',
        size='c',
        color='c',
        opacity=0.5,
        aggregate='count',
        width=500,
        height=500)

"""
---

# Advanced features

## Vega-Lite encoding dicts

You can use [Vega-Lite encoding dicts](https://vega.github.io/vega-lite/docs/encoding.html) for
the `x`, `y`, `color`, `size`, and `opacity` arguments to do all sorts of fun things. For example,
the chart below is computing the mean of the `y` values, grouped by month.
"""

with st.echo():
    plost.area_chart(
        data=datasets['seattle_weather'],
        x=dict(field='date', timeUnit='month'),
        y=dict(field='temp_max', aggregate='mean'),
        color='weather',
    )

"""
Plost also supports [Altair-style
shorthands](https://altair-viz.github.io/user_guide/encoding.html#encoding-data-types), like
"column_name:T" for temporal.
"""

"""
---

## Annotations

Use `x_annot` and `y_annot` to add vertical or horizontal lines with annotations:
"""

with st.echo():
    plost.area_chart(
        data=datasets['rand'],
        x='a',
        y=('b', 'c'),
        x_annot={
            12: "This is when things became random",
            33: "Actually they were always random. Back to normal now.",
        },
    )

"""
---

## Minimaps

You can add a minimap to many of the charts above my simply passing `pan_zoom='minimap'`.
"""

with st.echo():
    plost.line_chart(
        data=datasets['sp500'],
        x='date',
        y='price',
        width=500,
        pan_zoom='minimap')

"---"

with st.echo():
    plost.area_chart(
        data=datasets['sp500'],
        x='date',
        y='price',
        width=500,
        pan_zoom='minimap')

"---"

with st.echo():
    plost.scatter_chart(
        data=datasets['randn'],
        x='a',
        y='b',
        size='c',
        opacity='b',
        width=500,
        height=500,
        pan_zoom='minimap')

"---"

with st.echo():
    plost.bar_chart(
        data=datasets['pageviews'],
        bar='pagenum',
        value='pageviews',
        width=500,
        pan_zoom='minimap')

"---"

with st.echo():
    plost.bar_chart(
        data=datasets['pageviews'],
        bar='pagenum',
        value='pageviews',
        direction='horizontal',
        width=500,
        height=500,
        pan_zoom='minimap')

""
""
""
""
"🍅"


# -*- coding: utf-8 -*-
r"""
Implements a displayhook for Sage.

This displayhook has two new facilities, by default the displayhook contains a
new facility for displaying lists of matrices in an easier to read format::

    sage: [identity_matrix(i) for i in range(2,5)]
    [
                    [1 0 0 0]
           [1 0 0]  [0 1 0 0]
    [1 0]  [0 1 0]  [0 0 1 0]
    [0 1], [0 0 1], [0 0 0 1]
    ]

This facility uses :meth:`_repr_` (and a simple string) to try do a nice read
format (see :meth:`sage.structure.parent._repr_option` for details).

With this displayhook there exists an other way for displaying object and more
generally, all sage expression as an ASCII art object::

    sage: from sage.misc.interpreter import get_test_shell
    sage: shell = get_test_shell()
    sage: shell.run_cell('%display ascii_art')
    sage: shell.run_cell('integral(x^2/pi^x, x)')
     / 2    2                      \  -x*log(pi)
    -\x *log (pi) + 2*x*log(pi) + 2/*e
    --------------------------------------------
                         3
                      log (pi)
    sage: shell.run_cell("i = var('i')")
    sage: shell.run_cell('sum(i*x^i, i, 0, 10)')
        10      9      8      7      6      5      4      3      2
    10*x   + 9*x  + 8*x  + 7*x  + 6*x  + 5*x  + 4*x  + 3*x  + 2*x  + x
    sage: shell.run_cell('StandardTableaux(4).list()')
    [
    [                                                                  1  4    1  3
    [                 1  3  4    1  2  4    1  2  3    1  3    1  2    2       2
    [   1  2  3  4,   2      ,   3      ,   4      ,   2  4,   3  4,   3   ,   4
    <BLANKLINE>
                1 ]
        1  2    2 ]
        3       3 ]
    ,   4   ,   4 ]
    sage: shell.run_cell('%display simple')

This other facility uses a simple `AsciiArt` object
(see :class:`sage.misc.ascii_art.AsciiArt` and
:meth:`sage.structure.parent._ascii_art_`).

AUTHORS:

- Bill Cauchois (2009): initial version
- Jean-Baptiste Priez <jbp@kerios.fr> (2013): ASCII art
"""

import sys, __builtin__

# This is used to wrap lines when printing "tall" lists.
MAX_COLUMN = 70

def _check_tall_list_and_format(the_list):
    """
    First check whether a list is "tall" -- whether the reprs of the
    elements of the list will span multiple lines and cause the list
    to be printed awkwardly.  If not, this function returns ``None`` and
    does nothing; you should revert back to the normal method for
    printing an object (its repr). If so, return the string in the
    special format. Note that the special format isn't just for
    matrices. Any object with a multiline repr will be formatted.

    INPUT:

    - ``the_list`` - The list (or a tuple).

    TESTS::

        sage: from sage.misc.displayhook import format_obj

    We test _check_tall_list_and_format() indirectly by calling format_obj() on
    a list of matrices::

        sage: print sage.misc.displayhook.format_obj( \
                [matrix([[1, 2, 3, 4], [5, 6, 7, 8]]) for i in xrange(7)])
        [
        [1 2 3 4]  [1 2 3 4]  [1 2 3 4]  [1 2 3 4]  [1 2 3 4]  [1 2 3 4]
        [5 6 7 8], [5 6 7 8], [5 6 7 8], [5 6 7 8], [5 6 7 8], [5 6 7 8],
        <BLANKLINE>
        [1 2 3 4]
        [5 6 7 8]
        ]

    We return ``None`` if we don't have anything special to do::

        sage: format_obj('one-line string')
        sage: format_obj(matrix([[1,2,3]]))
    """
    # For every object to be printed, split its repr on newlines and store the
    # result in this list.
    split_reprs = []
    tall = False
    for elem in the_list:
        split_reprs.append(`elem`.split('\n'))
        if len(split_reprs[-1]) > 1:
            # Meanwhile, check to make sure the list is actually "tall".
            tall = True
    if not tall:
        return None
    # Figure out which type of parenthesis to use, based on the type of the_list.
    if isinstance(the_list, tuple):
        parens = '()'
    elif isinstance(the_list, list):
        parens = '[]'
    else:
        raise TypeError, 'expected list or tuple'

    # running_lines is a list of lines, which are stored as lists of strings
    # to be joined later. For each split repr, we add its lines to the
    # running_lines array. When current_column exceeds MAX_COLUMN, process
    # and output running_lines using _print_tall_list_row.
    running_lines = [[]]
    current_column = 0
    s = [parens[0]]
    for split_repr in split_reprs:
        width = max(len(x) for x in split_repr)
        if current_column + width > MAX_COLUMN and not (width > MAX_COLUMN):
            s.extend(_tall_list_row(running_lines))
            running_lines = [[]]
            current_column = 0
        current_column += width + 2
        # Add the lines from split_repr to the running_lines array. It may
        # be necessary to add or remove lines from either one so that the
        # number of lines matches up.
        for i in xrange(len(running_lines), len(split_repr)):
            running_lines.insert(0, [' ' * len(x) for x in running_lines[-1]])
        line_diff = len(running_lines) - len(split_repr)
        for i, x in enumerate(split_repr):
            running_lines[i + line_diff].append(x.ljust(width))
        for i in xrange(line_diff):
            running_lines[i].append(' ' * width)
    # Output any remaining entries.
    if len(running_lines[0]) > 0:
        s.extend(_tall_list_row(running_lines, True))
    s.append(parens[1])
    return "\n".join(s)

# This helper function for _print_tall_list processes and outputs the
# contents of the running_lines array.
def _tall_list_row(running_lines, last_row=False):
    s=[]
    for i, line in enumerate(running_lines):
        if i + 1 != len(running_lines):
            sep, tail = '  ', ''
        else:
            # The commas go on the bottom line of this row.
            sep, tail = ', ', '' if last_row else ','
        s.append(sep.join(line) + tail)
    # Separate rows with a newline to make them stand out.
    if not last_row:
        s.append("")
    return s

def format_obj(obj):
    """
    This function is used internally by the displayhook.

    We attempt to keep ascii art of list/tuple members intact as we
    print them. See :meth:`sage.structure.parent._repr_option` for
    details.

    OUTPUT:

    Return a string if we want to print it in a special way;
    otherwise, return ``None``.

    EXAMPLES::

        sage: import sage.misc.displayhook

    For most objects, nothing is done (``None`` is returned):

        sage: sage.misc.displayhook.format_obj('Hello, world!')
        sage: sage.misc.displayhook.format_obj((1, 2, 3, 4))

    We demonstrate the special format for lists of matrices::

        sage: sage.misc.displayhook.format_obj( \
                [matrix([[1], [2]]), matrix([[3], [4]])])
        '[\n[1]  [3]\n[2], [4]\n]'

    TESTS:

    In #14466 we override IPython's special printing of ``type`` objects
    and revert it to Python's standard string representation::

        sage: shell=sage.misc.interpreter.get_test_shell()
        sage: shell.displayhook(type)
        <type 'type'>

    """
    if isinstance(obj, type):
        return repr(obj)
    ascii_art = False
    if isinstance(obj, (tuple, list)) and len(obj) > 0:
        for o in obj:
            try:
                ascii_art = ascii_art or o.parent()._repr_option('element_ascii_art')
            except (AttributeError, TypeError):
                pass
    if ascii_art:
        return _check_tall_list_and_format(obj)
    else:
        return None

class DisplayHook(object):
    """
    Display hook for Sage.

    This is not used directly in interactive Sage (where we use the
    IPython system for display hooks).  This class provides a way to
    use the Sage display formatting when not using interactive Sage.
    """
    def __init__(self, oldhook = sys.__displayhook__):
        """
        Set the old display hook (default to repr)

        EXAMPLES::

            sage: from sage.misc.displayhook import DisplayHook
            sage: def f(o): print repr(o)[:5], "..."
            sage: d = DisplayHook(f)
            sage: d(range(10))
            [0, 1 ...
        """
        self.oldhook = oldhook

    def __call__(self, obj):
        """
        Format the object using Sage's formatting, or format it using the old
        display hook if Sage does not want to handle the object.

        EXAMPLES::

            sage: from sage.misc.displayhook import DisplayHook
            sage: d = DisplayHook()
            sage: d((identity_matrix(3), identity_matrix(3)))
            (
            [1 0 0]  [1 0 0]
            [0 1 0]  [0 1 0]
            [0 0 1], [0 0 1]
            )
        """
        s = format_obj(obj)
        if s is not None:
            print s
            __builtin__._ = obj
        else:
            self.oldhook(obj)

from IPython.core.formatters import PlainTextFormatter
from ascii_art import ascii_art
class SagePlainTextFormatter(PlainTextFormatter):
    r"""
    A replacement for the plain text formatter which can use two facilities:

    - correctly print lists of matrices or other objects (see
      :meth:`sage.structure.parent._repr_option`),
    - print ASCII art objects (like expressions) (see
      :meth:`sage.structure.parent._ascii_art_`).

    EXAMPLES::

        sage: from sage.misc.interpreter import get_test_shell
        sage: shell = get_test_shell()
        sage: shell.display_formatter.formatters['text/plain']
        <...displayhook.SagePlainTextFormatter object at 0x...>
        sage: shell.run_cell('a = identity_matrix(ZZ, 2); [a,a]')
        [
        [1 0]  [1 0]
        [0 1], [0 1]
        ]
    """
    def __call__(self, obj):
        r"""
        Computes the format data of ``result``.  If the
        :func:`sage.misc.displayhook.format_obj` writes a string, then
        we override IPython's :class:`DisplayHook` formatting.

        EXAMPLES::

            sage: from sage.misc.interpreter import get_test_shell
            sage: shell = get_test_shell()
            sage: shell.display_formatter.formatters['text/plain']
            <...displayhook.SagePlainTextFormatter object at 0x...>
            sage: shell.displayhook.compute_format_data(2)
            {u'text/plain': '2'}
            sage: a = identity_matrix(ZZ, 2)
            sage: shell.displayhook.compute_format_data([a,a])
            {u'text/plain': '[\n[1 0]  [1 0]\n[0 1], [0 1]\n]'}
            sage: from sage.misc.displayhook import SPTextFormatter
            sage: SPTextFormatter.set_display("ascii_art")
            sage: i = var('i')
            sage: shell.displayhook.compute_format_data(sum(i*x^i, i, 0, 10))
            {u'text/plain':     10      9      8      7      6      5      4      3      2
                            10*x   + 9*x  + 8*x  + 7*x  + 6*x  + 5*x  + 4*x  + 3*x  + 2*x  + x}
        """
        s = self._format_obj(obj)
        if s is None:
            s = super(SagePlainTextFormatter, self).__call__(obj)
        return s

    _format_obj = lambda _, obj: format_obj(obj)

    def set_display(self, mode="ascii_art"):
        r"""
        Method uses to config the formatting method
        (:meth:`simple_format_obj` or :func:`sage.misc.ascii_art.ascii_art`).

        TESTS::

            sage: [identity_matrix(i) for i in range(3,7)]
            [
                                             [1 0 0 0 0 0]
                                [1 0 0 0 0]  [0 1 0 0 0 0]
                     [1 0 0 0]  [0 1 0 0 0]  [0 0 1 0 0 0]
            [1 0 0]  [0 1 0 0]  [0 0 1 0 0]  [0 0 0 1 0 0]
            [0 1 0]  [0 0 1 0]  [0 0 0 1 0]  [0 0 0 0 1 0]
            [0 0 1], [0 0 0 1], [0 0 0 0 1], [0 0 0 0 0 1]
            ]
            sage: from sage.misc.displayhook import SPTextFormatter
            sage: SPTextFormatter.set_display("ascii_art")
            sage: from sage.misc.interpreter import get_test_shell
            sage: shell = get_test_shell()
            sage: shell.run_cell("i = var('i')")
            sage: shell.run_cell('sum(i*x^i, i, 0, 10)')
                10      9      8      7      6      5      4      3      2
            10*x   + 9*x  + 8*x  + 7*x  + 6*x  + 5*x  + 4*x  + 3*x  + 2*x  + x
        """
        self._format_obj = {
            "ascii_art": ascii_art,
            "simple": format_obj
        }[mode]

SPTextFormatter = None


"""
================
Annotating Plots
================

The following examples show how it is possible to annotate plots in Matplotlib.
This includes highlighting specific points of interest and using various
visual tools to call attention to this point. For a more complete and in-depth
description of the annotation and text tools in Matplotlib, see the
:doc:`tutorial on annotation </tutorials/text/annotations>`.
"""

import matplotlib.pyplot as plt
from matplotlib.patches import Ellipse
import numpy as np
from matplotlib.text import OffsetFrom

###############################################################################
# Specifying text points and annotation points
# --------------------------------------------
#
# You must specify an annotation point ``xy=(x, y)`` to annotate this point.
# Additionally, you may specify a text point ``xytext=(x, y)`` for the location
# of the text for this annotation.  Optionally, you can specify the coordinate
# system of *xy* and *xytext* with one of the following strings for *xycoords*
# and *textcoords* (default is 'data')::
#
#  'figure points'   : points from the lower left corner of the figure
#  'figure pixels'   : pixels from the lower left corner of the figure
#  'figure fraction' : (0, 0) is lower left of figure and (1, 1) is upper right
#  'axes points'     : points from lower left corner of axes
#  'axes pixels'     : pixels from lower left corner of axes
#  'axes fraction'   : (0, 0) is lower left of axes and (1, 1) is upper right
#  'offset points'   : Specify an offset (in points) from the xy value
#  'offset pixels'   : Specify an offset (in pixels) from the xy value
#  'data'            : use the axes data coordinate system
#
# Note: for physical coordinate systems (points or pixels) the origin is the
# (bottom, left) of the figure or axes.
#
# Optionally, you can specify arrow properties which draws and arrow
# from the text to the annotated point by giving a dictionary of arrow
# properties
#
# Valid keys are::
#
#   width : the width of the arrow in points
#   frac  : the fraction of the arrow length occupied by the head
#   headwidth : the width of the base of the arrow head in points
#   shrink : move the tip and base some percent away from the
#            annotated point and text
#   any key for matplotlib.patches.polygon  (e.g., facecolor)

# Create our figure and data we'll use for plotting
fig, ax = plt.subplots(figsize=(3, 3))

t = np.arange(0.0, 5.0, 0.01)
s = np.cos(2*np.pi*t)

# Plot a line and add some simple annotations
line, = ax.plot(t, s)
ax.annotate('figure pixels',
            xy=(10, 10), xycoords='figure pixels')
ax.annotate('figure points',
            xy=(80, 80), xycoords='figure points')
ax.annotate('figure fraction',
            xy=(.025, .975), xycoords='figure fraction',
            horizontalalignment='left', verticalalignment='top',
            fontsize=20)

# The following examples show off how these arrows are drawn.

ax.annotate('point offset from data',
            xy=(2, 1), xycoords='data',
            xytext=(-15, 25), textcoords='offset points',
            arrowprops=dict(facecolor='black', shrink=0.05),
            horizontalalignment='right', verticalalignment='bottom')

ax.annotate('axes fraction',
            xy=(3, 1), xycoords='data',
            xytext=(0.8, 0.95), textcoords='axes fraction',
            arrowprops=dict(facecolor='black', shrink=0.05),
            horizontalalignment='right', verticalalignment='top')

# You may also use negative points or pixels to specify from (right, top).
# E.g., (-10, 10) is 10 points to the left of the right side of the axes and 10
# points above the bottom

ax.annotate('pixel offset from axes fraction',
            xy=(1, 0), xycoords='axes fraction',
            xytext=(-20, 20), textcoords='offset pixels',
            horizontalalignment='right',
            verticalalignment='bottom')

ax.set(xlim=(-1, 5), ylim=(-3, 5))

###############################################################################
# Using multiple coordinate systems and axis types
# ------------------------------------------------
#
# You can specify the *xypoint* and the *xytext* in different positions and
# coordinate systems, and optionally turn on a connecting line and mark the
# point with a marker.  Annotations work on polar axes too.
#
# In the example below, the *xy* point is in native coordinates (*xycoords*
# defaults to 'data').  For a polar axes, this is in (theta, radius) space.
# The text in the example is placed in the fractional figure coordinate system.
# Text keyword args like horizontal and vertical alignment are respected.

fig, ax = plt.subplots(subplot_kw=dict(projection='polar'), figsize=(3, 3))
r = np.arange(0, 1, 0.001)
theta = 2*2*np.pi*r
line, = ax.plot(theta, r)

ind = 800
thisr, thistheta = r[ind], theta[ind]
ax.plot([thistheta], [thisr], 'o')
ax.annotate('a polar annotation',
            xy=(thistheta, thisr),  # theta, radius
            xytext=(0.05, 0.05),    # fraction, fraction
            textcoords='figure fraction',
            arrowprops=dict(facecolor='black', shrink=0.05),
            horizontalalignment='left',
            verticalalignment='bottom')

# You can also use polar notation on a cartesian axes.  Here the native
# coordinate system ('data') is cartesian, so you need to specify the
# xycoords and textcoords as 'polar' if you want to use (theta, radius).

el = Ellipse((0, 0), 10, 20, facecolor='r', alpha=0.5)

fig, ax = plt.subplots(subplot_kw=dict(aspect='equal'))
ax.add_artist(el)
el.set_clip_box(ax.bbox)
ax.annotate('the top',
            xy=(np.pi/2., 10.),      # theta, radius
            xytext=(np.pi/3, 20.),   # theta, radius
            xycoords='polar',
            textcoords='polar',
            arrowprops=dict(facecolor='black', shrink=0.05),
            horizontalalignment='left',
            verticalalignment='bottom',
            clip_on=True)  # clip to the axes bounding box

ax.set(xlim=[-20, 20], ylim=[-20, 20])

###############################################################################
# Customizing arrow and bubble styles
# -----------------------------------
#
# The arrow between *xytext* and the annotation point, as well as the bubble
# that covers the annotation text, are highly customizable. Below are a few
# parameter options as well as their resulting output.

fig, ax = plt.subplots(figsize=(8, 5))

t = np.arange(0.0, 5.0, 0.01)
s = np.cos(2*np.pi*t)
line, = ax.plot(t, s, lw=3)

ax.annotate(
    'straight',
    xy=(0, 1), xycoords='data',
    xytext=(-50, 30), textcoords='offset points',
    arrowprops=dict(arrowstyle="->"))
ax.annotate(
    'arc3,\nrad 0.2',
    xy=(0.5, -1), xycoords='data',
    xytext=(-80, -60), textcoords='offset points',
    arrowprops=dict(arrowstyle="->",
                    connectionstyle="arc3,rad=.2"))
ax.annotate(
    'arc,\nangle 50',
    xy=(1., 1), xycoords='data',
    xytext=(-90, 50), textcoords='offset points',
    arrowprops=dict(arrowstyle="->",
                    connectionstyle="arc,angleA=0,armA=50,rad=10"))
ax.annotate(
    'arc,\narms',
    xy=(1.5, -1), xycoords='data',
    xytext=(-80, -60), textcoords='offset points',
    arrowprops=dict(
        arrowstyle="->",
        connectionstyle="arc,angleA=0,armA=40,angleB=-90,armB=30,rad=7"))
ax.annotate(
    'angle,\nangle 90',
    xy=(2., 1), xycoords='data',
    xytext=(-70, 30), textcoords='offset points',
    arrowprops=dict(arrowstyle="->",
                    connectionstyle="angle,angleA=0,angleB=90,rad=10"))
ax.annotate(
    'angle3,\nangle -90',
    xy=(2.5, -1), xycoords='data',
    xytext=(-80, -60), textcoords='offset points',
    arrowprops=dict(arrowstyle="->",
                    connectionstyle="angle3,angleA=0,angleB=-90"))
ax.annotate(
    'angle,\nround',
    xy=(3., 1), xycoords='data',
    xytext=(-60, 30), textcoords='offset points',
    bbox=dict(boxstyle="round", fc="0.8"),
    arrowprops=dict(arrowstyle="->",
                    connectionstyle="angle,angleA=0,angleB=90,rad=10"))
ax.annotate(
    'angle,\nround4',
    xy=(3.5, -1), xycoords='data',
    xytext=(-70, -80), textcoords='offset points',
    size=20,
    bbox=dict(boxstyle="round4,pad=.5", fc="0.8"),
    arrowprops=dict(arrowstyle="->",
                    connectionstyle="angle,angleA=0,angleB=-90,rad=10"))
ax.annotate(
    'angle,\nshrink',
    xy=(4., 1), xycoords='data',
    xytext=(-60, 30), textcoords='offset points',
    bbox=dict(boxstyle="round", fc="0.8"),
    arrowprops=dict(arrowstyle="->",
                    shrinkA=0, shrinkB=10,
                    connectionstyle="angle,angleA=0,angleB=90,rad=10"))
# You can pass an empty string to get only annotation arrows rendered
ax.annotate('', xy=(4., 1.), xycoords='data',
            xytext=(4.5, -1), textcoords='data',
            arrowprops=dict(arrowstyle="<->",
                            connectionstyle="bar",
                            ec="k",
                            shrinkA=5, shrinkB=5))

ax.set(xlim=(-1, 5), ylim=(-4, 3))

# We'll create another figure so that it doesn't get too cluttered
fig, ax = plt.subplots()

el = Ellipse((2, -1), 0.5, 0.5)
ax.add_patch(el)

ax.annotate('$->$',
            xy=(2., -1), xycoords='data',
            xytext=(-150, -140), textcoords='offset points',
            bbox=dict(boxstyle="round", fc="0.8"),
            arrowprops=dict(arrowstyle="->",
                            patchB=el,
                            connectionstyle="angle,angleA=90,angleB=0,rad=10"))
ax.annotate('arrow\nfancy',
            xy=(2., -1), xycoords='data',
            xytext=(-100, 60), textcoords='offset points',
            size=20,
            # bbox=dict(boxstyle="round", fc="0.8"),
            arrowprops=dict(arrowstyle="fancy",
                            fc="0.6", ec="none",
                            patchB=el,
                            connectionstyle="angle3,angleA=0,angleB=-90"))
ax.annotate('arrow\nsimple',
            xy=(2., -1), xycoords='data',
            xytext=(100, 60), textcoords='offset points',
            size=20,
            # bbox=dict(boxstyle="round", fc="0.8"),
            arrowprops=dict(arrowstyle="simple",
                            fc="0.6", ec="none",
                            patchB=el,
                            connectionstyle="arc3,rad=0.3"))
ax.annotate('wedge',
            xy=(2., -1), xycoords='data',
            xytext=(-100, -100), textcoords='offset points',
            size=20,
            # bbox=dict(boxstyle="round", fc="0.8"),
            arrowprops=dict(arrowstyle="wedge,tail_width=0.7",
                            fc="0.6", ec="none",
                            patchB=el,
                            connectionstyle="arc3,rad=-0.3"))
ax.annotate('bubble,\ncontours',
            xy=(2., -1), xycoords='data',
            xytext=(0, -70), textcoords='offset points',
            size=20,
            bbox=dict(boxstyle="round",
                      fc=(1.0, 0.7, 0.7),
                      ec=(1., .5, .5)),
            arrowprops=dict(arrowstyle="wedge,tail_width=1.",
                            fc=(1.0, 0.7, 0.7), ec=(1., .5, .5),
                            patchA=None,
                            patchB=el,
                            relpos=(0.2, 0.8),
                            connectionstyle="arc3,rad=-0.1"))
ax.annotate('bubble',
            xy=(2., -1), xycoords='data',
            xytext=(55, 0), textcoords='offset points',
            size=20, va="center",
            bbox=dict(boxstyle="round", fc=(1.0, 0.7, 0.7), ec="none"),
            arrowprops=dict(arrowstyle="wedge,tail_width=1.",
                            fc=(1.0, 0.7, 0.7), ec="none",
                            patchA=None,
                            patchB=el,
                            relpos=(0.2, 0.5)))

ax.set(xlim=(-1, 5), ylim=(-5, 3))

###############################################################################
# More examples of coordinate systems
# -----------------------------------
#
# Below we'll show a few more examples of coordinate systems and how the
# location of annotations may be specified.

fig, (ax1, ax2) = plt.subplots(1, 2)

bbox_args = dict(boxstyle="round", fc="0.8")
arrow_args = dict(arrowstyle="->")

# Here we'll demonstrate the extents of the coordinate system and how
# we place annotating text.

ax1.annotate('figure fraction : 0, 0', xy=(0, 0), xycoords='figure fraction',
             xytext=(20, 20), textcoords='offset points',
             ha="left", va="bottom",
             bbox=bbox_args,
             arrowprops=arrow_args)

ax1.annotate('figure fraction : 1, 1', xy=(1, 1), xycoords='figure fraction',
             xytext=(-20, -20), textcoords='offset points',
             ha="right", va="top",
             bbox=bbox_args,
             arrowprops=arrow_args)

ax1.annotate('axes fraction : 0, 0', xy=(0, 0), xycoords='axes fraction',
             xytext=(20, 20), textcoords='offset points',
             ha="left", va="bottom",
             bbox=bbox_args,
             arrowprops=arrow_args)

ax1.annotate('axes fraction : 1, 1', xy=(1, 1), xycoords='axes fraction',
             xytext=(-20, -20), textcoords='offset points',
             ha="right", va="top",
             bbox=bbox_args,
             arrowprops=arrow_args)

# It is also possible to generate draggable annotations

an1 = ax1.annotate('Drag me 1', xy=(.5, .7), xycoords='data',
                   #xytext=(.5, .7), textcoords='data',
                   ha="center", va="center",
                   bbox=bbox_args,
                   #arrowprops=arrow_args
                   )

an2 = ax1.annotate('Drag me 2', xy=(.5, .5), xycoords=an1,
                   xytext=(.5, .3), textcoords='axes fraction',
                   ha="center", va="center",
                   bbox=bbox_args,
                   arrowprops=dict(patchB=an1.get_bbox_patch(),
                                   connectionstyle="arc3,rad=0.2",
                                   **arrow_args))
an1.draggable()
an2.draggable()

an3 = ax1.annotate('', xy=(.5, .5), xycoords=an2,
                   xytext=(.5, .5), textcoords=an1,
                   ha="center", va="center",
                   bbox=bbox_args,
                   arrowprops=dict(patchA=an1.get_bbox_patch(),
                                   patchB=an2.get_bbox_patch(),
                                   connectionstyle="arc3,rad=0.2",
                                   **arrow_args))

# Finally we'll show off some more complex annotation and placement

text = ax2.annotate('xy=(0, 1)\nxycoords=("data", "axes fraction")',
                    xy=(0, 1), xycoords=("data", 'axes fraction'),
                    xytext=(0, -20), textcoords='offset points',
                    ha="center", va="top",
                    bbox=bbox_args,
                    arrowprops=arrow_args)

ax2.annotate('xy=(0.5, 0)\nxycoords=artist',
             xy=(0.5, 0.), xycoords=text,
             xytext=(0, -20), textcoords='offset points',
             ha="center", va="top",
             bbox=bbox_args,
             arrowprops=arrow_args)

ax2.annotate('xy=(0.8, 0.5)\nxycoords=ax1.transData',
             xy=(0.8, 0.5), xycoords=ax1.transData,
             xytext=(10, 10),
             textcoords=OffsetFrom(ax2.bbox, (0, 0), "points"),
             ha="left", va="bottom",
             bbox=bbox_args,
             arrowprops=arrow_args)

ax2.set(xlim=[-2, 2], ylim=[-2, 2])
plt.show()