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
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., 3., 4., 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.utils import QuantumInstance

# Parameters Variables
n_trials = 50
max_depth = 7

# Auxiliary Variables
solution_dict = {}

# Define Optimizer
# optimizer = COBYLA()
optimizer = SPSA(maxiter=250, c1=.7, last_avg=25) 

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

for depth in range(1,max_depth):
    # Ansatz with diferent depth
    ansatz = TwoLocal(hamiltonian.num_qubits, ['ry','rz'], 'cx', reps=depth)
    es_1 = []
    es_2 = []
    
    for _ in tqdm(range(n_trials), desc=f"Depth {depth}"):

        # Run algorithm
        Algo = VQD(hamiltonian=hamiltonian,
                    ansatz=ansatz,
                    n_excited_states=1,
                    beta=10.,
                    optimizer=optimizer,
                    backend=backend,
                    overlap_method="amplitude",
                  )

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

mean = []
std = []
for i in range(1,max_depth):
    mean.append(solution_dict[i]['mean'])
    std.append(solution_dict[i]['std'])

solution_dict

ed_eig= (eigenvalues[1] - eigenvalues[0])/2

import seaborn as sns
sns.set()

from matplotlib.ticker import MaxNLocator

x_axis = [i for i in range(1,max_depth)]
plt.errorbar(x_axis,
             np.array(mean)/2, 
             yerr=np.array(std)/2,
             fmt='ro', ecolor='green')

plt.hlines(y=ed_eig, xmin=0.5, xmax=6.5, 
           label='Expected value',color='b')
plt.title('Varying Depth SPSA', size=18)
plt.xlabel('Depth', size= 14)
plt.ylabel('Gap', size=14)
plt.xticks(x_axis)
plt.legend()
plt.show()

from matplotlib.ticker import MaxNLocator

x_axis = [i for i in range(1,max_depth)]
plt.errorbar(x_axis,
             np.array(mean)/2, 
             yerr=np.array(std)/2,
             fmt='ro', ecolor='green')

plt.hlines(y=ed_eig, xmin=0.5, xmax=6.5, 
           label='Expected value',color='b')
plt.title('Varying Depth', size=18)
plt.xlabel('Depth', size= 14)
plt.ylabel('Gap', size=14)
plt.xticks(x_axis)
plt.legend()
plt.show()

max_depth = 7

spsa_dict = {1: {'mean': 7.63536, 'std': 2.024100436737268},
 2: {'mean': 6.421452, 'std': 3.009174019443874},
 3: {'mean': 5.842042, 'std': 3.040774217602484},
 4: {'mean': 5.632112000000001, 'std': 3.7485136763063838},
 5: {'mean': 5.302354, 'std': 4.1340626645086065},
 6: {'mean': 4.079198, 'std': 3.6543650323135477}}

spsa_dict = {1: {'mean': 5.335038, 'std': 1.708178187062462},
 2: {'mean': 4.527469999999999, 'std': 1.756017706431231},
 3: {'mean': 3.160082, 'std': 1.587476022331046},
 4: {'mean': 2.16104, 'std': 1.5474551641970113},
 5: {'mean': 1.3676439999999996, 'std': 1.5716504103852103},
 6: {'mean': 0.19817400000000002, 'std': 1.071217238623427}}

cobyla_dict = {1: {'mean': 3.1321660000000002, 'std': 0.24530721278429674},
 2: {'mean': 2.94829, 'std': 1.0438090011587366},
 3: {'mean': 2.912668, 'std': 0.9325122996379187},
 4: {'mean': 3.16006, 'std': 0.38378217728289615},
 5: {'mean': 3.048998, 'std': 0.8419937475991135},
 6: {'mean': 2.630514, 'std': 1.029439954734612}}

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

for i in range(1,max_depth):
    mean_spsa.append(spsa_dict[i]['mean'])
    std_spsa.append(spsa_dict[i]['std'])
    mean_cobyla.append(cobyla_dict[i]['mean'])
    std_cobyla.append(cobyla_dict[i]['std'])

from matplotlib.ticker import MaxNLocator

x_axis = [i for i in range(1,max_depth)]
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.hlines(y=ed_eig, xmin=0.5, xmax=6.5, 
           label='Expected value',color='b')

plt.title('Varying Depth', size=18)
plt.xlabel('Depth', size= 14)
plt.ylabel('Gap', size=14)
plt.ylim(0., 4.1)
plt.xticks(x_axis)
plt.legend(loc="upper left")
plt.savefig('Depth.png')
plt.show()

plt.plot(x_axis, np.array(mean_spsa)/2,  'or')
plt.fill_between(x_axis, np.array(mean_spsa)/2 - np.array(std_spsa)/2 , np.array(mean_spsa)/2 + np.array(std_spsa)/2,
                 color='gray', alpha=0.2)
plt.show()

