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

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

# Qiskit
from qiskit import BasicAer, Aer
from qiskit.aqua import QuantumInstance
from qiskit.aqua.components.optimizers import COBYLA, SPSA
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

def hamiltonian_varying_V(min_V, max_V, points, epsilons, n_trials):
    solution_VQD = {}
    energies_Classical = []
    
    V = np.linspace(min_V, max_V, points)
    
    backend = QuantumInstance(backend=BasicAer.get_backend('qasm_simulator'),
                              shots=10000)
    
    optimizer = SPSA(maxiter=250, c1=.7, last_avg=25) 
    #optimizer = COBYLA()
    
    for v in tqdm(V):
        hamiltonian = BCS_hamiltonian(epsilons, v)
        ansatz = TwoLocal(hamiltonian.num_qubits, ['ry','rz'], 'cx', reps=3)
        
        es_1 = []
        es_2 = []

        for _ in tqdm(range(n_trials), desc=f"V= {v}"):

            # 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)
    
        solution_VQD[v] = {'mean':np.mean(es_2 - es_1), 'std':np.std(es_2 - es_1)}
        
        classical, _ = classical_solver(hamiltonian)
        energies_Classical.append(classical[2]- classical[1])
    
    return solution_VQD, np.array(energies_Classical), V

min_V = -0.
max_V = -2.
points = 5
epsilons = [3,3]
n_trials = 50
solution_VQD, energy_classical, V = hamiltonian_varying_V(min_V, max_V, points, epsilons, n_trials)

solution_VQD

import seaborn as sns
sns.set()

mean = []
std = []
for _, sol in solution_VQD.items():
    mean.append(sol['mean'])
    std.append(sol['std'])

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

plt.plot(V, energy_classical/2, 'b-', label="Exact")
plt.xlabel('V')
plt.ylabel('Gap')
plt.legend()
plt.savefig('Var_V.png')
plt.show()

V

spsa_dict = {-0.0: {'mean': 2.2544600000000004, 'std': 3.6059371187529052},
 -0.5: {'mean': 2.201132, 'std': 3.5924339523749076},
 -1.0: {'mean': 2.3685160000000005, 'std': 3.2362400972956262},
 -1.5: {'mean': 2.262069, 'std': 3.6215824729514305},
 -2.0: {'mean': 2.151896, 'std': 3.285247713884601}}


cobyla_dict = {-0.0: {'mean': 0.2738039999999999, 'std': 0.5017597014348602},
 -0.5: {'mean': 1.1608589999999999, 'std': 0.7228235554884747},
 -1.0: {'mean': 2.0350799999999993, 'std': 0.8327042377699295},
 -1.5: {'mean': 3.137112, 'std': 0.6724548732487554},
 -2.0: {'mean': 4.314758, 'std': 0.5284894406097439}}

mean_spsa = []
std_spsa = []
mean_cobyla = []
std_cobyla = []

x_axis = np.array([-0. , -0.5, -1. , -1.5, -2. ])
for sol in x_axis:
    mean_spsa.append(spsa_dict[sol]['mean'])
    std_spsa.append(spsa_dict[sol]['std'])
    
    mean_cobyla.append(cobyla_dict[sol]['mean'])
    std_cobyla.append(cobyla_dict[sol]['std'])

from matplotlib.ticker import MaxNLocator

x_axis = np.array([-0. , -0.5, -1. , -1.5, -2. ])
energy_classical = np.array([0., 1., 2., 3., 4.])

plt.errorbar(x_axis,
             np.array(mean_spsa)/2, 
             yerr=np.array(std_spsa)/2,
             fmt='o', color='black', 
             ecolor='gray', elinewidth=2,
             label='SPSA')

plt.errorbar(x_axis,
             np.array(mean_cobyla)/2, 
             yerr=np.array(std_cobyla)/2,
             fmt='o', color='red',
             ecolor='green', elinewidth=2,
             label='COBYLA')

plt.plot(V, energy_classical/2, 'b-', label="Exact")
plt.xlabel('V')
plt.ylabel('Gap')
plt.legend()
plt.savefig('Var_V_optimizers.png')
plt.show()

