# Copyright 2018-2021 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 jax interface and its jittable variant with
a QNode"""
import pytest
from pennylane import numpy as np

import pennylane as qml
from pennylane import qnode, QNode
from pennylane.tape import QuantumTape
from pennylane.interfaces import InterfaceUnsupportedError

qubit_device_and_diff_method = [
    ["default.qubit", "backprop", "forward", "jax"],
    # Python
    ["default.qubit", "finite-diff", "backward", "jax-python"],
    ["default.qubit", "parameter-shift", "backward", "jax-python"],
    ["default.qubit", "adjoint", "forward", "jax-python"],
    ["default.qubit", "adjoint", "backward", "jax-python"],
    # Jit
    ["default.qubit", "finite-diff", "backward", "jax-jit"],
    ["default.qubit", "parameter-shift", "backward", "jax-jit"],
    ["default.qubit", "adjoint", "forward", "jax-jit"],
    ["default.qubit", "adjoint", "backward", "jax-jit"],
]

pytestmark = pytest.mark.jax

jax = pytest.importorskip("jax")
jnp = jax.numpy

from jax.config import config

config.update("jax_enable_x64", True)

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

    def test_execution_with_interface(self, dev_name, diff_method, mode, interface):
        """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, interface=interface, diff_method=diff_method)
        def circuit(a):
            qml.RY(a, wires=0)
            qml.RX(0.2, wires=0)
            return qml.expval(qml.PauliZ(0))

        a = np.array(0.1, requires_grad=True)
        circuit(a)

        assert circuit.interface == interface

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

        # gradients should work
        grad = jax.grad(circuit)(a)
        assert isinstance(grad, jnp.DeviceArray)
        assert grad.shape == tuple()

    def test_changing_trainability(self, dev_name, diff_method, mode, interface, mocker, tol):
        """Test 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 = jnp.array(0.1)
        b = jnp.array(0.2)

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

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

        grad_fn = jax.grad(circuit, argnums=[0, 1])
        spy = mocker.spy(qml.gradients.param_shift, "transform_fn")
        res = grad_fn(a, b)

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

        expected = [-np.sin(a) + np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]
        assert np.allclose(res, 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
        grad_fn = jax.grad(circuit, argnums=0)
        res = grad_fn(a, b)

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

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

        # The parameter-shift rule has been called only once
        assert len(spy.spy_return[0]) == 2

        # trainability also updates on evaluation
        a = np.array(0.54, requires_grad=False)
        b = np.array(0.8, requires_grad=True)
        circuit(a, b)
        assert circuit.qtape.trainable_params == [1]

    def test_classical_processing(self, dev_name, diff_method, mode, interface, tol):
        """Test classical processing within the quantum tape"""
        a = jnp.array(0.1)
        b = jnp.array(0.2)
        c = jnp.array(0.3)

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

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

        res = jax.grad(circuit, argnums=[0, 2])(a, b, c)

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

        assert len(res) == 2

    def test_matrix_parameter(self, dev_name, diff_method, mode, interface, tol):
        """Test that the jax interface works correctly
        with a matrix parameter"""
        U = jnp.array([[0, 1], [1, 0]])
        a = jnp.array(0.1)

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

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

        res = jax.grad(circuit, argnums=1)(U, a)
        assert np.allclose(res, np.sin(a), atol=tol, rtol=0)

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

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

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

                with QuantumTape() 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 = jnp.array(0.1)
        p = jnp.array([0.1, 0.2, 0.3])

        @qnode(dev, diff_method=diff_method, interface=interface, mode=mode)
        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)
        expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * (
            np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2])
        )
        assert np.allclose(res, expected, atol=tol, rtol=0)

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

vv_qubit_device_and_diff_method = [
    ["default.qubit", "backprop", "forward", "jax"],
    # Python
    ["default.qubit", "finite-diff", "backward", "jax-python"],
    ["default.qubit", "parameter-shift", "backward", "jax-python"],
    ["default.qubit", "adjoint", "forward", "jax-python"],
    ["default.qubit", "adjoint", "backward", "jax-python"],
]

@pytest.mark.parametrize("dev_name,diff_method,mode,interface", vv_qubit_device_and_diff_method)
class TestVectorValuedQNode:
    """Test that using vector-valued QNodes with JAX integrate with the
    PennyLane stack"""

    def test_jacobian(self, dev_name, diff_method, mode, interface, mocker, tol):
        """Test jacobian calculation"""
        if diff_method != "backprop" and mode == "forward":
            pytest.skip(
                "Computing the jacobian of vector-valued tapes is not supported currently in forward mode."
            )

        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 = np.array(0.1, requires_grad=True)
        b = np.array(0.2, requires_grad=True)

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

        @qnode(dev, diff_method=diff_method, interface=interface, mode=mode)
        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 res.shape == (2,)

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

        res = jax.jacobian(circuit, argnums=[0, 1])(a, b)
        expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]).T
        assert np.allclose(res, expected, atol=tol, rtol=0)

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

    def test_jacobian_forward_mode_raises(
        self, dev_name, diff_method, mode, interface, mocker, tol
    ):
        """Test jacobian calculation raises an error in forward mode for
        adjoint differentiation."""
        if diff_method != "adjoint" or mode != "forward":
            pytest.skip("Test only applicable for forward mode adjoint differentiation.")

        a = np.array(0.1, requires_grad=True)
        b = np.array(0.2, requires_grad=True)

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

        @qnode(dev, interface=interface, diff_method=diff_method, mode=mode)
        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 res.shape == (2,)

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

        with pytest.raises(InterfaceUnsupportedError):
            jax.jacobian(circuit, argnums=[0, 1])(a, b)

    def test_jacobian_no_evaluate(self, dev_name, diff_method, mode, interface, mocker, tol):
        """Test jacobian calculation when no prior circuit evaluation has been performed"""
        if mode == "forward":
            pytest.skip(
                "Computing the jacobian of vector-valued tapes is not supported currently in forward mode."
            )

        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 = np.array(0.1, requires_grad=True)
        b = np.array(0.2, requires_grad=True)

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

        @qnode(dev, diff_method=diff_method, interface=interface, mode=mode)
        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))]

        jac_fn = jax.jacobian(circuit, argnums=[0, 1])
        res = jac_fn(a, b)
        expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]).T
        assert np.allclose(res, expected, atol=tol, rtol=0)

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

        # call the Jacobian with new parameters
        a = np.array(0.6, requires_grad=True)
        b = np.array(0.832, requires_grad=True)

        res = jac_fn(a, b)
        expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]).T
        assert np.allclose(res, expected, atol=tol, rtol=0)

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

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

        a = np.array([0.1, 0.2], requires_grad=True)

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

        @qnode(dev, interface=interface, diff_method="finite-diff", 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))

        jax.jacobian(circuit)(a)

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

@pytest.mark.parametrize("interface", ["jax-jit", "jax-python"])
class TestShotsIntegration:
    """Test that the QNode correctly changes shot value, and
    remains differentiable."""

    def test_diff_method_None(self, interface):
        """Test jax device works with diff_method=None."""
        dev = qml.device("default.qubit.jax", wires=1, shots=10)

        @jax.jit
        @qml.qnode(dev, diff_method=None, interface=interface)
        def circuit(x):
            qml.RX(x, wires=0)
            return qml.expval(qml.PauliZ(0))

        assert jnp.allclose(circuit(jnp.array(0.0)), 1)

    def test_changing_shots(self, interface, mocker, tol):
        """Test that changing shots works on execution"""
        dev = qml.device("default.qubit", wires=2, shots=None)
        a, b = jnp.array([0.543, -0.654])

        @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface)
        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 np.allclose(res, -np.cos(a) * np.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 np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0)
        spy.assert_not_called()

    def test_gradient_integration(self, interface, tol, mocker):
        """Test that temporarily setting the shots works
        for gradient computations"""
        dev = qml.device("default.qubit", wires=2, shots=100)
        a, b = jnp.array([0.543, -0.654])

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

        @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface)
        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 = jax.grad(cost_fn, argnums=[0, 1])(a, b, shots=30000)
        assert dev.shots == 100

        expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]
        assert np.allclose(res, expected, atol=0.1, rtol=0)
        assert all(not isinstance(p, jnp.ndarray) for p in spy.call_args[0][0][0].get_parameters())

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

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

        # We're choosing interface="jax" such that backprop can be used in the
        # test later
        @qnode(dev, interface="jax")
        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

@pytest.mark.parametrize("dev_name,diff_method,mode,interface", 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, interface, tol):
        """Tests correct output shape and evaluation for a tape
        with a single prob output"""
        if diff_method == "adjoint":
            pytest.skip("Adjoint does not support probs")

        if interface == "jax-jit":
            pytest.skip(
                "Only Variance and Expectation returns are supported for the jittable JAX interface."
            )

        dev = qml.device(dev_name, wires=2)
        x = jnp.array(0.543)
        y = jnp.array(-0.654)

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

        res = jax.jacobian(circuit, argnums=[0, 1])(x, y)

        expected = np.array(
            [
                [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2],
                [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
            ]
        )
        assert np.allclose(res, expected.T, atol=tol, rtol=0)

    def test_multi_probs_diff(self, dev_name, diff_method, mode, interface, tol):
        """Tests correct output shape and evaluation for a tape
        with multiple prob outputs"""
        if diff_method == "adjoint":
            pytest.skip("Adjoint does not support probs")

        if interface == "jax-jit":
            pytest.skip(
                "Only Variance and Expectation returns are supported for the jittable JAX interface."
            )

        dev = qml.device(dev_name, wires=3)
        x = jnp.array(0.543)
        y = jnp.array(-0.654)

        @qnode(dev, diff_method=diff_method, interface=interface, mode=mode)
        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 / 2) ** 2, np.sin(x / 2) ** 2],
                [(1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(y)) / 2],
            ]
        )
        assert np.allclose(res, expected, atol=tol, rtol=0)
        res = jax.jacobian(circuit, argnums=[0, 1])(x, y)
        expected = np.array(
            [
                [
                    [-np.sin(x) / 2, np.sin(x) / 2],
                    [-np.cos(y) * np.sin(x) / 2, np.sin(x) * np.cos(y) / 2],
                ],
                [
                    [0, 0],
                    [-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2],
                ],
            ]
        )

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

    @pytest.mark.parametrize("ret", [qml.sample(qml.PauliZ(0)), qml.probs(wires=[1, 2])])
    def test_sample_probs_raises_jax_python(self, dev_name, diff_method, mode, ret, interface, tol):
        """Tests qml.sample and qml.probs cannot be used as multiple
        measurements with other measurement types with the JAX Python
        interface."""
        if diff_method == "backprop":
            pytest.skip("Backprop does not apply to this test")

        if ret.return_type is qml.measurements.Sample:
            dev = qml.device(dev_name, wires=3, shots=10)
            if diff_method == "adjoint":
                pytest.skip("Adjoint does not support finite shots")

        if ret.return_type is qml.measurements.Probability:
            if diff_method == "adjoint":
                pytest.skip("Adjoint does not support probs")

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

        x = jnp.array(0.543)
        y = jnp.array(-0.654)

        @qnode(dev, diff_method=diff_method, interface="jax-python", mode=mode)
        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.apply(ret)

        with pytest.raises(InterfaceUnsupportedError, match="sample and probability measurements"):
            circuit(x, y)

    def test_probs_diff_len_wires_raises(self, dev_name, diff_method, mode, interface, tol):
        """Tests multiple probs raise an error if the number of wires do not
        match with the JAX Python interface."""
        if diff_method == "backprop":
            pytest.skip("Backprop does not apply to this test")

        if diff_method == "adjoint":
            pytest.skip("Adjoint does not support probs")

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

        x = jnp.array(0.543)
        y = jnp.array(-0.654)

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

        with pytest.raises(
            InterfaceUnsupportedError,
            match="multiple probability measurements need to have the same number of wires specified",
        ):
            circuit(x, y)

    @pytest.mark.xfail(reason="Line 230 in QubitDevice: results = self._asarray(results) fails")
    def test_ragged_differentiation(self, dev_name, diff_method, mode, interface, tol):
        """Tests correct output shape and evaluation for a tape
        with prob and expval outputs"""
        if diff_method == "adjoint":
            pytest.skip("Adjoint does not support probs")

        dev = qml.device(dev_name, wires=2)
        x = jnp.array(0.543)
        y = jnp.array(-0.654)

        @qnode(dev, diff_method=diff_method, interface=interface, mode=mode)
        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), (1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(y)) / 2]
        )
        assert np.allclose(res, expected, atol=tol, rtol=0)

        res = jax.jacobian(circuit, argnums=[0, 1])(x, y)
        expected = np.array(
            [
                [-np.sin(x), 0],
                [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2],
                [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
            ]
        )
        assert np.allclose(res, expected, atol=tol, rtol=0)

    @pytest.mark.xfail(reason="Line 230 in QubitDevice: results = self._asarray(results) fails")
    def test_ragged_differentiation_variance(self, dev_name, diff_method, mode, interface, tol):
        """Tests correct output shape and evaluation for a tape
        with prob and variance outputs"""
        if diff_method == "adjoint":
            pytest.skip("Adjoint does not support probs")

        dev = qml.device(dev_name, wires=2)
        x = jnp.array(0.543)
        y = jnp.array(-0.654)

        @qnode(dev, diff_method=diff_method, interface=interface, 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.PauliZ(0)), qml.probs(wires=[1])]

        res = circuit(x, y)

        expected = np.array(
            [np.sin(x) ** 2, (1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(y)) / 2]
        )
        assert np.allclose(res, expected, atol=tol, rtol=0)

        res = jax.jacobian(circuit, argnums=[0, 1])(x, y)
        expected = np.array(
            [
                [2 * np.cos(x) * np.sin(x), 0],
                [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2],
                [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2],
            ]
        )
        assert np.allclose(res, expected, atol=tol, rtol=0)

    def test_sampling(self, dev_name, diff_method, mode, interface):
        """Test sampling works as expected"""
        if mode == "forward":
            pytest.skip("Sampling not possible with forward mode differentiation.")

        if diff_method == "adjoint":
            pytest.skip("Adjoint warns with finite shots")

        if interface == "jax-jit":
            pytest.skip(
                "Only Variance and Expectation returns are supported for the jittable JAX interface."
            )

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

        @qnode(dev, diff_method=diff_method, interface=interface, mode=mode)
        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, jnp.DeviceArray)

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

        class Template(qml.templates.StronglyEntanglingLayers):
            def expand(self):
                with qml.tape.QuantumTape() as tape:
                    qml.templates.StronglyEntanglingLayers(*self.parameters, self.wires)
                return tape

        @qnode(dev, interface=interface, diff_method=diff_method)
        def circuit1(weights):
            Template(weights, wires=[0, 1])
            return qml.expval(qml.PauliZ(0))

        @qnode(dev, interface=interface, diff_method=diff_method)
        def circuit2(data, weights):
            qml.templates.AngleEmbedding(jnp.stack([data, 0.7]), wires=[0, 1])
            Template(weights, wires=[0, 1])
            return qml.expval(qml.PauliX(0))

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

        w1 = qml.templates.StronglyEntanglingLayers.shape(n_wires=2, n_layers=3)
        w2 = qml.templates.StronglyEntanglingLayers.shape(n_wires=2, n_layers=4)

        weights = [
            jnp.array(np.random.random(w1)),
            jnp.array(np.random.random(w2)),
        ]

        grad_fn = jax.grad(cost)
        res = grad_fn(weights)

        assert len(res) == 2

    def test_second_derivative(self, dev_name, diff_method, mode, interface, tol):
        """Test second derivative calculation of a scalar valued QNode"""
        if diff_method not in {"backprop"}:
            pytest.skip("Test only supports backprop")

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

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

        x = jnp.array([1.0, 2.0])
        res = circuit(x)
        g = jax.grad(circuit)(x)
        g2 = jax.grad(lambda x: jnp.sum(jax.grad(circuit)(x)))(x)

        a, b = x

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

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

        expected_g2 = [
            -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(g2, expected_g2, atol=tol, rtol=0)

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

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

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

        x = jnp.array([1.0, 2.0])
        res = circuit(x)

        a, b = x

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

        grad_fn = jax.grad(circuit)
        g = grad_fn(x)

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

        hess = jax.jacobian(grad_fn)(x)

        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, expected_hess, atol=tol, rtol=0)

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

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

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

        x = jnp.array([1.0, 2.0])
        res = circuit(x)

        a, b = x

        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, expected_res, atol=tol, rtol=0)

        jac_fn = jax.jacobian(circuit)
        g = jac_fn(x)

        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, expected_g, atol=tol, rtol=0)

        hess = jax.jacobian(jac_fn)(x)

        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, expected_hess, atol=tol, rtol=0)

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

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

        # Test only applies to backprop -> interface="jax"
        @qnode(dev, diff_method=diff_method, interface="jax", mode=mode, max_diff=2)
        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))]

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

        x = jnp.array(
            [0.76, -0.87],
        )
        res = cost_fn(x)

        a, b = x

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

        grad_fn = jax.grad(cost_fn)
        g = grad_fn(x)

        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, expected_g, atol=tol, rtol=0)
        hess = jax.jacobian(grad_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, expected_hess, atol=tol, rtol=0)

    def test_hessian_vector_valued_separate_args(
        self, dev_name, diff_method, mode, interface, mocker, tol
    ):
        """Test hessian calculation of a vector valued QNode that has separate input arguments"""
        if diff_method not in {"backprop"}:
            pytest.skip("Test only supports backprop")

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

        @qnode(dev, diff_method=diff_method, interface=interface, mode=mode, max_diff=2)
        def circuit(a, b):
            qml.RY(a, wires=0)
            qml.RX(b, wires=0)
            return qml.probs(wires=0)

        a = jnp.array(1.0)
        b = jnp.array(2.0)
        res = circuit(a, b)

        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, expected_res, atol=tol, rtol=0)

        jac_fn = jax.jacobian(circuit, argnums=[0, 1])
        g = jac_fn(a, b)

        expected_g = np.array(
            [
                [-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, expected_g.T, atol=tol, rtol=0)

        spy = mocker.spy(qml.gradients.param_shift, "transform_fn")
        hess = jax.jacobian(jac_fn, argnums=[0, 1])(a, b)

        if diff_method == "backprop":
            spy.assert_not_called()
        elif diff_method == "parameter-shift":
            spy.assert_called()

        expected_hess = np.array(
            [
                [
                    [-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.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)],
                ],
            ]
        )
        assert np.allclose(hess, expected_hess, atol=tol, rtol=0)

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

        if interface == "jax-jit":
            pytest.skip(
                "Only Variance and Expectation returns are supported for the jittable JAX interface."
            )

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

        x = jnp.array(0.543)
        y = jnp.array(-0.654)

        @qnode(dev, diff_method=diff_method, interface=interface, 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 np.dtype("complex128")
            probs = jnp.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 = jax.grad(cost_fn, argnums=[0, 1])(x, y)
        expected = np.array([-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2])
        assert np.allclose(res, expected, atol=tol, rtol=0)

    def test_projector(self, dev_name, diff_method, mode, interface, 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 = jnp.array([1])
        x, y = 0.765, -0.654

        @qnode(dev, diff_method=diff_method, interface=interface, 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(x, y)
        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, expected, atol=tol, rtol=0)

        res = jax.grad(circuit, argnums=[0, 1])(x, y)
        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(res, expected, atol=tol, rtol=0)

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

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

        r = 0.543
        phi = -0.654

        @qnode(dev, interface=interface, 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 = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2
        assert np.allclose(res, expected, atol=tol, rtol=0)

        # circuit jacobians
        res = jax.grad(circuit, argnums=[0, 1])(r, phi)
        expected = np.array(
            [
                2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2,
                2 * np.sinh(2 * r) * np.sin(2 * phi),
            ]
        )
        assert np.allclose(res, expected, atol=tol, rtol=0)

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

        n = 0.12
        a = 0.765

        @qnode(dev, interface=interface, 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 + np.abs(a) ** 2 * (1 + 2 * n)
        assert np.allclose(res, expected, atol=tol, rtol=0)

        # circuit jacobians
        res = jax.grad(circuit, argnums=[0, 1])(n, a)
        expected = np.array([2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)])
        assert np.allclose(res, expected, atol=tol, rtol=0)

@pytest.mark.parametrize("interface", ["jax-jit", "jax-python"])
def test_adjoint_reuse_device_state(mocker, interface):
    """Tests that the jax interface reuses the device state for adjoint differentiation"""
    dev = qml.device("default.qubit", wires=1)

    @qnode(dev, interface=interface, diff_method="adjoint")
    def circ(x):
        qml.RX(x, wires=0)
        return qml.expval(qml.PauliZ(0))

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

    grad = jax.grad(circ)(1.0)
    assert circ.device.num_executions == 1

    spy.assert_called_with(mocker.ANY, use_device_state=True)

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

    @pytest.mark.parametrize("max_diff", [1, 2])
    def test_gradient_expansion_trainable_only(
        self, dev_name, diff_method, mode, max_diff, interface, 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")

        if max_diff > 1:
            pytest.skip("JAX only supports first derivatives")

        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=interface)
        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 = jnp.array(0.5)
        y = jnp.array(0.7)
        circuit(x, y)

        spy = mocker.spy(circuit.gradient_fn, "transform_fn")
        res = jax.grad(circuit, argnums=[0])(x, y)

        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, interface, mocker
    ):
        """Test that the Hamiltonian is not expanded if there
        are non-commuting groups and the number of shots is None
        and the first and second order gradients are correctly evaluated"""
        if diff_method == "adjoint":
            pytest.skip("The adjoint method does not yet support Hamiltonians")

        if max_diff > 1:
            pytest.skip("JAX only supports first derivatives")

        dev = qml.device(dev_name, wires=3, shots=None)
        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, interface=interface, diff_method=diff_method, mode=mode, max_diff=max_diff)
        def circuit(data, weights, coeffs):
            weights = weights.reshape(1, -1)
            qml.templates.AngleEmbedding(data, wires=[0, 1])
            qml.templates.BasicEntanglerLayers(weights, wires=[0, 1])
            return qml.expval(qml.Hamiltonian(coeffs, obs))

        d = jnp.array([0.1, 0.2])
        w = jnp.array([0.654, -0.734])
        c = jnp.array([-0.6543, 0.24, 0.54])

        # test output
        res = circuit(d, w, c)
        expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1])
        assert np.allclose(res, expected)
        spy.assert_not_called()

        # test gradients
        grad = jax.grad(circuit, argnums=[1, 2])(d, w, c)
        expected_w = [
            -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]),
            -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]),
        ]
        expected_c = [0, -np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]), np.cos(d[1] + w[1])]
        assert np.allclose(grad[0], expected_w)
        assert np.allclose(grad[1], expected_c)

        # test second-order derivatives
        if diff_method in ("parameter-shift", "backprop") and max_diff == 2:

            grad2_c = jax.jacobian(jax.grad(circuit, argnum=2), argnum=2)(d, w, c)
            assert np.allclose(grad2_c, 0)

            grad2_w_c = jax.jacobian(jax.grad(circuit, argnum=1), argnum=2)(d, w, c)
            expected = [0, -np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), 0], [
                0,
                -np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]),
                -np.sin(d[1] + w[1]),
            ]
            assert np.allclose(grad2_w_c, expected)

    # @pytest.mark.xfail(reason="Will fail since expval(H) expands to a vector valued return for finite-shots")
    @pytest.mark.parametrize("max_diff", [1, 2])
    def test_hamiltonian_expansion_finite_shots(
        self, dev_name, diff_method, mode, interface, 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")

        if interface == "jax-jit":
            pytest.skip(
                "Only Variance and Expectation returns are supported for the jittable JAX interface."
            )

        if max_diff > 1:
            pytest.skip("JAX only supports first derivatives")

        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, interface=interface, diff_method=diff_method, mode=mode, max_diff=max_diff)
        def circuit(data, weights, coeffs):
            weights = weights.reshape(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 = jnp.array([0.1, 0.2])
        w = jnp.array([0.654, -0.734])
        c = jnp.array([-0.6543, 0.24, 0.54])

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

        # test gradients
        grad = jax.grad(circuit, argnums=[1, 2])(d, w, c)
        expected_w = [
            -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]),
            -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]),
        ]
        expected_c = [0, -np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]), np.cos(d[1] + w[1])]
        assert np.allclose(grad[0], expected_w, atol=0.1)
        assert np.allclose(grad[1], expected_c, atol=0.1)

    #     # test second-order derivatives
    #     if diff_method == "parameter-shift" and max_diff == 2:

    #         grad2_c = jax.jacobian(jax.grad(circuit, argnum=2), argnum=2)(d, w, c)
    #         assert np.allclose(grad2_c, 0, atol=0.1)

    #         grad2_w_c = jax.jacobian(jax.grad(circuit, argnum=1), argnum=2)(d, w, c)
    #         expected = [0, -np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), 0], [
    #             0,
    #             -np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]),
    #             -np.sin(d[1] + w[1]),
    #         ]
    #         assert np.allclose(grad2_w_c, expected, atol=0.1)

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

@pytest.mark.parametrize("dev_name,diff_method,mode", jit_qubit_device_and_diff_method)
class TestJIT:
    """Test JAX JIT integration with the QNode and automatic resolution of the
    correct JAX interface variant."""

    def test_gradient(self, dev_name, diff_method, mode, tol):
        """Test derivative calculation of a scalar valued QNode"""
        dev = qml.device(dev_name, wires=1)

        if diff_method == "adjoint":
            pytest.xfail(reason="The adjoint method is not using host-callback currently")

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

        x = jnp.array([1.0, 2.0])
        res = circuit(x)
        g = jax.grad(circuit)(x)

        a, b = x

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

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

    @pytest.mark.filterwarnings(
        "ignore:Requested adjoint differentiation to be computed with finite shots."
    )
    @pytest.mark.parametrize("shots", [10, 1000])
    def test_hermitian(self, dev_name, diff_method, mode, shots):
        """Test that the jax device works with qml.Hermitian and jitting even
        when shots>0.

        Note: before a fix, the cases of shots=10 and shots=1000 were failing due
        to different reasons, hence the parametrization in the test.
        """
        dev = qml.device(dev_name, wires=2, shots=shots)

        if diff_method == "backprop":
            pytest.skip("Backpropagation is unsupported if shots > 0.")

        if diff_method == "adjoint" and mode == "forward":
            pytest.skip("Computing the gradient for Hermitian is not supported with adjoint.")

        projector = np.array(qml.matrix(qml.PauliZ(0) @ qml.PauliZ(1)))

        @jax.jit
        @qml.qnode(dev, interface="jax", diff_method=diff_method, mode=mode)
        def circ(projector):
            return qml.expval(qml.Hermitian(projector, wires=range(2)))

        assert jnp.allclose(circ(projector), 1)

    @pytest.mark.filterwarnings(
        "ignore:Requested adjoint differentiation to be computed with finite shots."
    )
    @pytest.mark.parametrize("shots", [10, 1000])
    def test_probs_obs_none(self, dev_name, diff_method, mode, shots):
        """Test that the jax device works with qml.probs, a MeasurementProcess
        that has obs=None even when shots>0."""
        dev = qml.device(dev_name, wires=2, shots=shots)

        if diff_method == "backprop":
            pytest.skip("Backpropagation is unsupported if shots > 0.")

        @qml.qnode(dev, interface="jax", diff_method="parameter-shift")
        def circuit():
            return qml.probs(wires=0)

        assert jnp.allclose(circuit(), jnp.array([1.0, 0.0]))

    @pytest.mark.xfail(
        reason="Non-trainable parameters are not being correctly unwrapped by the interface"
    )
    def test_gradient_subset(self, dev_name, diff_method, mode, tol):
        """Test derivative calculation of a scalar valued QNode with respect
        to a subset of arguments"""
        a = jnp.array(0.1)
        b = jnp.array(0.2)

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

        @jax.jit
        @qnode(dev, diff_method=diff_method, interface="jax", mode=mode)
        def circuit(a, b):
            qml.RY(a, wires=0)
            qml.RX(b, wires=0)
            qml.RZ(c, wires=0)
            return qml.expval(qml.PauliZ(0))

        res = jax.grad(circuit, argnums=[0, 1])(a, b, 0.0)

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

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


# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py:light
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.7.1
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Video Actor Synchroncy and Causality (VASC)
# ## RAEng: Measuring Responsive Caregiving Project
# ### Caspar Addyman, 2020
# ### https://github.com/infantlab/VASC
#
# # Step 2: Reorganise the OpenPose JSON wire frame data
#
# This script uses output from [OpenPose](https://github.com/CMU-Perceptual-Computing-Lab/openpose) human figure recognition neural network to create labeled wireframes for each figure in each frame of a video.
#
#
# The `write_json flag` saves the people pose data using a custom JSON writer. Each JSON file has a set of coordinates and confidence scores for each person identified in the frame. For a given person there is:
#
# > An array pose_keypoints_2d containing the body part locations and detection confidence formatted as x1,y1,c1,x2,y2,c2,.... The coordinates x and y can be normalized to the range [0,1], [-1,1], [0, source size], [0, output size], etc., depending on the flag keypoint_scale (see flag for more information), while c is the confidence score in the range [0,1].
#
# <img src="keypoints_pose_25.png" alt="BODY-25 mapping" width="240"/>

# ## 2.1 - import modules and initialise variables

# +
import os                #operating system functions
import math              #simple math
import glob              #file listing
import json              #importing and exporting json files
import cv2               #computervision toolkit
import numpy as np       #tools for numerical data
import pandas as pd
import logging
import ipywidgets as widgets  #let's us add buttons and sliders to this page.
from ipycanvas import Canvas

import matplotlib.pyplot as plt
# %matplotlib inline

import vasc #a module of our own functions (found in vasc.py in this folder)

#turn on debugging
logger = logging.getLogger()
logger.setLevel(logging.INFO)
# %pdb on
# -

# #### 2.1.1 - anonymise the videos?
#
# Setting the `anon` flag to
#
# * `True` - we will not display just the wireframes on black backround without the underlying images from the video.
# * `False` - we will *attempt to* draw video images - If videos are not available we fall back to anonymous mode

anon = False

# ## 2.2 - Where are the data?
#
# This routine only needs to know where to find the processed data  and what are the base names. The summary information is listed in the `videos.json` file we created. The raw numerical data is in `allframedata.npz`.

# +
# where's the project folder? (without trailing slash)
projectpath = "C:\\Users\\cas\\OneDrive - Goldsmiths College\\Projects\\Measuring Responsive Caregiving\\VASCTutorial"
#where are the videos
videos_in = "C:\\Users\\cas\\OneDrive - Goldsmiths College\\Projects\\Measuring Responsive Caregiving\\VASCTutorial\\demovideos"

# locations of videos and output
videos_in = projectpath
videos_out   = projectpath + "\\out"
videos_out_openpose   = videos_out + "\\openpose"
videos_out_timeseries = videos_out + "\\timeseries"
videos_out_analyses   = videos_out + "\\analyses"

print(videos_out_openpose)
print(videos_out_timeseries)
print(videos_out_analyses)
# -

#retrieve the list of base names of processed videos.
try:
    with open(videos_out + '\\videos.json') as json_file:
        videos = json.load(json_file)
        print("Existing videos.json found..")
except:
    videos = {}
    print("videos.json not found in ", videos_out)

#optional - check the json
for vid in videos:
    print(vid)
    for cam in videos[vid]:
        print(videos[vid][cam])

# ### Step 2.2.1 Load or reload the raw/cleaned data.
#
# At this step we can either load the unprocessed data from Step 1 - `allframedata.npz` or we can load a cleaned or partially cleaned set of data that has already by processed in an earlier session - `cleandata.npz`.

# EITHER
# can reload the original values without recomputing
reloaded = np.load(videos_out_timeseries + '\\allframedata.npz')
# OR
# load clean or partially cleaned data from a previous session
#reloaded = np.load(videos_out_timeseries + '\\cleandata.npz')
keypoints_original = reloaded["keypoints_array"] #the unprocessed data

keypoints_array = keypoints_original  #an array where we clean the data.

#check the shape
keypoints_array.shape

# ## Step 2.3 Clean the data
#
# We now have an numpy array called `keypoints_array` containing all the openpose numbers for all videos. Now we need to do some cleaning of the data. We provide set of tools to do this. There are several tasks we need to do.
#
# 1. Pick camera with best view of both participants - swap this to camera 1 (if multiple cameras).
# 2. You might delete sets for whom all data is too poor quality. But they can also be excluded in **Step 3** by a flag in the data spreadsheet.
# 3. Tag the adult & infant in first frame of interest. So both individuals should be in first frame.
# 4. Try to automatically tag then in subsequent frames.
# 5. Manually fix anything the automatic process gets wrong.
# 6. Exclude other detected people (3rd parties & false positives)
#
# We do all of this with the control panel below.
#

# ### Step 2.3.1: Which is best camera angle?
#
# If we have just one camera then use that. If there are multiple angles, pick the best one and swap it to be "camera1".
#

# ### Step 2.3.2: Where does the interesting data start and end?
#
# For many videos, the period of interest might start (and end) some time into the video. For now we are using the whole video .
# TODO  - We wiil give the user the opportunity to set these.
#

#let's loop through the processed list and set and startframe and endframe for each video
# for the moment we'll just use the full video.
# TODO - we will let the use specify this per video
for vid in videos:
    for cam in videos[vid]:
        videos[vid][cam]["start"] = 0
        videos[vid][cam]["end"] = videos[vid][cam]["frames"]

# ### Step 2.3.2: Tag the actors of interest at start
#
# We want to know which person is the adult and which is the infant in the first frame. We want the child to be Person 0 and Adult to be Person 1. The buttons below provide the choice to swap data series so that this is correct.
#
# For example, if child data starts in series 3, we pick 3 in the drop down list next to button `Swap to child (0)` and then press the button. This will swap these two series.
#
# This function operates beyond the current frame so it's possible to use it multiple times if data jumps around. However, there is a short cut for part of this process..

# ### Step 2.3.3: Fix by location - Track actors frame by frame
#
# At present OpenPose doesn't track individuals from one frame to the next (I believe they are working on this). It just labels each person in each frame. This means that Person 1 in frame 1 might become Person 2 by frame 100. Here we provide some tools that automatically trying to guess who is who. This is tricky so we also ask for human input.
#
# Once the child and adult data start in series 0 and 1 respectively, press the `Fix by location` button.
#
# If this leaves a few errors we can move slider to affected frame and use the swap series function to manually correct.

# ### Step 2.3.4: Fix by size - identify data by size of wireframes.
#
# In many of our cases of interest we have an adult and a young child interacting. Therefore, it is handy to try autolabelling based on a sorting of the size of their wireframes.
#
# Pressing the `Fix by size` button, makes person 0 the smallest person in the frame, person 1 the next smallest and so on.

# ### Step 2.3.4: Exclude other people
#
# Finally we can delete any people in background or false positives (ghosts) detected by OpenPose. We simply set these to zero.

# ## CONTROL PANEL
#
# Run this BIG block of code to provide controls to edit and reorganise the data.
# If anything goes wrong you can revert to the original data.
#
# This needs `ipywidgets` and `ipycanvas` to be installed. (See [Step 0](Step0.GettingStarted.ipynb))

# +
###############################################################
# All the code in this block draws our data editing control panel
###############################################################

###############################################################
## a canvas object to show current frame of the video.
canvas = Canvas(width=800, height=600)

###############################################################
## dropbown lists to select the video and the camera (if multiple angles)
## next to the video dropdown we have Delete button to remove that vid
## next to the camera dropdown we have a swap button to choose a primary camera.
vidlist = [] #used to fill dropdown options
camlist = [] #used to fill dropdown options
for vid in videos:
    vidlist.append(vid)

for cam in videos[vid]:
    camlist.append(cam)

pickvid = widgets.Dropdown(
    options= vidlist,
    value= vidlist[0],
    description='Select subject:'
)
button_exclude =  widgets.Button(description='DELETE THIS ONE!')

pickcam = widgets.Dropdown(
    options= camlist,
    value= camlist[0],
    description='Select camera:'
)
button_swapcam = widgets.Button(description="Swap this to camera1")
cambox = widgets.HBox([pickcam, button_swapcam])

###############################################################
## Who is who?
## In processed video we want child in index 0 and adult in index 1.
## so need ability to swap series and delete unwanted data.
## pressing button_swapchild swaps selected set of data to be index 0 - default for child.

button_swapchild = widgets.Button(description="Swap to child (0)")
child = widgets.Dropdown(
    options = list(range(10)),
    value= 0,
    description='Set: '
)
babybox = widgets.HBox([button_swapchild, child])

adult = widgets.Dropdown(
    options = list(range(10)),
    value= 1,
    description='Set: '
)
button_swapadult = widgets.Button(description="Swap to adult (1)")
adultbox = widgets.HBox([button_swapadult,adult])

button_remove = widgets.Button(description="Remove these data")
remove = widgets.Dropdown(
    options = list(range(10)),
    value= 2,
    description='Set: '
)
removebox = widgets.HBox([button_remove,remove])

###############################################################
## What frame is displayed in the canvas?
## all swap and delete operation work on data AFTER this frame.
## include a few buttons to adjust the frame forward or backwards slightly
slider = widgets.IntSlider(
    value=0,
    min=0,
    max=161,
    step=1,
    description='Frame:',
    continuous_update=False,
    orientation='horizontal',
    readout=True,
    readout_format='d',
    layout=widgets.Layout(width='800px')
)

#buttons to adjust the slider in small increments
minus1pct = widgets.Button(description="-1%")
minus10 = widgets.Button(description="-10")
minus1 = widgets.Button(description="-1")
plus1 = widgets.Button(description="+1")
plus10 = widgets.Button(description="+10")
plus1pct = widgets.Button(description="+1%")

def minus1pct_clicked(output):
    slider.value = max(0,slider.value - 0.01 * slider.max)
def minus10_clicked(output):
    slider.value = max(0,slider.value - 10)
def minus1_clicked(output):
    slider.value = max(0,slider.value - 1)
def plus1_clicked(output):
    slider.value = min(slider.max,slider.value + 1)
def plus10_clicked(output):
    slider.value = min(slider.max,slider.value + 10)
def plus1pct_clicked(output):
    slider.value = min(slider.max,slider.value + 0.01 * slider.max)

minus1pct.on_click(minus1pct_clicked)
minus10.on_click(minus10_clicked)
minus1.on_click(minus1_clicked)
plus1.on_click(plus1_clicked)
plus10.on_click(plus10_clicked)
plus1pct.on_click(plus1pct_clicked)

adjustbox  = widgets.HBox([minus1pct,minus10,minus1,plus1,plus10,plus1pct])

###############################################################
## Action buttons
## To redraw everything it's current state, to attempt autofixing or to undo some or all our changes

button_update = widgets.Button(description="Redraw")
button_fixlocations = widgets.Button(description="Fix by location",tooltip="match each person to nearest person in next frame")
button_fixsizes = widgets.Button(description="Fix by size",tooltip="label people sequentially by size of their wireframe")
button_reset_one = widgets.Button(description="Reset this video")
button_reset_all = widgets.Button(description="Reset all")
buttonbox = widgets.HBox([button_update,button_fixlocations,button_fixsizes,button_exclude,button_reset_one,button_reset_all])
output = widgets.Output()

###############################################################
## Widget 'Event' codes
## watches each widget waiting for something to change and then executes these bits of code.

def pickvid_change(change):
    if change['name'] == 'value' and (change['new'] != change['old']):
        updateAll(True)

def pickcam_change(change):
    if change['name'] == 'value' and (change['new'] != change['old']):
        updateAll(True)

def slider_change(slider):
    updateAll(False)

def on_button_clicked(output):
    logging.info('button_update_all clicked')
    updateAll(True)

def on_reset_all(output):
    global keypoints_array
    logging.info('button_reset_all clicked')
    keypoints_array = np.copy(keypoints_original)
    updateAll(True)

def on_fixlocations(output):
    global keypoints_array
    logging.info('on_fixlocations')
    v = videos[pickvid.value][pickcam.value]["v"]
    c = videos[pickvid.value][pickcam.value]["c"]
    end  = videos[pickvid.value][pickcam.value]["end"]
    window = 10
    vasc.fixpeopleSeries(keypoints_array,v,c,[0,1],slider.value, end, window)
    updateAll(True)

def on_fixsizes(output):
    global keypoints_array
    logging.info('on_fixsizes')
    v = videos[pickvid.value][pickcam.value]["v"]
    c = videos[pickvid.value][pickcam.value]["c"]
    N = videos[vid][cam]["maxpeople"]
    end  = videos[pickvid.value][pickcam.value]["end"]
    vasc.sortpeoplebySize(keypoints_array,v,c,N,slider.value, end)
    updateAll(True)

def on_deleteparticipant(output):
    global keypoints_array
    global videos
    logging.info('on_deleteparticipant')
    for cam in videos[pickvid.value]: #loop through delete all cameras for this video.
        v = videos[pickvid.value][cam]["v"]
        c = videos[pickvid.value][cam]["c"]
        end  = videos[pickvid.value][cam]["end"]
        vasc.deleteSeries(keypoints_array,v,c,remove.value,0, end)
    #now remove this from videos object
    if pickvid.value in videos:
        logging.info(pickvid.value)
        del videos[pickvid.value]
    #repopulate the dropdown
    for vid in videos:
        vidlist.append(vid)
    pickvid.options = vidlist
    updateAll(True)

def on_deleteseries(output):
    global keypoints_array
    global videos
    logging.info('on_fixseries')
    v = videos[pickvid.value][pickcam.value]["v"]
    c = videos[pickvid.value][pickcam.value]["c"]
    end  = videos[pickvid.value][pickcam.value]["end"]
    vasc.deleteSeries(keypoints_array,v,c,remove.value,slider.value, end)
    updateAll(True)

def on_swapcam(output):
    global keypoints_array
    global videos
    logging.info('on_swapcam')
    print(videos[pickvid.value][pickcam.value])
    videos, keypoints_array = vasc.swapCameras(videos, keypoints_array,pickvid.value,pickcam.value,"camera1")
    updateAll(True)

def on_swapchild(output):
    global keypoints_array
    global videos
    logging.info('on_swapchild')
    v = videos[pickvid.value][pickcam.value]["v"]
    c = videos[pickvid.value][pickcam.value]["c"]
    end  = videos[pickvid.value][pickcam.value]["end"]
    vasc.swapSeries(keypoints_array,v,c,0,child.value,slider.value,end)
    updateAll(True)

def on_swapadult(output):
    global keypoints_array
    global videos
    logging.info('on_swapadult')
    v = videos[pickvid.value][pickcam.value]["v"]
    c = videos[pickvid.value][pickcam.value]["c"]
    end  = int(videos[pickvid.value][pickcam.value]["end"])
    vasc.swapSeries(keypoints_array,v,c,1,adult.value,slider.value,end)
    updateAll(True)

slider.observe(slider_change, 'value')
pickvid.observe(pickvid_change, 'value')
pickcam.observe(pickcam_change, 'value')
button_exclude.on_click(on_deleteparticipant)
button_swapcam.on_click(on_swapcam)
button_swapchild.on_click(on_swapchild)
button_swapadult.on_click(on_swapadult)
button_fixsizes.on_click(on_fixsizes)
button_fixlocations.on_click(on_fixlocations)
button_remove.on_click(on_deleteseries)
button_update.on_click(on_button_clicked)
button_reset_all.on_click(on_reset_all)

###############################################################
## ## functions to draw complicated stuff..
def drawOneFrame(vid, cam, frameNum):
    # which subarray of data do we need?
    v = videos[vid][cam]["v"]
    c = videos[vid][cam]["c"]
    if anon == True:
        #draw a black image
        frame = np.zeros((videos[vid][cam]["height"], videos[vid][cam]["width"], 3), dtype = "uint8")
    else:
        vidpath = videos[pickvid.value][pickcam.value]["fullpath"]
        frame = vasc.getframeimage(vidpath,frameNum)
    vasc.drawPoints(frame,keypoints_array[v,c,frameNum,:,:],videos[vid][cam]["maxpeople"])
    vasc.drawLines(frame,keypoints_array[v,c,frameNum,:,:],videos[vid][cam]["maxpeople"])
    vasc.drawBodyCG(frame,keypoints_array[v,c,frameNum,:,:],videos[vid][cam]["maxpeople"])
    #send the image to the canvas
    img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    hiddencanvas = Canvas(width=img.shape[1], height=img.shape[0])
    hiddencanvas.put_image_data(img, 0, 0)
    canvas.draw_image(hiddencanvas,0,0,canvas.width,canvas.height)
    canvas.restore()

def drawMovementGraph(vid, cam, points, frame = 0, average = True):
    v = videos[vid][cam]["v"]
    c = videos[vid][cam]["c"]
    N = videos[vid][cam]["frames"]
    t = np.zeros([N,1])
    t[:,0]= list(range(N))

    #variable to track the centre of gravity for each person
    ceegees = np.zeros([N,videos[vid][cam]["maxpeople"]])

    for frameNum in range(N):
        for p in range(videos[vid][cam]["maxpeople"]):
            personkeypoints = keypoints_array[v,c,frameNum,p,:]
            avx = vasc.averagePoint(personkeypoints,vasc.xs)
            if (avx > 0):
                ceegees[frameNum,p] = avx
            else:
                ceegees[frameNum,p] = None

    plt.figure(figsize=(12, 4))
    plt.plot(t,ceegees)
    plt.axvline(x=frame,c='tab:cyan')
    plt.title('Horizontal movement of people (average) over time.)')
    plt.legend([0, 1, 2, 3])
    plt.show()

###############################################################
## Handy update routine to run each time something has changed
def updateAll(forceUpdate = False):
    output.clear_output(wait = True)
    if forceUpdate:
        slider.value = 0
        slider.max = videos[pickvid.value][pickcam.value]["end"]
    with output:
        display(canvas,pickvid,cambox, babybox,adultbox,removebox, buttonbox, slider, adjustbox)
        drawOneFrame(pickvid.value,pickcam.value,slider.value)
        drawMovementGraph(pickvid.value,pickcam.value,vasc.xs,slider.value,True)

#draw everything for first time
updateAll(True)
output
# -

# ### Step 2.4: TODO - Correct for camera motion?
#
# Some video sets the camera is not fixed. Any camera movements will cause perfectly correlated movements in the pair of signals. We need to decide what (if anything) to do about this. (Not yet implemented.)
#
#

# ### Step 2.5: TODO - Interpolate missing data
#
# There are still likely to be gaps. We need to decide what to do about those.  At the moment interpolation is done by scipy in the Step 3 code.
#
# #### Step 2.5.1. TODO - autofix to cope with missing data
#
# Missing data currently confuses autofix and on it's own interpolation won't help here. Because you can't interpolate until you know who is who. Our current approach is to let autofix by location use a moving average of several previous frames.

# ### Step 2.6: TODO - Save your game
#
# Ought to be able to save the array when you half way through cleaning it. So you don't lose progress can come back another time.
# You can do this manually at the moment by running step 2.7 to save current progress and then restarting by reloading `cleandata.npz` in step 2.2.1

keypoints_array.shape

# ## Step 2.7: Save the numpy data!
#
# ### *Warning, with BIG datasets, these steps can take multiple minutes each...*
#
# Saving the data at this stage so we don't have to repeat these steps again if we reorganise or reanalyse the data.
#
# We create a compressed NumPy array `cleandata.npz` containing the person location data for all the videos.
#
# We also update the `videos.json` file with more info about the videos. in a new file called `clean.json`.
#

# +
#update the json file in the video out directory
with open(videos_out + '\\clean.json', 'w') as outfile:
    json.dump(videos, outfile)

# in the time series folder we save the data file.
#in a compressed format as it has a lot of empty values
np.savez_compressed(videos_out_timeseries + '\\cleandata.npz', keypoints_array=keypoints_array)
# -

# ## Step 2.8: Save a pandas dataframe version too.
#
# Most of our analysis will be done with SciPy which uses pandas dataframes as its main data format. So let's build a multiindex dataframe containing just the data we need.
#
# The rows will have three levels of hierarchy (video x person x BODY25-coordinate). The rows are the individual frames. So a single column will contain the complete time-series of a single dimension of a single point of one person.  So in this example:
# ```
# rows 0-411 represent the 412 frames of data.
#
# col 0 is x-coordinate of point 0 (nose) of infant in video 'lookit.01'
# col 1 is y-coordinate of point 0 (nose) of infant in video 'lookit.01'
# col 2 is openpose confidence score for how well it identified that point.
# ```
#
# <img src="multiindexdataframe.png" alt="multiindex" width="871"/>
#

#optional
#can reload the clean values without recomputing steps above
reloaded = np.load(videos_out_timeseries + '\\cleandata.npz')
keypoints_array = reloaded["keypoints_array"] #the unprocessed data
keypoints_array.shape

#delete all cameras except 0
keypoints_array = np.delete(keypoints_array,np.s_[1:],1)
#delete all people except 0 & 1
keypoints_array = np.delete(keypoints_array,np.s_[2:],3)

#truncate the timeseries - many videos are longer than we need
keypoints_array = np.delete(keypoints_array,np.s_[10000:],2)
shp = keypoints_array.shape
keypoints_array.shape

# Another save point - this array is much smaller so will load / Save quicker.

np.savez_compressed(videos_out_timeseries + '\\trimdata.npz', keypoints_array=keypoints_array)

#Another save point here if it helps.
trimmed = np.load(videos_out_timeseries + '\\trimdata.npz')
keypoints_array = trimmed["keypoints_array"] #the unprocessed data
keypoints_array.shape

# Now we reorganise the data in a multiindex pandas array and save using `pyarrow`.
# First create an empty dataframe with right shape

# +
#first list the three levels of row hierarchy
toplevel = videos.keys()
participants = ["infant","parent"]
coords = list(range(3*vasc.nPoints)) #we have 3 x 25 coordinates to store

#columns are frames
timeseries = list(range(shp[2])) #how big is third dimension of the array?

col_names = ['video','person','coord']
#row_names = ['frames']

col_index = pd.MultiIndex.from_product([toplevel,participants,coords], names=col_names)

cleandf = pd.DataFrame(columns=col_index, index = timeseries)
#cleandf.head()
# -

# Then populate the dataframe row by row.
#
# *This step is particularly SLOW*

for vid in videos:
    for p in range(2) :
        v = videos[vid]["camera1"]["v"]
        part = participants[p]
        for r in range(3*vasc.nPoints):
            cleandf[(vid, part, r)] = keypoints_array[v,0,:,p,r]

#Sort the columns into alphabetical order (helps with step 3 calculations.)
cleandf = cleandf.sort_index(axis = 1)

# ### Finally save this to a compressed file.
#
# We use the fast `parquet` format with library `pyarrow` in order to preserve our hierarchical index in a compressed format. We save into the timeseries sub-folder.
#

import pyarrow.parquet as pq
import pyarrow as pa

pq.write_table(pa.Table.from_pandas(cleandf), videos_out_timeseries + '\\cleandata.parquet')

print('reading parquet file:')
pqdf = pq.read_table(videos_out_timeseries + '\\cleandata.parquet').to_pandas()
print(pqdf.head())

# #### That's it.
#
# Now go onto [Step 3 - Analyse the data](Step3.AnalyseData.scipy)


import pandas as pd
import numpy as np
import io
import pandas_profiling
retail_raw = pd.read_csv('https://storage.googleapis.com/dqlab-dataset/retail_raw_reduced_data_quality.csv')

# Cetak tipe data di setiap kolom retail_raw
print(retail_raw.dtypes)


# -*- coding: utf-8 -*-

import pandas as pd
import numpy as np

data = np.array(["Engin","Derin","Salih"])
s = pd.Series(data, index=[1,2,3])
print(s)
print(s[1])

data2 = {"matematik":10, "fizik":20, "beden eğitimi":100}
s2 = pd.Series(data2, index = ["fizik","matematik","beden eğitimi"])
print(s2)

print(s2[0])
print(s2["matematik"])

s3 = pd.Series(5,index=[1,2,3,4,5])
print(s3)


#!/usr/bin/python
# -*- coding utf-8 -*-

#
#  Start-Programm von zufall
#
#  Copyright (C) 2019 Holger Böttcher <hbomat@posteo.de>
#

_version_ = '0.1.0'

# -----------------------------------------------------------------------------
# Variable zur Steuerung des Testbetriebes
# -----------------------------------------------------------------------------
#
# für den Programmierer; hier im Quelltext ein-/ausschalten
# _TEST = True:  es werden die vollständigen Python-Fehlermitteilungen angezeigt
# _TEST = False: es werden (teilweise) zufall-eigene Fehlermitteilungen angezeigt

_TEST = True

# -----------------------------------------------------------------------------
# Zurücksetzen
# -----------------------------------------------------------------------------

_ip = get_ipython()
_ip.magic('reset -sf')

# -----------------------------------------------------------------------------
# Importe
# -----------------------------------------------------------------------------

import copy
import importlib

from IPython.core.inputtransformer import InputTransformer

from sympy.core.expr import Expr
from sympy.matrices import Matrix as SympyMatrix

# ausgewählte SymPy-Objekte und -funktionen zur unmittelbaren Verfügung des Nutzers
from sympy.abc import *
from sympy.core.numbers import Integer, Rational, Float, pi, E, I
from sympy.core.symbol import Symbol, symbols
from sympy.core.function import Lambda
from sympy import (solve, solveset, S, expand, collect, factor, simplify,
    nsimplify, N, Not, And, Or, factorial, binomial)

# Initialisierung der Latex-Ausgabe
from sympy.interactive.printing import init_printing
from sympy.printing.latex import LatexPrinter
LatexPrinter._default_settings["mat_delim"] = "("
init_printing(use_latex='mathjax')

# Import ausgewählter zufall-Elemente
from zufall.lib.objekte.umgebung import UMG
from zufall.lib.objekte.ausnahmen import ZufallError
import zufall

from agla.lib.objekte.ausnahmen import AglaError
import agla

# Alle anderen Objekte und Funktionen werden dynamisch nachgeladen (s.u.)

# -----------------------------------------------------------------------------
# Variable zur Steuerung von internen Vereinfachungen
# -----------------------------------------------------------------------------
#
# Systemvariable, zur interaktiven Benutzung durch den Anwender
# UMG.SIMPL = True:  in den Programmen vorgesehene Vereinfachungen werden
#                    automatisch durchgeführt
# UMG.SIMPL = False: es werden keine automatischen Vereinfachungen durch-
#                    geführt

UMG.SIMPL = True

# -----------------------------------------------------------------------------
# Variable zur Steuerung der Berechnung von hyperbolischen Objekten
# -----------------------------------------------------------------------------
#
# Systemvariable, zur interaktiven Benutzung durch den Anwender
# UMG.EXAKT = True:  die Berechnungen werden exakt (mit SymPY / SympyEngine)
#                    durchgeführt
# UMG.EXAKT = False: die Berechnungen werden mit float-Werten durchgeführt

UMG.EXAKT = False

# -----------------------------------------------------------------------------
# Dynamisches Importieren
# -----------------------------------------------------------------------------

_attr = 'zufall.lib.objekte.'
_fkt = 'zufall.lib.funktionen.funktionen'
_agla = 'agla.lib.objekte.'

_to_load = {

    'v'             : _agla + 'vektor',
    'Vektor'        : _agla + 'vektor',
	'Punkt'         : _agla + 'vektor',
	'O'             : _agla + 'vektor',
    'O2'            : _agla + 'vektor',
	'X2'            : _agla + 'vektor',
    '_ZeilenVektor' : _agla + 'vektor',

	'SympyMatrix'   : _attr + 'markoff_kette',
	'Matrix'        : _attr + 'markoff_kette',
    'NullMat'       : _attr + 'markoff_kette',
    'NullMat2'      : _attr + 'markoff_kette',
    'EinhMat'       : _attr + 'markoff_kette',
    'EinhMat2'      : _attr + 'markoff_kette',

    'DatenReihe'                  : _attr + 'datenreihe',
    'DR'                          : _attr + 'datenreihe',
    'ZufallsGroesse'              : _attr + 'zufalls_groesse',
	'ZG'                          : _attr + 'zufalls_groesse',
	'BinomialVerteilung'          : _attr + 'binomial_verteilung',
	'BV'                          : _attr + 'binomial_verteilung',
	'HyperGeometrischeVerteilung' : _attr + 'hyper_geometrische_verteilung',
	'HGV'                         : _attr + 'hyper_geometrische_verteilung',
	'GeometrischeVerteilung'      : _attr + 'geometrische_verteilung',
	'GV'                          : _attr + 'geometrische_verteilung',
	'PoissonVerteilung'           : _attr + 'poisson_verteilung',
	'PV'                          : _attr + 'poisson_verteilung',
	'GleichVerteilung'            : _attr + 'gleich_verteilung',
	'GLV'                         : _attr + 'gleich_verteilung',
	'NormalVerteilung'            : _attr + 'normal_verteilung',
	'NV'                          : _attr + 'normal_verteilung',
	'ExponentialVerteilung'       : _attr + 'exponential_verteilung',
	'EV'                          : _attr + 'exponential_verteilung',
	'BernoulliKette'              : _attr + 'bernoulli_kette',
	'BK'                          : _attr + 'bernoulli_kette',
	'ZufallsExperiment'           : _attr + 'zufalls_experiment',
	'ZE'                          : _attr + 'zufalls_experiment',
	'ZufallsVersuch'              : _attr + 'zufalls_experiment',
	'ZV'                          : _attr + 'zufalls_experiment',
	'Urne'                        : _attr + 'urne',
	'GluecksRad'                  : _attr + 'gluecks_rad',
	'Rad'                         : _attr + 'gluecks_rad',
	'GR'                          : _attr + 'gluecks_rad',
	'Muenze'                      : _attr + 'muenze',
	'Wuerfel'                     : _attr + 'wuerfel',
	'HaeufigkeitsBaum'            : _attr + 'haeufigkeits_baum',
	'HB'                          : _attr + 'haeufigkeits_baum',
	'VierFelderTafel'             : _attr + 'vier_felder_tafel',
	'VT'                          : _attr + 'vier_felder_tafel',
    'KonfidenzIntervall'          : _attr + 'konfidenz_intervall',
	'KI'                          : _attr + 'konfidenz_intervall',
	'EreignisAlgebra'             : _attr + 'ereignis_algebra',
	'EA'                          : _attr + 'ereignis_algebra',
    'AlternativTest'              : _attr + 'alternativ_test',
	'AT'                          : _attr + 'alternativ_test',
	'SignifikanzTestP'            : _attr + 'signifikanz_test_p',
	'STP'                         : _attr + 'signifikanz_test_p',
	'SkatBlatt'                   : _attr + 'skat_blatt',
	'Skat'                        : _attr + 'skat_blatt',
	'Roulette'                    : _attr + 'roulette',
	'Craps'                       : _attr + 'craps',
	'Toto'                        : _attr + 'toto',
	'Lotto'                       : _attr + 'lotto',
	'Chuck'                       : _attr + 'chuck',
	'ChuckALuck'                  : _attr + 'chuck',
	'MarkoffKette'                : _attr + 'markoff_kette',
	'MK'                          : _attr + 'markoff_kette',

    'permutationen'  : _fkt,
    'perm'           : _fkt,
	'kombinationen'  : _fkt,
    'komb'           : _fkt,
    'variationen'    : _fkt,
	'anzahl'         : _fkt,
    'anzahl_treffer' : _fkt,
	'anzahlTreffer'  : _fkt,
	'summe'          : _fkt,
	'zuf_zahl'       : _fkt,
	'zufZahl'        : _fkt,
	'fakultaet'      : _fkt,
	'fak'            : _fkt,
	'binomial'       : _fkt,
	'B'              : _fkt,
	'loese'          : _fkt,
	'Hilfe'          : _fkt,
	'auswahlen'      : _fkt,
	'gesetze'        : _fkt,
	'stochastisch'   : _fkt,
	'einfach'        : _fkt,
	'is_zahl'        : _fkt,
	'isZahl'         : _fkt,
	'mit_param'   	 : _fkt,
	'mitParam'       : _fkt,
    'jaNein'         : _fkt,
    'kurz_form'      : _fkt,
    'kurzForm'       : _fkt,
	'ja_nein'        : _fkt,
	'ja'             : _fkt,
	'Ja'             : _fkt,
	'nein'           : _fkt,
	'Nein'           : _fkt,
	'mit'            : _fkt,
	'Mit'            : _fkt,
	'ohne'           : _fkt,
	'Ohne'           : _fkt,

    'abs'            : _fkt,
	'sqrt'           : _fkt,
	'exp'            : _fkt,
	'ln'             : _fkt,
	'lg'             : _fkt,
	'log'            : _fkt,
    'sin'            : _fkt,
	'cos'            : _fkt,
	'tan'            : _fkt,
	'cot'            : _fkt,
	'sing'           : _fkt,
	'cosg'           : _fkt,
	'tang'           : _fkt,
	'cotg'           : _fkt,
    'arcsin'         : _fkt,
	'arccos'         : _fkt,
	'arctan'         : _fkt,
	'arccot'         : _fkt,
	'asin'           : _fkt,
	'acos'           : _fkt,
	'atan'           : _fkt,
	'acot'           : _fkt,
    'arcsing'        : _fkt,
	'arccosg'        : _fkt,
	'arctang'        : _fkt,
	'arccotg'        : _fkt,
	'asing'          : _fkt,
	'acosg'          : _fkt,
	'atang'          : _fkt,
	'acotg'          : _fkt,
    'sinh'           : _fkt,
	'cosh'           : _fkt,
	'tanh'           : _fkt,
	'arsinh'         : _fkt,
	'arcosh'         : _fkt,
	'artanh'         : _fkt,
	'asinh'          : _fkt,
	'acosh'          : _fkt,
	'atanh'          : _fkt,
    're'             : _fkt,
	'im'             : _fkt,
	'conjugate'      : _fkt,
	'konjugiert'     : _fkt,
	'max'            : _fkt,
	'min'            : _fkt,
	'deg'            : _fkt,
	'grad'           : _fkt,
	'rad'            : _fkt,
	'bog'            : _fkt

	}

# -----------------------------------------------------------------------------
# Geschützte Namen
# -----------------------------------------------------------------------------

zufall_namen = list(_to_load.keys())
del zufall_namen[zufall_namen.index('B')]

_protected_names = zufall_namen + [

        # Python-Namen
        'and', 'as', 'assert', 'break', 'class', 'continue',
        'def', 'del', 'elif', 'else', 'except', 'exec',
        'finally', 'float', 'for', 'from', 'global', 'if', 'import',
        'in', 'int', 'is', 'lambda', 'not', 'or', 'pass', 'print',
		'raise', 'return', 'try', 'while', 'with', 'yield',

        # SymPy-Namen
        'Rational', 'pi', 'Symbol', 'symbols', 'solve', 'solveset',
        'nsolve', 'expand', 'simplify', 'nsimplify', 'collect', 'diff',
        'factor', 're', 'im', 'min', 'max', 'conjugate', 'sympify',

        # griechische Buchstaben
        'alpha', 'beta', 'gamma', 'delta', 'epsilon', 'zeta', 'eta',
		 'theta', 'iota', 'kappa', 'lamda', 'mu', 'nu', 'xi', 'omicron',
		 'pi', 'rho', 'sigma', 'tau', 'upsilon', 'phi', 'chi', 'psi',
		 'omega'
      ]

# -----------------------------------------------------------------------------
# Funktion zum Laden der Moduln
# -----------------------------------------------------------------------------

def _load_modules(line):

    for obj in _to_load:
        if obj in line and _to_load[obj] != False:
            name = _to_load[obj]
            modul = importlib.import_module(name)
            attribut = getattr(modul, obj)
            _ip.push({obj : attribut})
            _to_load[obj] = False

# -----------------------------------------------------------------------------
# Custom InputTransformer
# -----------------------------------------------------------------------------
#
#   Für
#
# - Dynamisches Laden der Moduln
#
# - Ersetzen der deutschen Umlaute und ß
#
#  - Operator '^'   für das Potenzieren (zusätzlich zu '**')
#    Operator '°'   für das Skalarprodukt und die Verknüpfung
#                   von Abbildungen (zusätzlich zu '*')
#    Operator '><'  für das Kreuzprodukt (Vektorprodukt)
#
# - Schreibwarnung für die zufall-, Python- und SymPy-Namen
#
#  (s.a. 'ipython/core/test/test_interactiveshell.py')

class ZufallInputTransformer(InputTransformer):

    def push(self, line):
        """für InputTransformer erforderlich"""

		# Operatoren, Umlaute
        #
        if line.find('^') >= 0:
            line = line.replace('^', '**')
        if line.find('°') >= 0:
            line = line.replace('°', '*')
        if line.find('><') >= 0:
            line = line.replace('><', '&')

        if line.find('ä') >= 0:
            line = line.replace('ä', 'ae')
        if line.find('ö') >= 0:
            line = line.replace('ö', 'oe')
        if line.find('ü') >= 0:
            line = line.replace('ü', 'ue')
        if line.find('ß') >= 0:
            line = line.replace('ß', 'ss')

        # Dynamisches Laden der Moduln
        #
        _load_modules(line)

        # Schreibblockade
        #

        def klammer(line):    # Index der 1. Klammer oder None
            for i, z in enumerate(line):
                if z in '([{':
                    return i
        def gleich(line):   # Indices aller '='-Zeichen vor der 1. Klammer
            line1 = line.replace('==', 'xx')	 # '==' ausschließen
            k = klammer(line1)
            g = []
            for i, z in enumerate(line1):
                if i == k:
                    break
                if z == '=':
                    g.append(i)
            return g

        if '#' in line:
            k = line.find('#')
            kl = False
            for i in range(k):
                if line[i] in '([{)]}':
                    kl = not kl
            if not kl:		# innere Kommentare bleiben erhalten
               line = line[:k]
        anw_liste = line.split(';')  # Behandlung von Mehrfachzuweisungen

        anw_liste_neu = []
        for i, anw in enumerate(anw_liste):
            gl = gleich(anw)
            no1 = anw.find("'")
            no2 = anw.find('"')
            for i, g in enumerate(gl):
                if abs(no1) < gl[i] or abs(no2) < gl[i]:
                    gl = gl[:i]
                    break
            if len(gl) <= 1:
                anw_liste_neu.append(anw)
                continue
            lhs, rhs = anw[:gl[-1]], anw[gl[-1]:]
            ind = ''
            for i in lhs:
                if i == ' ':
                    ind += ' '
                else:
                    break
            lhs = lhs.split('=')
            anw_neu = []
            for i, name in enumerate(lhs):
                if i == 0:
                    anw_liste_neu.append(name + rhs)
                else:
                    anw_liste_neu.append(ind + name[1:] + rhs)
        anw_liste = anw_liste_neu

        aus_liste = []
        for anw in anw_liste:
            anw = anw.rstrip()
            g = anw.find('=')
            if g <= 0:
                aus_liste += [anw]
                continue
            try:
                if anw[g+1] == '=':
                    aus_liste += [anw]
                    continue
            except IndexError:
                pass
            klammer = False
            for i in range(len(anw)):
                if i == g:
                    break
                if anw[i] in '([{)]}':
                    klammer = not klammer
            if klammer:
                aus_liste += [anw]
                continue
            if anw[g-1] in ('<', '>', '!'):
                aus_liste += [anw]
                continue
            bezeichner = anw[:g].strip()
            fehler = False
            if bezeichner in _protected_names:
                print("zufall: der Wert von " + bezeichner + " kann nicht überschrieben werden")
                fehler = True
            if anw[0] == '_':
                print("zufall: ein Bezeichner darf nicht mit einem Unterstrich beginnen")
                fehler = True
            p = anw.find('.')
            if 0 <= p < g and klammer and not 'UMG.' in anw:
                print("zufall: auf der linken Seite einer Zuweisung darf kein '.' auftreten")
                fehler = True
            if not fehler:
                aus_liste += [anw]
                continue
            else:
                return

        return ';'.join(aus_liste)

    def reset(self):
        """für InputTransformer erforderlich"""
        pass

transformer = ZufallInputTransformer()
_ip.input_splitter.python_line_transforms.append(transformer)
_ip.input_transformer_manager.python_line_transforms.append(transformer)

# -----------------------------------------------------------------------------
# Eigene Fehlermeldungen
# -----------------------------------------------------------------------------

# Extra-Behandlung von SyntaxError (hat kein traceback)
#

import linecache
from IPython.core.ultratb import SyntaxTB
from IPython.utils import py3compat
from IPython.utils import ulinecache
def new_structured_traceback(self, etype, value, elist, tb_offset=None,
                             context=5):
    if isinstance(value, SyntaxError) \
            and isinstance(value.filename, py3compat.string_types) \
            and isinstance(value.lineno, int):
        linecache.checkcache(value.filename)
        newtext = ulinecache.getline(value.filename, value.lineno)
        if newtext:
            value.text = newtext
    nr = value.lineno
    print('SyntaxError:  ' + value.text[:-1] + (('  /Zeile ' + str(nr) + '/') \
	      if nr > 1 else ''))
    return

if not _TEST:
    SyntaxTB.structured_traceback = new_structured_traceback

# Custom exc_handler (für ausgewählte Ausnahmen)
#
def exc_handler(shell, etype, value, tb, tb_offset=None):
    stype, sval = str(etype), str(value)
    if 'ZufallError' in stype or 'AglaError' in stype:
        print(sval)
    elif 'NameError' in stype:
        txt = sval[sval.find("'")+1 : sval.rfind("'")]
        print('NameError: ' + txt + ' ist nicht definiert')
    elif 'KeyError' in stype:
        print('KeyError: ' + sval)
    elif 'IndexError' in stype:
        print('IndexError: ' + 'der Index ist außerhalb des Bereiches')
    elif 'AttributeError' in stype:
        no = []
        for i, z in enumerate(sval):
            if z == "'":
                no += [i]
        if len(no) != 4:
            print('AttributeError')
        else:
            obj, eig = sval[no[0]+1:no[1]], sval[no[2]+1:no[3]]
            print('AttributeError:  ' + 'ein ' + obj + '-Objekt ' + \
                 'hat keine Eigenschaft/Methode \nmit dem Namen ' + eig)
    elif 'TypeError' in stype:
        print('TypeError: ' + 'Funktionsaufruf nicht möglich oder anderer Fehler')
    elif 'MemoryError' in stype:
        print('MemoryError: ' + 'die Zahlen sind zu groß')
    elif 'ImportError' in stype:
        i = sval.find("'")
        j = sval[i+1:].find("'")
        txt = sval[i+1:i+j+1]
        print('ImportError: ' + txt +  ' kann nicht importiert werden')
    elif 'ZeroDivisionError' in stype:
        print('ZeroDivisionError: ' +	 'Division durch Null ist nicht erlaubt')
    elif 'IndentationError' in stype:
        print('IndentationError: ' +	 'Syntaxfehler - falsche Einrückung')
    elif 'RuntimeError' in stype:
        print('RuntimeError: ' +	    'Ein Laufzeitfehler ist aufgetreten')
    elif 'ValueError' in stype:
        print('ValueError: ' +	    'Ein ValueError ist aufgetreten')
    return

if not _TEST:
    _ip.set_custom_exc((ZufallError, AglaError, NameError, KeyError, IndexError, \
      AttributeError, TypeError, MemoryError, ImportError, ZeroDivisionError, \
      IndentationError, RuntimeError), exc_handler)

# -----------------------------------------------------------------------------
# Ausführung der (linksseitigen) Multiplikation  Zahl * Vektor
# -----------------------------------------------------------------------------

"""
Die rechtsseitige Multiplikation eines Vektors mit einer Zahl
(Vektor * Zahl) funktioniert für alle infrage kommenden SymPy-Klassen,
das Ergebnis ist ein Vektor-Objekt
Hinsichtlich der linksseitigen Multiplikation (Zahl * Vektor) verhalten
sich die Klassen unterschiedlich. Bei einigen gibt die __mul__-Methode
NotImplemented zurück, sodass automatisch die __rmul__-Methode von Vektor
aufgerufen wird. Für die anderen wird dieses Verhalten mittels Überladen
der __mul__-Methode des Expr-Objektes bzw. ihrer eigenen __mul__-Methode
(Symbol) erreicht
"""

# -----------------------------------------------------------------------------
# Überladen der Expr.__mul__- Methode
# -----------------------------------------------------------------------------

_expr_mul = copy.deepcopy(Expr.__mul__)
def _expr_mul_anpassung(self, other):
    Vektor = importlib.import_module('agla.lib.objekte.vektor').Vektor
    if isinstance(other, (Vektor, SympyMatrix)):
        return NotImplemented
    return _expr_mul(self, other)
Expr.__mul__ = _expr_mul_anpassung

# -----------------------------------------------------------------------------
# Überladen der Symbol.__mul__- Methode
# -----------------------------------------------------------------------------

_symbol_mul = copy.deepcopy(Symbol.__mul__)
def _symbol_mul_anpassung(self, other):
    Vektor = importlib.import_module('agla.lib.objekte.vektor').Vektor
    if isinstance(other, (Vektor, SympyMatrix)):
        return NotImplemented
    return _symbol_mul(self, other)
Symbol.__mul__ = _symbol_mul_anpassung

# -----------------------------------------------------------------------------
# Float.__mul__ - Überladung
# -----------------------------------------------------------------------------

_Float_mul = copy.deepcopy(Float.__mul__)

def _Float_mul_anpassung(self, other):
    if isinstance(other, SympyMatrix):
        ve = [self*v for v in other.vekt]
        return Matrix(*ve)
    return _Float_mul(self, other)

Float.__mul__ = _Float_mul_anpassung