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

# Python imports
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
sns.set()

# 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., 3., 4., 3.]
V = -2
hamiltonian = BCS_hamiltonian(EPSILONS, V)
print(f"Hamiltonian: {hamiltonian}\n")
eigenvalues, eigenvectors = classical_solver(hamiltonian)
print(f"Eigenvalues: {eigenvalues}")

from itertools import product
# Get the ideal count
ideal_dict = {}

bin_combinations = list(product(['0','1'], repeat=hamiltonian.num_qubits))
for i, eigvect in enumerate(np.abs(eigenvectors)**2):
    dic = {}
    for ind, val in enumerate(bin_combinations):
        val = ''.join(val)
        dic[val] = eigvect[ind]
    ideal_dict[i] = dic

from qiskit.aqua import QuantumInstance

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

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

# Define ansatz
ansatz = TwoLocal(hamiltonian.num_qubits, ['ry','rz'], 'cx', reps=3)

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

Algo.run()      
vqd_energies = Algo.energies
vqd_states = Algo.states

from copy import copy
# Create states and measure them
states = []
for ind, state in enumerate(vqd_states):
    states.append(copy(state))
    states[ind].measure_all()

from qiskit import execute
from qiskit.visualization import plot_histogram



fig, axs = plt.subplots(2, 2, figsize=(30,15))

for i, ax in zip(range(len(states)), axs.flat):
    count = execute(states[i], backend=backend, shots=10000).result().get_counts()
    plot_histogram([ideal_dict[i], count],legend=['Ideal','Result'], ax=ax)

print(f'Ideal eigenvalues: {eigenvalues}')
print(f'Obtained eigenvalues: {vqd_energies}')

n_1_sim, n_2_sim = 1, 2
n_1_vqd, n_2_vqd = 1, 2
print(f"Gap Exact: {(eigenvalues[n_2_sim] - eigenvalues[n_1_sim])/2}")
print(f"Gap VQD: {(vqd_energies[n_2_vqd] - vqd_energies[n_1_vqd])/2}")

from qiskit.providers.aer.noise import depolarizing_error
from qiskit.providers.aer.noise import thermal_relaxation_error
from qiskit.providers.aer.noise import NoiseModel

from qiskit.providers.aer.noise import depolarizing_error
from qiskit.providers.aer.noise import thermal_relaxation_error
from qiskit.providers.aer.noise import NoiseModel

# T1 and T2 values for qubits 0-n
n_qubits = 4
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 qiskit.aqua import QuantumInstance
from qiskit import Aer

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

# Define Backend
backend = QuantumInstance(backend=Aer.get_backend('qasm_simulator'),
                          noise_model=noise_thermal,
                          shots=10000)

# Define ansatz
ansatz = TwoLocal(hamiltonian.num_qubits, ['ry','rz'], 'cx', reps=3)

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

Algo.run()      
vqd_energies = Algo.energies
vqd_states = Algo.states

from copy import copy
# Create states and measure them
states = []
for ind, state in enumerate(vqd_states):
    states.append(copy(state))
    states[ind].measure_all()

from qiskit import execute
from qiskit.visualization import plot_histogram

import seaborn as sns
sns.set()

fig, axs = plt.subplots(2, 2, figsize=(30,15))

backend = QuantumInstance(backend=Aer.get_backend('qasm_simulator'),
                          noise_model=noise_thermal,
                          shots=10000)

for i, ax in zip(range(len(states)), axs.flat):
    count = backend.execute(states[i]).get_counts()
    plot_histogram([ideal_dict[i], count],legend=['Ideal','Result'], ax=ax)


print(f'Ideal eigenvalues: {eigenvalues}')
print(f'Obtained eigenvalues: {vqd_energies}')

n_1_sim, n_2_sim = 1, 2
n_1_vqd, n_2_vqd = 1, 2
print(f"Gap Exact: {(eigenvalues[n_2_sim] - eigenvalues[n_1_sim])/2}")
print(f"Gap VQD: {(vqd_energies[n_2_vqd] - vqd_energies[n_1_vqd])/2}")

from tqdm import tqdm
from qiskit.aqua import QuantumInstance

es_1 = []
es_2 = []

n = 50

for _ in tqdm(range(n)):
    # Define Optimizer
    optimizer = COBYLA()

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

    # Define ansatz
    ansatz = TwoLocal(hamiltonian.num_qubits, ['ry','rz'], 'cx', reps=2)

    # 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])

np.savetxt("groundstate.txt", es_1)
np.savetxt("excitedstate.txt", es_2)

print(f"Ground State: {np.round(np.mean(es_1),3)} +/- {np.round(np.std(es_1),3)} | Expected: {eigenvalues[1]}")
print(f"Excited State: {np.round(np.mean(es_2),3)} +/- {np.round(np.std(es_2),3)} | Expected: {eigenvalues[2]}")

np.round(np.mean(es_2) - np.mean(es_1), 3)/2

import seaborn as sns

sns.boxplot(x=es_1)
ax=sns.swarmplot(x=es_1, color="0.25")
ax.set_title("1st Excited State")
ax.vlines(x=eigenvalues[1], ymin=-10, ymax=10, label='Expected Eigenvalue',color='r')
plt.legend()
plt.show()

sns.boxplot(x=es_2)
ax=sns.swarmplot(x=es_2, color="0.25")
ax.set_title("2nd Excited State")
ax.vlines(x=eigenvalues[2], ymin=-10, ymax=10, label='Expected Eigenvalue',color='r')
plt.legend()
plt.show()

sns.boxplot(x=(np.array(es_2)- np.array(es_1))/2)
ax=sns.swarmplot(x=(np.array(es_2)- np.array(es_1))/2, color="0.25")
ax.set_title("Gap")
ax.vlines(x=(eigenvalues[2] - eigenvalues[1])/2, ymin=-10, ymax=10, label='Expected value',color='r')
plt.legend()
plt.show()

v_dependence =[]
v_values = np.linspace(0, 5, 20)
for v in v_values:
    EPSILONS = [3., 3., 3., 4., 3.]
    hamiltonian = BCS_hamiltonian(EPSILONS, -float(v))
    eigenvalues, eigenvectors = classical_solver(hamiltonian)
    v_dependence.append(eigenvalues[1] - eigenvalues[0])

plt.plot(v_values, v_dependence, 'o');

