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.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 qiskit.providers.aer.noise import depolarizing_error
from qiskit.providers.aer.noise import thermal_relaxation_error
from qiskit.providers.aer.noise import NoiseModel

# Defining noise model

# T1 and T2 values for qubits 0-n
n_qubits = 2*len(EPSILONS)

# T1s = np.random.normal(50e3, 10e3, n_qubits) # Sampled from normal distribution mean 50 microsec
# T2s = np.random.normal(70e3, 10e3, n_qubits)  # Sampled from normal distribution mean 50 microsec

T1s = [30e3]*n_qubits
T2s = [20e3]*n_qubits

# Truncate random T2s <= T1s
T2s = np.array([min(T2s[j], 2 * T1s[j]) for j in range(n_qubits)])

# Instruction times (in nanoseconds)
time_u1 = 0   # virtual gate
time_u2 = 50  # (single X90 pulse)
time_u3 = 100 # (two X90 pulses)
time_cx = 300
time_reset = 1000  # 1 microsecond
time_measure = 1000 # 1 microsecond

# QuantumError objects
errors_reset = [thermal_relaxation_error(t1, t2, time_reset)
                for t1, t2 in zip(T1s, T2s)]
errors_measure = [thermal_relaxation_error(t1, t2, time_measure)
                  for t1, t2 in zip(T1s, T2s)]
errors_u1  = [thermal_relaxation_error(t1, t2, time_u1)
              for t1, t2 in zip(T1s, T2s)]
errors_u2  = [thermal_relaxation_error(t1, t2, time_u2)
              for t1, t2 in zip(T1s, T2s)]
errors_u3  = [thermal_relaxation_error(t1, t2, time_u3)
              for t1, t2 in zip(T1s, T2s)]
errors_cx = [[thermal_relaxation_error(t1a, t2a, time_cx).expand(
             thermal_relaxation_error(t1b, t2b, time_cx))
              for t1a, t2a in zip(T1s, T2s)]
               for t1b, t2b in zip(T1s, T2s)]

# Add errors to noise model
noise_thermal = NoiseModel()

for j in range(n_qubits):
    noise_thermal.add_quantum_error(errors_reset[j], "reset", [j])
    noise_thermal.add_quantum_error(errors_measure[j], "measure", [j])
    noise_thermal.add_quantum_error(errors_u1[j], "u1", [j])
    noise_thermal.add_quantum_error(errors_u2[j], "u2", [j])
    noise_thermal.add_quantum_error(errors_u3[j], "u3", [j])
    for k in range(n_qubits):
        noise_thermal.add_quantum_error(errors_cx[j][k], "cx", [j, k])

print(noise_thermal)

from tqdm import tqdm
from qiskit.aqua import QuantumInstance

# Parameters Variables
n_trials = 50
max_depth = 7

# Auxiliary Variables
solution_dict = {}

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

# Define Backend
backend = QuantumInstance(backend= Aer.get_backend('qasm_simulator'),
                          noise_model=noise_thermal,
                          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=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[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

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=2, 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()

spsa_dict = {1: {'mean': 4.540763883482028, 'std': 0.2166500619449026},
 2: {'mean': 3.6292261939489365, 'std': 1.3793621699692478},
 3: {'mean': 3.5182758038620277, 'std': 0.8910628276050163},
 4: {'mean': 3.292976405037763, 'std': 1.3966224263108362},
 5: {'mean': 3.5144598497611055, 'std': 1.0963827019837287},
 6: {'mean': 2.6975338496665677, 'std': 1.3507496759669415}}


cobyla_dict = {1: {'mean': 2.3633300482974127, 'std': 2.1367566366817043},
 2: {'mean': 4.113586221793872, 'std': 0.7316145547992721},
 3: {'mean': 4.321305470168272, 'std': 0.4919574179280678},
 4: {'mean': 4.3014277420825975, 'std': 0.4023952320110141},
 5: {'mean': 4.157595644181709, 'std': 0.33724862839247316},
 6: {'mean': 3.968890438601869, 'std': 0.29850260762472425}}

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=2, 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.5)
plt.xticks(x_axis)
plt.legend()
plt.savefig('Depth.png')
plt.show()

