# Copyright 2018-2020 Xanadu Quantum Technologies Inc.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Unit tests for the :mod:`pennylane` optimizers.
"""
# pylint: disable=redefined-outer-name
import itertools as it

import numpy as onp
import pytest

import pennylane as qml
from pennylane import numpy as np
from pennylane.optimize import (
    GradientDescentOptimizer,
    MomentumOptimizer,
    NesterovMomentumOptimizer,
    AdagradOptimizer,
    RMSPropOptimizer,
    AdamOptimizer,
    RotoselectOptimizer,
    RotosolveOptimizer,
)


x_vals = np.linspace(-10, 10, 16, endpoint=False)

# Hyperparameters for optimizers
stepsize = 0.1
gamma = 0.5
delta = 0.8


# function arguments in various formats
mixed_list = [(0.2, 0.3), np.array([0.4, 0.2, 0.4]), 0.1]
mixed_tuple = (np.array([0.2, 0.3]), [0.4, 0.2, 0.4], 0.1)
flat_list = [0.2, 0.3, 0.1, 0.4, -0.1]
multid_array = np.array([[0.1, 0.2], [-0.1, -0.4]])
multid_list = [[0.1, 0.2], [-0.1, -0.4]]


# functions and their gradients
fnames = ["test_function_1", "test_function_2", "test_function_3"]
univariate_funcs = [np.sin, lambda x: np.exp(x / 10.0), lambda x: x ** 2]
grad_uni_fns = [lambda x: (np.cos(x),), lambda x: (np.exp(x / 10.0) / 10.0,), lambda x: (2 * x,)]

multivariate_funcs = [
    lambda x: np.sin(x[0]) + np.cos(x[1]),
    lambda x: np.exp(x[0] / 3) * np.tanh(x[1]),
    lambda x: np.sum([x_ ** 2 for x_ in x]),
]
grad_multi_funcs = [
    lambda x: (np.array([np.cos(x[0]), -np.sin(x[1])]),),
    lambda x: (
        np.array(
            [np.exp(x[0] / 3) / 3 * np.tanh(x[1]), np.exp(x[0] / 3) * (1 - np.tanh(x[1]) ** 2)]
        ),
    ),
    lambda x: (np.array([2 * x_ for x_ in x]),),
]

mvar_mdim_funcs = [
    lambda x: np.sin(x[0, 0]) + np.cos(x[1, 0]) - np.sin(x[0, 1]) + x[1, 1],
    lambda x: np.exp(x[0, 0] / 3) * np.tanh(x[0, 1]),
    lambda x: np.sum([x_[0] ** 2 for x_ in x]),
]
grad_mvar_mdim_funcs = [
    lambda x: (np.array([[np.cos(x[0, 0]), -np.cos(x[0, 1])], [-np.sin(x[1, 0]), 1.0]]),),
    lambda x: (
        np.array(
            [
                [
                    np.exp(x[0, 0] / 3) / 3 * np.tanh(x[0, 1]),
                    np.exp(x[0, 0] / 3) * (1 - np.tanh(x[0, 1]) ** 2),
                ],
                [0.0, 0.0],
            ]
        ),
    ),
    lambda x: (np.array([[2 * x_[0], 0.0] for x_ in x]),),
]


@qml.qnode(qml.device("default.qubit", wires=1))
def quant_fun(*variables):
    qml.RX(variables[0][1], wires=[0])
    qml.RY(variables[1][2], wires=[0])
    qml.RY(variables[2], wires=[0])
    return qml.expval(qml.PauliZ(0))


@qml.qnode(qml.device("default.qubit", wires=1))
def quant_fun_flat(var):
    qml.RX(var[0], wires=[0])
    qml.RY(var[1], wires=[0])
    qml.RY(var[2], wires=[0])
    qml.RX(var[3], wires=[0])
    return qml.expval(qml.PauliZ(0))


@qml.qnode(qml.device("default.qubit", wires=1))
def quant_fun_mdarr(var):
    qml.RX(var[0, 1], wires=[0])
    qml.RY(var[1, 0], wires=[0])
    qml.RY(var[1, 1], wires=[0])
    return qml.expval(qml.PauliZ(0))


@qml.qnode(qml.device("default.qubit", wires=1))
def quant_fun_mdlist(var):
    qml.RX(var[0][1], wires=[0])
    qml.RY(var[1][0], wires=[0])
    qml.RY(var[1][1], wires=[0])
    return qml.expval(qml.PauliZ(0))


@pytest.fixture(scope="function")
def bunch():
    class A:
        sgd_opt = GradientDescentOptimizer(stepsize)
        mom_opt = MomentumOptimizer(stepsize, momentum=gamma)
        nesmom_opt = NesterovMomentumOptimizer(stepsize, momentum=gamma)
        adag_opt = AdagradOptimizer(stepsize)
        rms_opt = RMSPropOptimizer(stepsize, decay=gamma)
        adam_opt = AdamOptimizer(stepsize, beta1=gamma, beta2=delta)
        rotosolve_opt = RotosolveOptimizer()
        rotoselect_opt = RotoselectOptimizer()

    return A()


class TestOptimizer:
    """Basic optimizer tests."""

    def test_mixed_inputs_for_hybrid_optimization(self, bunch, tol):
        """Tests that gradient descent optimizer treats parameters of mixed types the same
        for hybrid optimization tasks."""

        def hybrid_fun(variables):
            return quant_fun(*variables) + variables[0][1]

        hybrid_list = bunch.sgd_opt.step(hybrid_fun, mixed_list)
        hybrid_tuple = bunch.sgd_opt.step(hybrid_fun, mixed_tuple)

        assert hybrid_list[0] == pytest.approx(hybrid_tuple[0], abs=tol)
        assert hybrid_list[1] == pytest.approx(hybrid_tuple[1], abs=tol)
        assert hybrid_list[2] == pytest.approx(hybrid_tuple[2], abs=tol)

    def test_mixed_inputs_for_classical_optimization(self, bunch, tol):
        """Tests that gradient descent optimizer treats parameters of mixed types the same
        for purely classical optimization tasks."""

        def class_fun(var):
            return var[0][1] * 2.0 + var[1][2] + var[2]

        class_list = bunch.sgd_opt.step(class_fun, mixed_list)
        class_tuple = bunch.sgd_opt.step(class_fun, mixed_tuple)

        assert class_list[0] == pytest.approx(class_tuple[0], abs=tol)
        assert class_list[1] == pytest.approx(class_tuple[1], abs=tol)
        assert class_list[2] == pytest.approx(class_tuple[2], abs=tol)

    def test_mixed_inputs_for_quantum_optimization(self, bunch, tol):
        """Tests that gradient descent optimizer treats parameters of mixed types the same
        for purely quantum optimization tasks."""

        quant_list = bunch.sgd_opt.step(quant_fun, *mixed_list)
        quant_tuple = bunch.sgd_opt.step(quant_fun, *mixed_tuple)

        assert quant_list[0] == pytest.approx(quant_tuple[0], abs=tol)
        assert quant_list[1] == pytest.approx(quant_tuple[1], abs=tol)
        assert quant_list[2] == pytest.approx(quant_tuple[2], abs=tol)

    def test_array_and_list_return_same_update(self, bunch, tol):
        """Tests that gradient descent optimizer has the same output for
        lists and arrays."""

        def hybrid_fun_mdarr(var):
            return quant_fun_mdarr(var) + var[0, 0]

        def hybrid_fun_mdlist(var):
            return quant_fun_mdlist(var) + var[0][0]

        array = bunch.sgd_opt.step(hybrid_fun_mdarr, multid_array)
        ls = bunch.sgd_opt.step(hybrid_fun_mdlist, multid_list)

        assert array == pytest.approx(np.asarray(ls), abs=tol)

    def test_step_and_cost_autograd_sgd_mixed_list(self, bunch):
        """Test that the correct cost is returned via the step_and_cost method for the
        gradient-descent optimizer"""
        _, res = bunch.sgd_opt.step_and_cost(quant_fun, *mixed_list)
        expected = quant_fun(*mixed_list)

        assert np.all(res == expected)

    def test_step_and_cost_autograd_sgd_multid_array(self, bunch):
        """Test that the correct cost is returned via the step_and_cost method for the
        gradient-descent optimizer"""
        _, res = bunch.sgd_opt.step_and_cost(quant_fun_mdarr, multid_array)
        expected = quant_fun_mdarr(multid_array)

        assert np.all(res == expected)

    def test_step_and_cost_autograd_nesterov_mixed_list(self, bunch):
        """Test that the correct cost is returned via the step_and_cost method for the
        Nesterov momentum optimizer"""
        _, res = bunch.nesmom_opt.step_and_cost(quant_fun, *mixed_list)
        expected = quant_fun(*mixed_list)

        assert np.all(res == expected)

    def test_step_and_cost_autograd_nesterov_multid_array(self, bunch):
        """Test that the correct cost is returned via the step_and_cost method for the
        Nesterov momentum optimizer"""
        _, res = bunch.nesmom_opt.step_and_cost(quant_fun_mdarr, multid_array)
        expected = quant_fun_mdarr(multid_array)

        assert np.all(res == expected)

    def test_step_and_cost_autograd_rotosolve_mixed_list(self, bunch):
        """Test that the correct cost is returned via the step_and_cost method for the
        Rotosolve optimizer"""
        _, res = bunch.rotosolve_opt.step_and_cost(quant_fun, *mixed_list)
        expected = quant_fun(*mixed_list)

        assert np.all(res == expected)

    def test_step_and_cost_autograd_rotosolve_multid_array(self, bunch):
        """Test that the correct cost is returned via the step_and_cost method for the
        Rotosolve optimizer"""
        _, res = bunch.rotosolve_opt.step_and_cost(quant_fun_mdarr, multid_array)
        expected = quant_fun_mdarr(multid_array)

        assert np.all(res == expected)

    @pytest.mark.parametrize("params", [[1.7, 2.2], [-1.42, 0.1], [0.05, -0.8]])
    def test_step_and_cost_autograd_rotoselect(self, bunch, params):
        """Test that the correct cost is returned via the step_and_cost method for the
        Rotoselect momentum optimizer"""
        generators = [qml.RY, qml.RX]
        possible_generators = [qml.RX, qml.RY, qml.RZ]
        bunch.rotoselect_opt.possible_generators = possible_generators

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

        def ansatz(params, generators):
            generators[0](params[0], wires=0)
            generators[1](params[1], wires=1)
            qml.CNOT(wires=[0, 1])

        @qml.qnode(dev)
        def circuit_1(params, generators=None):  # generators will be passed as a keyword arg
            ansatz(params, generators)
            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))

        @qml.qnode(dev)
        def circuit_2(params, generators=None):  # generators will be passed as a keyword arg
            ansatz(params, generators)
            return qml.expval(qml.PauliX(0))

        def cost_fn(params, generators):
            Z_1, Y_2 = circuit_1(params, generators=generators)
            X_1 = circuit_2(params, generators=generators)
            return 0.5 * Y_2 + 0.8 * Z_1 - 0.2 * X_1

        _, _, res = bunch.rotoselect_opt.step_and_cost(cost_fn, params, generators)
        expected = cost_fn(params, generators)

        assert np.all(res == expected)

    @pytest.mark.parametrize("func, f_grad", list(zip(univariate_funcs, grad_uni_fns)))
    @pytest.mark.parametrize("var", [0, -3, 42])
    def test_step_and_cost_supplied_grad(self, bunch, func, var, f_grad):
        """Test that returned cost is correct if gradient function is supplied"""
        _, res = bunch.sgd_opt.step_and_cost(func, var, grad_fn=f_grad)
        expected = func(var)

        assert np.all(res == expected)

    @pytest.mark.parametrize("x_start", x_vals)
    def test_gradient_descent_optimizer_univar(self, x_start, bunch, tol):
        """Tests that basic stochastic gradient descent takes gradient-descent steps correctly
        for uni-variate functions."""

        # TODO parametrize this for also
        for gradf, f, name in zip(grad_uni_fns, univariate_funcs, fnames):
            x_new = bunch.sgd_opt.step(f, x_start)
            x_correct = x_start - gradf(x_start)[0] * stepsize
            assert x_new == pytest.approx(x_correct, abs=tol)

    def test_gradient_descent_optimizer_multivar(self, bunch, tol):
        """Tests that basic stochastic gradient descent takes gradient-descent steps correctly
        for multi-variate functions."""

        for gradf, f, name in zip(grad_multi_funcs, multivariate_funcs, fnames):
            for jdx in range(len(x_vals[:-1])):
                x_vec = x_vals[jdx : jdx + 2]
                x_new = bunch.sgd_opt.step(f, x_vec)
                x_correct = x_vec - gradf(x_vec)[0] * stepsize
                assert x_new == pytest.approx(x_correct, abs=tol)

    def test_gradient_descent_optimizer_multivar_multidim(self, bunch, tol):
        """Tests that basic stochastic gradient descent takes gradient-descent steps correctly
        for multi-variate functions and with higher dimensional inputs."""

        for gradf, f, name in zip(grad_mvar_mdim_funcs, mvar_mdim_funcs, fnames):
            for jdx in range(len(x_vals[:-3])):
                x_vec = x_vals[jdx : jdx + 4]
                x_vec_multidim = np.reshape(x_vec, (2, 2))
                x_new = bunch.sgd_opt.step(f, x_vec_multidim)
                x_correct = x_vec_multidim - gradf(x_vec_multidim)[0] * stepsize
                x_new_flat = x_new.flatten()
                x_correct_flat = x_correct.flatten()
                assert x_new_flat == pytest.approx(x_correct_flat, abs=tol)

    @pytest.mark.parametrize("x_start", x_vals)
    def test_gradient_descent_optimizer_usergrad(self, x_start, bunch, tol):
        """Tests that basic stochastic gradient descent takes gradient-descent steps correctly
        using user-provided gradients."""

        for gradf, f, name in zip(grad_uni_fns[::-1], univariate_funcs, fnames):
            x_new = bunch.sgd_opt.step(f, x_start, grad_fn=gradf)
            x_correct = x_start - gradf(x_start)[0] * stepsize
            assert x_new == pytest.approx(x_correct, abs=tol)

    @pytest.mark.parametrize("x_start", x_vals)
    def test_momentum_optimizer_univar(self, x_start, bunch, tol):
        """Tests that momentum optimizer takes one and two steps correctly
        for uni-variate functions."""

        for gradf, f, name in zip(grad_uni_fns, univariate_funcs, fnames):
            bunch.mom_opt.reset()

            x_onestep = bunch.mom_opt.step(f, x_start)
            x_onestep_target = x_start - gradf(x_start)[0] * stepsize
            assert x_onestep == pytest.approx(x_onestep_target, abs=tol)

            x_twosteps = bunch.mom_opt.step(f, x_onestep)
            momentum_term = gamma * gradf(x_start)[0]
            x_twosteps_target = x_onestep - (gradf(x_onestep)[0] + momentum_term) * stepsize
            assert x_twosteps == pytest.approx(x_twosteps_target, abs=tol)

    def test_momentum_optimizer_multivar(self, bunch, tol):
        """Tests that momentum optimizer takes one and two steps correctly
        for multi-variate functions."""

        for gradf, f, name in zip(grad_multi_funcs, multivariate_funcs, fnames):
            for jdx in range(len(x_vals[:-1])):
                bunch.mom_opt.reset()

                x_vec = x_vals[jdx : jdx + 2]
                x_onestep = bunch.mom_opt.step(f, x_vec)
                x_onestep_target = x_vec - gradf(x_vec)[0] * stepsize
                assert x_onestep == pytest.approx(x_onestep_target, abs=tol)

                x_twosteps = bunch.mom_opt.step(f, x_onestep)
                momentum_term = gamma * gradf(x_vec)[0]
                x_twosteps_target = x_onestep - (gradf(x_onestep)[0] + momentum_term) * stepsize
                assert x_twosteps == pytest.approx(x_twosteps_target, abs=tol)

    @pytest.mark.parametrize("x_start", x_vals)
    def test_nesterovmomentum_optimizer_univar(self, x_start, bunch, tol):
        """Tests that nesterov momentum optimizer takes one and two steps correctly
        for uni-variate functions."""

        for gradf, f, name in zip(grad_uni_fns, univariate_funcs, fnames):
            bunch.nesmom_opt.reset()

            x_onestep = bunch.nesmom_opt.step(f, x_start)
            x_onestep_target = x_start - gradf(x_start)[0] * stepsize
            assert x_onestep == pytest.approx(x_onestep_target, abs=tol)

            x_twosteps = bunch.nesmom_opt.step(f, x_onestep)
            momentum_term = gamma * gradf(x_start)[0]
            shifted_grad_term = gradf(x_onestep - stepsize * momentum_term)[0]
            x_twosteps_target = x_onestep - (shifted_grad_term + momentum_term) * stepsize
            assert x_twosteps == pytest.approx(x_twosteps_target, abs=tol)

    def test_nesterovmomentum_optimizer_multivar(self, bunch, tol):
        """Tests that nesterov momentum optimizer takes one and two steps correctly
        for multi-variate functions."""

        for gradf, f, name in zip(grad_multi_funcs, multivariate_funcs, fnames):
            for jdx in range(len(x_vals[:-1])):
                bunch.nesmom_opt.reset()

                x_vec = x_vals[jdx : jdx + 2]
                x_onestep = bunch.nesmom_opt.step(f, x_vec)
                x_onestep_target = x_vec - gradf(x_vec)[0] * stepsize
                assert x_onestep == pytest.approx(x_onestep_target, abs=tol)

                x_twosteps = bunch.nesmom_opt.step(f, x_onestep)
                momentum_term = gamma * gradf(x_vec)[0]
                shifted_grad_term = gradf(x_onestep - stepsize * momentum_term)[0]
                x_twosteps_target = x_onestep - (shifted_grad_term + momentum_term) * stepsize
                assert x_twosteps == pytest.approx(x_twosteps_target, abs=tol)

    @pytest.mark.parametrize("x_start", x_vals)
    def test_nesterovmomentum_optimizer_usergrad(self, x_start, bunch, tol):
        """Tests that nesterov momentum optimizer takes gradient-descent steps correctly
        using user-provided gradients."""

        for gradf, f, name in zip(grad_uni_fns[::-1], univariate_funcs, fnames):
            bunch.nesmom_opt.reset()

            x_onestep = bunch.nesmom_opt.step(f, x_start, grad_fn=gradf)
            x_onestep_target = x_start - gradf(x_start)[0] * stepsize
            assert x_onestep == pytest.approx(x_onestep_target, abs=tol)

            x_twosteps = bunch.nesmom_opt.step(f, x_onestep, grad_fn=gradf)
            momentum_term = gamma * gradf(x_start)[0]
            shifted_grad_term = gradf(x_onestep - stepsize * momentum_term)[0]
            x_twosteps_target = x_onestep - (shifted_grad_term + momentum_term) * stepsize
            assert x_twosteps == pytest.approx(x_twosteps_target, abs=tol)

    @pytest.mark.parametrize("x_start", x_vals)
    def test_adagrad_optimizer_univar(self, x_start, bunch, tol):
        """Tests that adagrad optimizer takes one and two steps correctly
        for uni-variate functions."""

        for gradf, f, name in zip(grad_uni_fns, univariate_funcs, fnames):
            bunch.adag_opt.reset()

            x_onestep = bunch.adag_opt.step(f, x_start)
            past_grads = gradf(x_start)[0] * gradf(x_start)[0]
            adapt_stepsize = stepsize / np.sqrt(past_grads + 1e-8)
            x_onestep_target = x_start - gradf(x_start)[0] * adapt_stepsize
            assert x_onestep == pytest.approx(x_onestep_target, abs=tol)

            x_twosteps = bunch.adag_opt.step(f, x_onestep)
            past_grads = (
                gradf(x_start)[0] * gradf(x_start)[0] + gradf(x_onestep)[0] * gradf(x_onestep)[0]
            )
            adapt_stepsize = stepsize / np.sqrt(past_grads + 1e-8)
            x_twosteps_target = x_onestep - gradf(x_onestep)[0] * adapt_stepsize
            assert x_twosteps == pytest.approx(x_twosteps_target, abs=tol)

    def test_adagrad_optimizer_multivar(self, bunch, tol):
        """Tests that adagrad optimizer takes one and two steps correctly
        for multi-variate functions."""

        for gradf, f, name in zip(grad_multi_funcs, multivariate_funcs, fnames):
            for jdx in range(len(x_vals[:-1])):
                bunch.adag_opt.reset()

                x_vec = x_vals[jdx : jdx + 2]
                x_onestep = bunch.adag_opt.step(f, x_vec)
                past_grads = gradf(x_vec)[0] * gradf(x_vec)[0]
                adapt_stepsize = stepsize / np.sqrt(past_grads + 1e-8)
                x_onestep_target = x_vec - gradf(x_vec)[0] * adapt_stepsize
                assert x_onestep == pytest.approx(x_onestep_target, abs=tol)

                x_twosteps = bunch.adag_opt.step(f, x_onestep)
                past_grads = (
                    gradf(x_vec)[0] * gradf(x_vec)[0] + gradf(x_onestep)[0] * gradf(x_onestep)[0]
                )
                adapt_stepsize = stepsize / np.sqrt(past_grads + 1e-8)
                x_twosteps_target = x_onestep - gradf(x_onestep)[0] * adapt_stepsize
                assert x_twosteps == pytest.approx(x_twosteps_target, abs=tol)

    @pytest.mark.parametrize("x_start", x_vals)
    def test_rmsprop_optimizer_univar(self, x_start, bunch, tol):
        """Tests that rmsprop optimizer takes one and two steps correctly
        for uni-variate functions."""

        for gradf, f, name in zip(grad_uni_fns, univariate_funcs, fnames):
            bunch.rms_opt.reset()

            x_onestep = bunch.rms_opt.step(f, x_start)
            past_grads = (1 - gamma) * gradf(x_start)[0] * gradf(x_start)[0]
            adapt_stepsize = stepsize / np.sqrt(past_grads + 1e-8)
            x_onestep_target = x_start - gradf(x_start)[0] * adapt_stepsize
            assert x_onestep == pytest.approx(x_onestep_target, abs=tol)

            x_twosteps = bunch.rms_opt.step(f, x_onestep)
            past_grads = (1 - gamma) * gamma * gradf(x_start)[0] * gradf(x_start)[0] + (
                1 - gamma
            ) * gradf(x_onestep)[0] * gradf(x_onestep)[0]
            adapt_stepsize = stepsize / np.sqrt(past_grads + 1e-8)
            x_twosteps_target = x_onestep - gradf(x_onestep)[0] * adapt_stepsize
            assert x_twosteps == pytest.approx(x_twosteps_target, abs=tol)

    def test_rmsprop_optimizer_multivar(self, bunch, tol):
        """Tests that rmsprop optimizer takes one and two steps correctly
        for multi-variate functions."""

        for gradf, f, name in zip(grad_multi_funcs, multivariate_funcs, fnames):
            for jdx in range(len(x_vals[:-1])):
                bunch.rms_opt.reset()

                x_vec = x_vals[jdx : jdx + 2]
                x_onestep = bunch.rms_opt.step(f, x_vec)
                past_grads = (1 - gamma) * gradf(x_vec)[0] * gradf(x_vec)[0]
                adapt_stepsize = stepsize / np.sqrt(past_grads + 1e-8)
                x_onestep_target = x_vec - gradf(x_vec)[0] * adapt_stepsize
                assert x_onestep == pytest.approx(x_onestep_target, abs=tol)

                x_twosteps = bunch.rms_opt.step(f, x_onestep)
                past_grads = (1 - gamma) * gamma * gradf(x_vec)[0] * gradf(x_vec)[0] + (
                    1 - gamma
                ) * gradf(x_onestep)[0] * gradf(x_onestep)[0]
                adapt_stepsize = stepsize / np.sqrt(past_grads + 1e-8)
                x_twosteps_target = x_onestep - gradf(x_onestep)[0] * adapt_stepsize
                assert x_twosteps == pytest.approx(x_twosteps_target, abs=tol)

    @pytest.mark.parametrize("x_start", x_vals)
    def test_adam_optimizer_univar(self, x_start, bunch, tol):
        """Tests that adam optimizer takes one and two steps correctly
        for uni-variate functions."""

        for gradf, f, name in zip(grad_uni_fns, univariate_funcs, fnames):
            bunch.adam_opt.reset()

            x_onestep = bunch.adam_opt.step(f, x_start)
            adapted_stepsize = stepsize * np.sqrt(1 - delta) / (1 - gamma)
            firstmoment = gradf(x_start)[0]
            secondmoment = gradf(x_start)[0] * gradf(x_start)[0]
            x_onestep_target = x_start - adapted_stepsize * firstmoment / (
                np.sqrt(secondmoment) + 1e-8
            )
            assert x_onestep == pytest.approx(x_onestep_target, abs=tol)

            x_twosteps = bunch.adam_opt.step(f, x_onestep)
            adapted_stepsize = stepsize * np.sqrt(1 - delta ** 2) / (1 - gamma ** 2)
            firstmoment = gamma * gradf(x_start)[0] + (1 - gamma) * gradf(x_onestep)[0]
            secondmoment = (
                delta * gradf(x_start)[0] * gradf(x_start)[0]
                + (1 - delta) * gradf(x_onestep)[0] * gradf(x_onestep)[0]
            )
            x_twosteps_target = x_onestep - adapted_stepsize * firstmoment / (
                np.sqrt(secondmoment) + 1e-8
            )
            assert x_twosteps == pytest.approx(x_twosteps_target, abs=tol)

    def test_adam_optimizer_multivar(self, bunch, tol):
        """Tests that adam optimizer takes one and two steps correctly
        for multi-variate functions."""

        for gradf, f, name in zip(grad_multi_funcs, multivariate_funcs, fnames):
            for jdx in range(len(x_vals[:-1])):
                bunch.adam_opt.reset()

                x_vec = x_vals[jdx : jdx + 2]
                x_onestep = bunch.adam_opt.step(f, x_vec)
                adapted_stepsize = stepsize * np.sqrt(1 - delta) / (1 - gamma)
                firstmoment = gradf(x_vec)[0]
                secondmoment = gradf(x_vec)[0] * gradf(x_vec)[0]
                x_onestep_target = x_vec - adapted_stepsize * firstmoment / (
                    np.sqrt(secondmoment) + 1e-8
                )
                assert x_onestep == pytest.approx(x_onestep_target, abs=tol)

                x_twosteps = bunch.adam_opt.step(f, x_onestep)
                adapted_stepsize = stepsize * np.sqrt(1 - delta ** 2) / (1 - gamma ** 2)
                firstmoment = gamma * gradf(x_vec)[0] + (1 - gamma) * gradf(x_onestep)[0]
                secondmoment = (
                    delta * gradf(x_vec)[0] * gradf(x_vec)[0]
                    + (1 - delta) * gradf(x_onestep)[0] * gradf(x_onestep)[0]
                )
                x_twosteps_target = x_onestep - adapted_stepsize * firstmoment / (
                    np.sqrt(secondmoment) + 1e-8
                )
                assert x_twosteps == pytest.approx(x_twosteps_target, abs=tol)

    @staticmethod
    def rotosolve_step(f, x):
        """Helper function to test the Rotosolve and Rotoselect optimizers"""
        # make sure that x is an array
        if np.ndim(x) == 0:
            x = np.array([x])

        # helper function for x[d] = theta
        def insert(xf, d, theta):
            xf[d] = theta
            return xf

        for d, _ in enumerate(x):
            H_0 = float(f(insert(x, d, 0)))
            H_p = float(f(insert(x, d, np.pi / 2)))
            H_m = float(f(insert(x, d, -np.pi / 2)))
            a = onp.arctan2(2 * H_0 - H_p - H_m, H_p - H_m)

            x[d] = -np.pi / 2 - a

            if x[d] <= -np.pi:
                x[d] += 2 * np.pi
        return x

    @pytest.mark.parametrize("x_start", x_vals)
    def test_rotosolve_optimizer_univar(self, x_start, bunch, tol):
        """Tests that rotosolve optimizer takes one and two steps correctly
        for uni-variate functions."""

        for f in univariate_funcs:
            x_onestep = bunch.rotosolve_opt.step(f, x_start)
            x_onestep_target = self.rotosolve_step(f, x_start)

            assert np.allclose(x_onestep, x_onestep_target, atol=tol, rtol=0)

            x_twosteps = bunch.rotosolve_opt.step(f, x_onestep)
            x_twosteps_target = self.rotosolve_step(f, x_onestep_target)

            assert np.allclose(x_twosteps, x_twosteps_target, atol=tol, rtol=0)

    @pytest.mark.parametrize(
        "x_start",
        [
            [1.2, 0.2],
            [-0.62, -2.1],
            [0.05, 0.8],
            [[0.3], [0.25]],
            [[-0.6], [0.45]],
            [[1.3], [-0.9]],
        ],
    )
    def test_rotosolve_optimizer_multivar(self, x_start, bunch, tol):
        """Tests that rotosolve optimizer takes one and two steps correctly
        for multi-variate functions."""

        for func in multivariate_funcs:
            # alter multivariate_func to accept nested lists of parameters
            f = lambda x: func(np.ravel(x))

            x_onestep = bunch.rotosolve_opt.step(f, x_start)
            x_onestep_target = self.rotosolve_step(f, x_start)

            assert x_onestep == pytest.approx(x_onestep_target, abs=tol)

            x_twosteps = bunch.rotosolve_opt.step(f, x_onestep)
            x_twosteps_target = self.rotosolve_step(f, x_onestep_target)

            assert x_twosteps == pytest.approx(x_twosteps_target, abs=tol)

    @pytest.mark.parametrize("x_start", [[1.2, 0.2], [-0.62, -2.1], [0.05, 0.8]])
    @pytest.mark.parametrize(
        "generators", [list(tup) for tup in it.product([qml.RX, qml.RY, qml.RZ], repeat=2)]
    )
    def test_rotoselect_optimizer(self, x_start, generators, bunch, tol):
        """Tests that rotoselect optimizer finds the optimal generators and parameters for the VQE circuit
        defined in `this rotoselect tutorial <https://pennylane.ai/qml/demos/tutorial_rotoselect.html>`_."""

        # the optimal generators for the 2-qubit VQE circuit
        # H = 0.5 * Y_2 + 0.8 * Z_1 - 0.2 * X_1
        optimal_generators = [qml.RY, qml.RX]
        possible_generators = [qml.RX, qml.RY, qml.RZ]
        bunch.rotoselect_opt.possible_generators = possible_generators

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

        def ansatz(params, generators):
            generators[0](params[0], wires=0)
            generators[1](params[1], wires=1)
            qml.CNOT(wires=[0, 1])

        @qml.qnode(dev)
        def circuit_1(params, generators=None):  # generators will be passed as a keyword arg
            ansatz(params, generators)
            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))

        @qml.qnode(dev)
        def circuit_2(params, generators=None):  # generators will be passed as a keyword arg
            ansatz(params, generators)
            return qml.expval(qml.PauliX(0))

        def cost_fn(params, generators):
            Z_1, Y_2 = circuit_1(params, generators=generators)
            X_1 = circuit_2(params, generators=generators)
            return 0.5 * Y_2 + 0.8 * Z_1 - 0.2 * X_1

        f_best_gen = lambda x: cost_fn(x, optimal_generators)
        optimal_x_start = x_start.copy()

        # after four steps the optimzer should find the optimal generators/x_start values
        for _ in range(4):
            x_start, generators = bunch.rotoselect_opt.step(cost_fn, x_start, generators)
            optimal_x_start = self.rotosolve_step(f_best_gen, optimal_x_start)

        assert x_start == pytest.approx(optimal_x_start, abs=tol)
        assert generators == optimal_generators

    @pytest.mark.parametrize("x_start", [[1.2, 0.2], [-0.62, -2.1], [0.05, 0.8]])
    def test_keywords_rotoselect(self, bunch, x_start, tol):
        """test rotoselect accepts keywords"""

        generators = [qml.RY, qml.RX]
        possible_generators = [qml.RX, qml.RY, qml.RZ]
        bunch.rotoselect_opt.possible_generators = possible_generators

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

        def ansatz(params, generators):
            generators[0](params[0], wires=0)
            generators[1](params[1], wires=1)
            qml.CNOT(wires=[0, 1])

        @qml.qnode(dev)
        def circuit_1(params, generators=None):  # generators will be passed as a keyword arg
            ansatz(params, generators)
            return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))

        @qml.qnode(dev)
        def circuit_2(params, generators=None):  # generators will be passed as a keyword arg
            ansatz(params, generators)
            return qml.expval(qml.PauliX(0))

        def cost_fn(params, generators, shift=0.0):
            Z_1, Y_2 = circuit_1(params, generators=generators)
            X_1 = circuit_2(params, generators=generators)
            return 0.5 * (Y_2 - shift) ** 2 + 0.8 * (Z_1 - shift) ** 2 - 0.2 * (X_1 - shift) ** 2

        params_new, _, res_new = bunch.rotoselect_opt.step_and_cost(
            cost_fn, x_start, generators, shift=0.0
        )
        params_new2, _, res_new2 = bunch.rotoselect_opt.step_and_cost(
            cost_fn, x_start, generators, shift=1.0
        )

        assert params_new != pytest.approx(params_new2, abs=tol)
        assert res_new2 == pytest.approx(cost_fn(x_start, generators, shift=1.0), abs=tol)

    def test_update_stepsize(self):
        """Tests that the stepsize correctly updates"""

        eta = 0.5
        opt = AdamOptimizer(eta)
        assert opt._stepsize == eta

        eta2 = 0.1
        opt.update_stepsize(eta2)
        assert opt._stepsize == eta2


def reset(opt):
    if getattr(opt, "reset", None):
        opt.reset()


@pytest.fixture
def opt(opt_name):
    if opt_name == "gd":
        return GradientDescentOptimizer(stepsize)

    if opt_name == "nest":
        return NesterovMomentumOptimizer(stepsize, momentum=gamma)

    if opt_name == "moment":
        return MomentumOptimizer(stepsize, momentum=gamma)

    if opt_name == "ada":
        return AdagradOptimizer(stepsize)

    if opt_name == "rms":
        return RMSPropOptimizer(stepsize, decay=gamma)

    if opt_name == "adam":
        return AdamOptimizer(stepsize, beta1=gamma, beta2=delta)

    if opt_name == "roto":
        return RotosolveOptimizer()


@pytest.mark.parametrize(
    "opt_name",
    [
        "gd",
        "moment",
        "nest",
        "ada",
        "rms",
        "adam",
        "roto",
    ],
)
class TestOverOpts:
    """Tests keywords, multiple arguements, and non-training arguments in relevent optimizers"""

    def test_kwargs(self, mocker, opt, opt_name, tol):
        """Test that the keywords get passed and alter the function"""

        class func_wrapper:
            @staticmethod
            def func(x, c=1.0):
                return (x - c) ** 2

        x = 1.0

        wrapper = func_wrapper()
        spy = mocker.spy(wrapper, "func")

        x_new_two = opt.step(wrapper.func, x, c=2.0)
        reset(opt)

        args2, kwargs2 = spy.call_args_list[-1]

        x_new_three_wc, cost_three = opt.step_and_cost(wrapper.func, x, c=3.0)
        reset(opt)

        args3, kwargs3 = spy.call_args_list[-1]

        if opt_name != "roto":
            assert args2 == (x,)
            assert args3 == (x,)
        else:
            assert x_new_two != pytest.approx(x, abs=tol)
            assert x_new_three_wc != pytest.approx(x, abs=tol)

        assert kwargs2 == {"c": 2.0}
        assert kwargs3 == {"c": 3.0}

        assert cost_three == pytest.approx(wrapper.func(x, c=3.0), abs=tol)

    def test_multi_args(self, mocker, opt, opt_name, tol):
        """Test passing multiple arguments to function"""

        class func_wrapper:
            @staticmethod
            def func(x, y, z):
                return x[0] * y[0] + z[0]

        wrapper = func_wrapper()
        spy = mocker.spy(wrapper, "func")

        x = np.array([1.0])
        y = np.array([2.0])
        z = np.array([3.0])

        (x_new, y_new, z_new), cost = opt.step_and_cost(wrapper.func, x, y, z)
        reset(opt)
        args_called1, kwargs1 = spy.call_args_list[-1]  # just take last call

        x_new2, y_new2, z_new2 = opt.step(wrapper.func, x_new, y_new, z_new)
        reset(opt)
        args_called2, kwargs2 = spy.call_args_list[-1]  # just take last call

        if opt_name != "roto":
            assert args_called1 == (x, y, z)
            assert args_called2 == (x_new, y_new, z_new)
        else:
            assert x_new != pytest.approx(x, abs=tol)
            assert y_new != pytest.approx(y, abs=tol)
            assert z_new != pytest.approx(z, abs=tol)

        assert kwargs1 == {}
        assert kwargs2 == {}

        assert cost == pytest.approx(wrapper.func(x, y, z), abs=tol)

    def test_nontrainable_data(self, opt, opt_name, tol):
        """Check non-trainable argument does not get updated"""

        def func(x, data):
            return x[0] * data[0]

        x = np.array([1.0])
        data = np.array([1.0], requires_grad=False)

        args_new = opt.step(func, x, data)
        reset(opt)
        args_new_wc, cost = opt.step_and_cost(func, *args_new)
        reset(opt)

        assert len(args_new) == pytest.approx(2, abs=tol)
        assert args_new[0] != pytest.approx(x, abs=tol)
        assert args_new[1] == pytest.approx(data, abs=tol)

        assert cost == pytest.approx(func(*args_new), abs=tol)

    def test_steps_the_same(self, opt, opt_name, tol):
        """Tests whether separating the args into different inputs affects their
        optimization step. Assumes single argument optimization is correct, as tested elsewhere."""

        def func1(x, y, z):
            return x[0] * y[0] * z[0]

        def func2(args):
            return args[0][0] * args[1][0] * args[2][0]

        x = np.array([1.0])
        y = np.array([2.0])
        z = np.array([3.0])
        args = (x, y, z)

        x_seperate, y_seperate, z_seperate = opt.step(func1, x, y, z)
        reset(opt)

        args_new = opt.step(func2, args)
        reset(opt)

        assert x_seperate == pytest.approx(args_new[0], abs=tol)
        assert y_seperate == pytest.approx(args_new[1], abs=tol)
        assert z_seperate == pytest.approx(args_new[2], abs=tol)

    def test_one_trainable_one_non_trainable(self, opt, opt_name, tol):
        """Tests that a cost function that takes one trainable and one
        non-trainable parameter executes well."""
        dev = qml.device("default.qubit", wires=2)

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

        def cost(x, target):
            return (circuit(x) - target) ** 2

        ev = np.tensor(0.7781, requires_grad=False)
        x = np.tensor(0.0, requires_grad=True)

        original_ev = ev

        (x, ev), cost = opt.step_and_cost(cost, x, ev)

        # check that the argument to RX doesn't change, as the X rotation doesn't influence <Z>
        assert x == 0
        assert ev == original_ev

    def test_one_non_trainable_one_trainable(self, opt, opt_name, tol):
        """Tests that a cost function that takes one non-trainable and one
        trainable parameter executes well."""
        dev = qml.device("default.qubit", wires=2)

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

        def cost(target, x):  # Note: the order of the arguments has been swapped
            return (circuit(x) - target) ** 2

        ev = np.tensor(0.7781, requires_grad=False)
        x = np.tensor(0.0, requires_grad=True)

        original_ev = ev

        (ev, x), cost = opt.step_and_cost(cost, ev, x)
        # check that the argument to RX doesn't change, as the X rotation doesn't influence <Z>
        assert x == 0
        assert ev == original_ev

    def test_two_trainable_args(self, opt, opt_name, tol):
        """Tests that a cost function that takes at least two trainable
        arguments executes well."""
        dev = qml.device("default.qubit", wires=2)

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

        def cost(x, y, target):
            return (circuit(x, y) - target) ** 2

        ev = np.tensor(0.7781, requires_grad=False)
        x = np.tensor(0.0, requires_grad=True)
        y = np.tensor(0.0, requires_grad=True)

        original_ev = ev

        (x, y, ev), cost = opt.step_and_cost(cost, x, y, ev)

        # check that the argument to RX doesn't change, as the X rotation doesn't influence <Z>
        assert x == 0
        assert ev == original_ev
