import warnings
import unittest
import numpy as np
import networkx as nx
import copy
import inspect

from openqaoa import QAOA
import openqaoa.optimizers.pennylane as pl
from openqaoa.backends import create_device
from openqaoa.optimizers.training_vqa import PennyLaneOptimizer
from openqaoa.optimizers.pennylane.optimization_methods_pennylane import AVAILABLE_OPTIMIZERS
from openqaoa.derivatives.derivative_functions import derivative
from openqaoa.optimizers.logger_vqa import Logger
from openqaoa.qaoa_components import create_qaoa_variational_params, QAOADescriptor, PauliOp, Hamiltonian
from openqaoa.utilities import X_mixer_hamiltonian
from openqaoa.backends.qaoa_backend import get_qaoa_backend
from openqaoa.optimizers import get_optimizer
from openqaoa.derivatives.qfim import qfim as Qfim
from openqaoa.problems import QUBO, MinimumVertexCover


#list of optimizers to test, pennylane optimizers
list_optimizers = PennyLaneOptimizer.PENNYLANE_OPTIMIZERS

#create a problem
g = nx.circulant_graph(4, [1])
problem = MinimumVertexCover(g, field =1.0, penalty=10)
qubo_problem_1 = problem.qubo
qubo_problem_2 = QUBO.random_instance(5)
qubo_problem_3 = QUBO.random_instance(6)


class TestPennylaneOptimizers(unittest.TestCase):

    def setUp(self):
        
        self.log = Logger({'func_evals': 
                           {
                               'history_update_bool': False, 
                               'best_update_string': 'HighestOnly'
                           },
                           'jac_func_evals':
                           {
                               'history_update_bool': False, 
                               'best_update_string': 'HighestOnly'
                           },
                           'qfim_func_evals': 
                           {
                               'history_update_bool': False, 
                               'best_update_string': 'HighestOnly'
                           }
                          }, 
                          {
                              'root_nodes': ['func_evals', 'jac_func_evals', 
                                             'qfim_func_evals'], 
                              'best_update_structure': []
                          })
        
        self.log.log_variables({'func_evals': 0, 'jac_func_evals': 0, 'qfim_func_evals': 0})

    def _run_method_workflows(self, method, problem):
        " helper function to run the test for any method using workflows"
        q = QAOA()
        q.set_classical_optimizer(method=method, maxiter=3, jac='finite_difference')
        q.compile(problem) 
        q.optimize()

        assert len(q.result.most_probable_states['solutions_bitstrings'][0]) > 0

    def _run_method_manual(self, method, problem):
        " helper function tu run the test for any method using manual mode"

        cost_hamil = problem.hamiltonian
        mixer_hamil = X_mixer_hamiltonian(n_qubits=problem.n)
        qaoa_descriptor = QAOADescriptor(cost_hamil, mixer_hamil, p=2)
        device = create_device('local','vectorized')
        backend_obj_vectorized = get_qaoa_backend(qaoa_descriptor,device)
        variate_params = create_qaoa_variational_params(qaoa_descriptor, 'standard', 'ramp')
        niter = 5
        grad_stepsize = 0.0001
        stepsize = 0.001

        # declare needed functions
        jac = derivative(backend_obj_vectorized, variate_params, self.log, 
                    'gradient', 'finite_difference', 
                    {'stepsize': grad_stepsize})
        qfim = Qfim(backend_obj_vectorized, variate_params, self.log)   


        # Optimize
        vector_optimizer = get_optimizer(backend_obj_vectorized, variate_params, optimizer_dict={
                                        'method': method, 'jac': jac, 'maxiter': niter, 'qfim': qfim,
                                        'optimizer_options' : {'stepsize': stepsize}})
        vector_optimizer()

        # saving the results
        results = vector_optimizer.qaoa_result

        assert len(results.most_probable_states['solutions_bitstrings'][0]) == problem.n

    def test_pennylane_optimizers_workflows(self):
        " function to run the tests for pennylane optimizers, workflows"

        i = 0
        for problem in [qubo_problem_3, qubo_problem_2, qubo_problem_1]:
            for opt in list_optimizers:
                self._run_method_workflows(opt, problem)
                i += 1

        assert i == 3*len(list_optimizers)

    def test_pennylane_optimizers_manual(self):
        " function to run the tests for pennylane optimizers, manual mode"

        i = 0
        for problem in [qubo_problem_3, qubo_problem_2, qubo_problem_1]:
            for opt in list_optimizers:
                self._run_method_manual(opt, problem)
                i += 1

        assert i == 3*len(list_optimizers)

    def _pennylane_step(self, params_array, cost, optimizer, method, jac, qfim):
        " helper function to run a setp of the pennylane optimizer"
        params_array = pl.numpy.array(params_array, requires_grad=True)
        if method in ['pennylane_adagrad', 'pennylane_adam', 'pennylane_vgd', 'pennylane_momentum', 'pennylane_nesterov_momentum', 'pennylane_rmsprop']:
            x, y = optimizer.step_and_cost(cost, params_array, grad_fn=jac)
        if method in ['pennylane_rotosolve']: 
            x, y = optimizer.step_and_cost(
                                            cost, params_array,
                                            nums_frequency={'params': {(i,):1 for i in range(params_array.size)}},
                                           )
        if method in ['pennylane_spsa']:       
            x, y = optimizer.step_and_cost(cost, params_array)

        return x, y

    def test_step_and_cost(self):
        " function to run the tests for steps of pennylane optimizers "

        # define some problem
        cost_hamil = Hamiltonian([PauliOp('ZZ', (0, 1)), PauliOp('ZZ', (1, 2)), PauliOp(
            'ZZ', (0, 3)), PauliOp('Z', (2,)), PauliOp('Z', (1,))], [1, 1.1, 1.5, 2, -0.8], 0.8)
        mixer_hamil = X_mixer_hamiltonian(n_qubits=4)
        qaoa_descriptor = QAOADescriptor(cost_hamil, mixer_hamil, p=2)
        device = create_device('local','vectorized')
        backend_obj_vectorized = get_qaoa_backend(qaoa_descriptor,device)
        variate_params = create_qaoa_variational_params(qaoa_descriptor, 'standard', 'ramp')
        niter = 5
        grad_stepsize = 0.0001
        stepsize = 0.001

        # declare needed functions
        jac = derivative(backend_obj_vectorized, variate_params, self.log, 
                    'gradient', 'finite_difference', 
                    {'stepsize': grad_stepsize})
        qfim = Qfim(backend_obj_vectorized, variate_params, self.log)   
        def cost(params):
            variate_params.update_from_raw(params)
            return np.real(backend_obj_vectorized.expectation(variate_params))

        i = 0
        for method in list_optimizers:
        
            pennylane_method = method

            # copy the parameters
            x0 = copy.deepcopy(variate_params.raw().copy())

            # Optimize with the implemented optimizer in OpenQAOA
            vector_optimizer = get_optimizer(backend_obj_vectorized, variate_params, optimizer_dict={
                                            'method': method, 'jac': jac, 'maxiter': niter, 'qfim': qfim,
                                            'optimizer_options' : {'stepsize': stepsize}})
            vector_optimizer()

            # formatting the data
            y_opt = vector_optimizer.qaoa_result.intermediate['cost'][1:4]
            if pennylane_method in ['pennylane_rotosolve']: y_opt = vector_optimizer.qaoa_result.intermediate['cost'][4:40:12]

            # get optimizer to try
            optimizer = AVAILABLE_OPTIMIZERS[pennylane_method]
            #get optimizer arguments
            arguments = inspect.signature(optimizer).parameters.keys()

            #check if stepsize is in the optimizer arguments
            options = {}
            if 'stepsize' in arguments: options['stepsize'] = stepsize
            if 'maxiter'  in arguments: options['maxiter'] = niter

            #pass the argument to the optimizer
            optimizer = optimizer(**options) 

            # reinitialize variables
            variate_params.update_from_raw(x0)
            x0 = variate_params.raw().copy()
            y0 = cost(x0)

            # compute steps (depends on the optimizer)
            x1, y1 = self._pennylane_step(x0, cost, optimizer, pennylane_method, jac, qfim)
            x2, y2 = self._pennylane_step(x1, cost, optimizer, pennylane_method, jac, qfim)
            x3, y3 = self._pennylane_step(x2, cost, optimizer, pennylane_method, jac, qfim)

            # list of results
            y = [y1, y2, y3]

            # check that the results are ok
            if pennylane_method in ['pennylane_spsa']: 
                assert np.sum(np.abs(np.array(y)) >= 0) == 3
            else:
                for yi, y_opt_i in zip(y, y_opt):
                    assert np.isclose(yi, y_opt_i, rtol=0.001, atol=0.001)

            i += 1                

        assert i == len(list_optimizers)


if __name__ == "__main__":
    with warnings.catch_warnings():
        warnings.simplefilter('ignore', category=PendingDeprecationWarning)
        unittest.main()
