#   Copyright 2022 Entropica Labs
#
#   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.

"""
A set of functions to implement pennylane optimization algorithms.
Read https://docs.pennylane.ai/en/stable/introduction/interfaces.html#optimizers
Optimisers requiring a pennylane backend haven't been implemented yet.
Similarly as with the custom optimization methods Scipy `minimize` is used. Extends available scipy methods.
"""

from openqaoa.optimizers import pennylane as pl
import inspect
from scipy.optimize import OptimizeResult
import numpy as np

AVAILABLE_OPTIMIZERS = {  # optimizers implemented
    "pennylane_adagrad": pl.AdagradOptimizer,
    "pennylane_adam": pl.AdamOptimizer,
    "pennylane_vgd": pl.GradientDescentOptimizer,
    "pennylane_momentum": pl.MomentumOptimizer,
    "pennylane_nesterov_momentum": pl.NesterovMomentumOptimizer,
    "pennylane_rmsprop": pl.RMSPropOptimizer,
    "pennylane_rotosolve": pl.RotosolveOptimizer,
    "pennylane_spsa": pl.SPSAOptimizer,
}


def pennylane_optimizer(
    fun,
    x0,
    args=(),
    maxfev=None,
    pennylane_method="vgd",
    maxiter=100,
    tol=10 ** (-6),
    jac=None,
    callback=None,
    nums_frequency=None,
    spectra=None,
    shifts=None,
    **options
):
    """
    Minimize a function `fun` using some pennylane method.
    To check available methods look at the available_methods_dict variable.
    Read https://docs.pennylane.ai/en/stable/introduction/interfaces.html#optimizers

    PARAMETERS
    ----------
    fun : callable
        Function to minimize
    x0 : ndarray
        Initial guess.
    args : sequence, optional
        Arguments to pass to `func`.
    maxfev : int, optional
        Maximum number of function evaluations.
    pennylane_method : string, optional
        Optimizer method to compute the steps.
    maxiter : int, optional
        Maximum number of iterations.
    tol : float
        Tolerance before the optimizer terminates; if `tol` is larger than the difference between two steps, terminate optimization.
    jac : callable, optinal
        Callable gradient function. Required for all methods but rotosolve and spsa.
    callback : callable, optional
        Called after each iteration, as ``callback(xk)``, where ``xk`` is the
        current parameter vector.
    options : dict, optional
        Dictionary where keys are the arguments for the optimizers object, and
        the values are the values to pass to these arguments.
        To know all the possible options see https://docs.pennylane.ai/en/stable/introduction/interfaces.html#optimizers.
    nums_frequency : dict[dict], optional
        It is required for rotosolve method
        The number of frequencies in the fun per parameter.
    spectra : dict[dict], optional
        It is required for rotosolve method
        Frequency spectra in the objective_fn per parameter.
    shifts : dict[dict], optional
        It is required for rotosolve method
        Shift angles for the reconstruction per parameter.
        Read https://docs.pennylane.ai/en/stable/code/api/pennylane.RotosolveOptimizer.html#pennylane.RotosolveOptimizer.step for more information.


    RETURNS
    -------
    OptimizeResult : OptimizeResult
        Scipy OptimizeResult object.
    """

    def cost(
        params, **k
    ):  # define a function to convert the params list from pennylane to numpy
        return fun(np.array(params), *k)

    optimizer = AVAILABLE_OPTIMIZERS[pennylane_method]  # define the optimizer

    # get optimizer arguments
    arguments = inspect.signature(optimizer).parameters.keys()
    options_keys = list(options.keys())

    # check which values of the options dict can be passed to the optimizer (pop the others)
    for key in options_keys:
        if key not in arguments:
            options.pop(key)
        if "maxiter" in arguments:
            options["maxiter"] = maxiter

    optimizer = optimizer(**options)  # pass the arguments

    bestx = pl.numpy.array(x0, requires_grad=True)
    besty = cost(x0, *args)
    funcalls = 1  # tracks no. of function evals.
    niter = 0
    improved = True
    stop = False

    testx = np.copy(bestx)
    testy = np.real(besty)
    while improved and not stop and niter < maxiter:
        improved = False

        # compute step (depends on the optimizer)
        if pennylane_method in [
            "pennylane_adagrad",
            "pennylane_adam",
            "pennylane_vgd",
            "pennylane_momentum",
            "pennylane_nesterov_momentum",
            "pennylane_rmsprop",
        ]:
            testx, testy = optimizer.step_and_cost(cost, bestx, *args, grad_fn=jac)
        elif pennylane_method in ["pennylane_rotosolve"]:
            testx, testy = optimizer.step_and_cost(
                cost,
                bestx,
                *args,
                nums_frequency={"params": {(i,): 1 for i in range(bestx.size)}}
                if not nums_frequency
                else nums_frequency,
                spectra=spectra,
                shifts=shifts,
                full_output=False,
            )
        elif pennylane_method in ["pennylane_spsa"]:
            testx, testy = optimizer.step_and_cost(cost, bestx, *args)

        # check if stable
        if np.abs(besty - testy) < tol and niter > 1:
            improved = False

        else:
            besty = testy
            bestx = testx
            improved = True

        if callback is not None:
            callback(bestx)
        if maxfev is not None and funcalls >= maxfev:
            stop = True
            break

        niter += 1

    return OptimizeResult(
        fun=besty, x=np.array(bestx), nit=niter, nfev=funcalls, success=(niter > 1)
    )
