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]
V = -2
hamiltonian = BCS_hamiltonian(EPSILONS, V)
print(hamiltonian)

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=1000)

# 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=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

backend = BasicAer.get_backend('qasm_simulator')
count = execute(states[0], backend=backend, shots=10000).result().get_counts()
plot_histogram([ideal_dict[0], count],legend=['Ideal','Result'])

count = execute(states[1], backend=backend, shots=10000).result().get_counts()
plot_histogram([ideal_dict[1], count],legend=['Ideal','Result'])

count = execute(states[2], backend=backend, shots=10000).result().get_counts()
plot_histogram([ideal_dict[2], count],legend=['Ideal','Result'])

count = execute(states[3], backend=backend, shots=10000).result().get_counts()
plot_histogram([ideal_dict[3], count],legend=['Ideal','Result'])

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

def hamiltonian_varying_eps(eps, V):
    energies_VQD = []
    energies_Classical = []
    
    for epsilons in tqdm(eps):
        hamiltonian = BCS_hamiltonian(epsilons, V)
        
        optimizer = COBYLA(maxiter=10000)
        backend = BasicAer.get_backend('qasm_simulator')

        Algo = VQD(hamiltonian=hamiltonian,
                    n_excited_states=1,
                    beta=3.,
                    optimizer=optimizer,
                    backend=backend)
        
        Algo.run(0)      
        uantum = Algo.energies
        energies_VQD.append([quantum[1], quantum[2]])
        classical = classical_solver(hamiltonian)
        energies_Classical.append([classical[1], classical[2]])
    
    return energies_VQD, energies_Classical

eps = [[1, 1.5], [1, 2.], [1, 2.5]]
V = -1.
energy_vqd, energy_classical = hamiltonian_varying_eps(eps, V)

plt.plot([str(v) for v in eps], [value[1] - value[0] for value in energy_vqd], 'ro', label="VQD")
plt.plot([str(v) for v in eps], [value[1] - value[0] for value in energy_classical], 'go', label="Exact")
plt.xlabel('$\epsilon$')
plt.ylabel('Gap')
plt.legend()
plt.show()

