import sys
sys.path.append('../../')

# Python imports
import numpy as np
import matplotlib.pyplot as plt

# Qiskit
from qiskit import BasicAer
from qiskit.aqua.components.optimizers import COBYLA, SPSA, L_BFGS_B
from qiskit.circuit.library import TwoLocal

# VOLTA
from volta.vqd import VQD
from volta.utils import classical_solver
from volta.hamiltonians import BCS_hamiltonian

%load_ext autoreload
%autoreload 2

EPSILONS = [3, 3]
V = -2
hamiltonian = BCS_hamiltonian(EPSILONS, V)
print(hamiltonian)

eigenvalues, eigenvectors = classical_solver(hamiltonian)
print(f"Eigenvalues: {eigenvalues}")

from tqdm import tqdm
from qiskit.aqua import QuantumInstance

# Parameters Variables
n_trials = 50

# Auxiliary Variables
solution_dict = {}

# Define Optimizer
optimizers = [COBYLA(), 
              SPSA(maxiter=250, c1=.7, last_avg=25), 
              #L_BFGS_B(),
             ]

optimizer_names = ['COBYLA', 
                   'SPSA', 
                   #'L_BFGS_B'
                  ]

# Define Backend
backend = QuantumInstance(backend=BasicAer.get_backend('qasm_simulator'),
                          shots=10_000)

# Ansatz with diferent depth
ansatz = TwoLocal(hamiltonian.num_qubits, ['ry','rz'], 'cx', reps=3)

for i, optimizer in enumerate(optimizers):
    
    es_1 = []
    es_2 = []
    
    for _ in tqdm(range(n_trials), desc=f"Optimizer {optimizer_names[i]}"):

        # Run algorithm
        Algo = VQD(hamiltonian=hamiltonian,
                    ansatz=ansatz,
                    n_excited_states=2,
                    beta=10.,
                    optimizer=optimizer,
                    backend=backend)

        Algo.run(0)      
        vqd_energies = Algo.energies
        es_1.append(vqd_energies[1])
        es_2.append(vqd_energies[2])
    es_1 = np.array(es_1)
    es_2 = np.array(es_2)
    
    # Maybe use a pd.dataframe
    solution_dict[i] = {'mean': np.mean(es_2 - es_1), 
                        'std':np.std(es_2 - es_1),
                       }

solution_dict

mean = []
std = []
for i in range(len(optimizers)):
    mean.append(solution_dict[i]['mean'])
    std.append(solution_dict[i]['std'])

import seaborn as sns
sns.set()

from matplotlib.ticker import MaxNLocator

x_axis = optimizer_names

plt.errorbar(x_axis,
             np.array(mean)/2, 
             yerr=np.array(std)/2,
             fmt='ro', ecolor='green')

plt.hlines(y=2, xmin=0., xmax=len(optimizers)+ 0.25, 
           label='Expected value',color='b')

plt.title('Optimizers', size=18)
plt.xlabel('Optimizer', size= 14)
plt.ylabel('Gap', size=14)
plt.xticks(x_axis)
plt.legend()
plt.show()

