import numpy as np
from scipy.sparse import diags
from scipy.sparse import identity
from scipy.sparse import kron

L = 2  # cutoff for Fock space -> can not go larger than 2 in qiskit without having problems for minimal BMN


a_b = diags(np.sqrt(np.linspace(1,L-1,L-1)),offsets=1)
a_b.todense()

i_b = identity(L)
i_b.todense()

a_f = diags(np.sqrt(np.linspace(1,1,1)),offsets=1)
a_f.todense()

sz = diags([1.0,-1.0])
sz.todense()

i_f = identity(2)
i_f.todense()

N_bos = 6  # number of boson sites -> fixed for mini-BMN 2

product_list = [i_b] * N_bos  # only the identity for bosons repeated N_bos times
a_b_list = []  # this will contain a1...a6
for i in np.arange(0,N_bos):  # loop over all bosonic operators
    operator_list = product_list.copy()  # all elements are the identity operator
    operator_list[i] = a_b  # the i^th element is now the annihilation operator for a single boson
    a_b_list.append(operator_list[0]) # start taking tensor products from first element
    for a in operator_list[1:]:
        a_b_list[i] = kron(a_b_list[i],a)  # do the outer product between each operator_list element

len(a_b_list)

N_f = 3  # number of fermionic sites -> fixed for the supersymmetric model mini-BMN with 2 matrices 

product_list = [i_f] * N_f  # only the identity for fermions repeated N_f times
a_f_list = []  # this will contain f1...f3
for i in np.arange(0,N_f):  # loop over all bosonic operators
    operator_list = product_list.copy()  # all elements are the identity operator
    operator_list[i] = a_f  # the i^th element is now the annihilation operator for a single fermion
    for j in np.arange(0,i):  # the 0:(i-1) elements are replaced by sigma_Z
        operator_list[j] = sz
    a_f_list.append(operator_list[0])  # start taking tensor products from the first operator
    for a in operator_list[1:]:
        a_f_list[i] = kron(a_f_list[i],a)  # do the outer product 

i_b_tot = identity(L**N_bos) 

i_f_tot = identity(2**N_f)

op_list = []
for a in a_b_list:
    op_list.append(kron(a,i_f_tot))
for a in a_f_list:
    op_list.append(kron(i_b_tot,a))

len(op_list)

op_list[0]

x_list = []
# only use the bosonic operators
bosons = op_list[:N_bos]
for op in bosons:
    x_list.append(1/np.sqrt(2)*(op.conjugate().transpose() + op))

# save the fermions operators for convenience in a new list
fermions = op_list[-N_f:]

# Create the simple quartic Hamiltonian.
H_k = 0

for a in op_list[:N_bos]:
    H_k = H_k + a.conjugate().transpose() * a

for a in op_list[-N_f:]:
    H_k = H_k + (3.0 / 2) * a.conjugate().transpose() * a

# vacuum energy
H_k = H_k + 0.25 * (2 * N_bos - 3 * N_f - 3)



from scipy.sparse.linalg import eigsh
eigv = eigsh(H_k,k=10,which='SA',return_eigenvectors=False,tol=0)

print(f"Free Hamiltonian eigenvalues: {eigv[::-1]}")

V_b = (
    x_list[2] * x_list[2] * x_list[3] * x_list[3]
    + x_list[2] * x_list[2] * x_list[4] * x_list[4]
    + x_list[1] * x_list[1] * x_list[3] * x_list[3]
    + x_list[1] * x_list[1] * x_list[5] * x_list[5]
    + x_list[0] * x_list[0] * x_list[4] * x_list[4]
    + x_list[0] * x_list[0] * x_list[5] * x_list[5]
    - 2 * x_list[0] * x_list[2] * x_list[3] * x_list[5]
    - 2 * x_list[0] * x_list[1] * x_list[3] * x_list[4]
    - 2 * x_list[1] * x_list[2] * x_list[4] * x_list[5]
)

V_bf = (2j / np.sqrt(2)) * (
    (x_list[0] - 1j * x_list[3])
    * fermions[1].conjugate().transpose()
    * fermions[2].conjugate().transpose()
    + (x_list[1] - 1j * x_list[4])
    * fermions[2].conjugate().transpose()
    * fermions[0].conjugate().transpose()
    + (x_list[2] - 1j * x_list[5])
    * fermions[0].conjugate().transpose()
    * fermions[1].conjugate().transpose()
    - (x_list[0] + 1j * x_list[3]) * fermions[2] * fermions[1]
    - (x_list[1] + 1j * x_list[4]) * fermions[0] * fermions[2]
    - (x_list[2] + 1j * x_list[5]) * fermions[1] * fermions[0]
)

# t'Hooft coupling
g2N = 0.2
# Number of colors
N = 2
# full hamiltonian
H = H_k + g2N/N*V_b + np.sqrt(g2N/N)*V_bf

H.shape

eigv = eigsh(H,k=10,which='SM',return_eigenvectors=False,tol=0)

print(f"Full Hamiltonian with lambda={g2N} : {eigv[::-1]}")

import warnings
import time
from qiskit.opflow import MatrixOp

start_time = time.time()
warnings.filterwarnings("ignore")

qubitOp = MatrixOp(primitive=H)

print("Size of Hamiltonian:", np.shape(H))
end_time = time.time()

runtime = end_time-start_time
print('Program runtime: ',runtime)

from qiskit.circuit.library import EfficientSU2

var_form = EfficientSU2(qubitOp.num_qubits, su2_gates=['ry','rz'], entanglement="full", reps=1)

var_form.draw(output='mpl')

from qiskit.algorithms import NumPyEigensolver

mes = NumPyEigensolver(k=4)  # k is the number of eigenvalues to compute

result = mes.compute_eigenvalues(qubitOp)
print(result.eigenvalues)
print("Exact Result of qubit hamiltonian:", np.real(result.eigenvalues))
print("Exact Result of discrete hamiltonian:", np.real(eigv[::-1][0:4]))

from qiskit import Aer  # change this to Aer for C++ compiled code and the option to run on the GPU
from qiskit.algorithms import VQE

# start a quantum instance
# fix the random seed of the simulator to make values reproducible
from qiskit.utils import algorithm_globals, QuantumInstance
seed = 50
algorithm_globals.random_seed = seed


# with GPU
#backend = Aer.get_backend("statevector_simulator", method="statevector_gpu")
# with parallel computation restricted to 2 OpenMP threads (0: maximum)
backend = Aer.get_backend("statevector_simulator", max_parallel_threads=6, max_parallel_experiments=0)

q_instance = QuantumInstance(backend, seed_transpiler=seed, seed_simulator=seed)


from qiskit.algorithms.optimizers import SLSQP, COBYLA, L_BFGS_B, NELDER_MEAD, SPSA

start_time = time.time()
# choose one classical optimizer
# optimizer = L_BFGS_B(maxiter=5000)
# optimizer = SPSA(maxiter=5000)
optimizer = COBYLA(maxiter=5000)

counts = []
values = []
# callback functions to store the counts from each iteration of the VQE
def store_intermediate_result(eval_count, parameters, mean, std):
    counts.append(eval_count)
    values.append(mean)

# initital points
#random_init = np.random.random(var_form.num_parameters)
#fixed_init = result.optimal_point

# Setup the VQE algorithm
#vqe = VQE(ansatz=var_form, optimizer=optimizer, initial_point=random_init, quantum_instance=q_instance, callback=store_intermediate_result)
#vqe = VQE(ansatz=var_form, optimizer=optimizer, initial_point=fixed_init, quantum_instance=q_instance, callback=store_intermediate_result)
vqe = VQE(ansatz=var_form, optimizer=optimizer, quantum_instance=q_instance, callback=store_intermediate_result)

# run the VQE with out Hamiltonian operator
# since the VQE is a ground state solver, the syntax is the same as before ;-)
result = vqe.compute_minimum_eigenvalue(qubitOp)
vqe_result = np.real(result.eigenvalue)
print("VQE Result:", vqe_result)

end_time = time.time()
runtime = end_time-start_time
print('Program runtime:',runtime, "s")

# Plot convergence plot
from matplotlib import pyplot as plt
plt.figure(figsize=(12,10))
plt.plot(counts, values, label="SLSQP")
plt.xlabel('Eval count')
plt.ylabel('Energy')
plt.title('Energy convergence')
plt.legend(loc='upper right')
plt.show();

# start a quantum instance
# fix the random seed of the simulator to make values reproducible
from qiskit.utils import algorithm_globals, QuantumInstance
seed = 50
algorithm_globals.random_seed = seed

# with GPU
# backend = Aer.get_backend('qasm_simulator', method="statevector_gpu")
# with parallel computation restricted to 1 OpenMP threads (0: maximum)
backend = Aer.get_backend('qasm_simulator', max_parallel_threads=1, max_parallel_experiments=0)

q_instance = QuantumInstance(backend, seed_transpiler=seed, seed_simulator=seed, shots=1024)

from qiskit.aqua.components.optimizers import SPSA
from qiskit.aqua import QuantumInstance

warnings.filterwarnings("ignore")

optimizer = SPSA(maxiter=6000)

counts = []
values = []
# callback functions to store the counts from each iteration of the VQE
def store_intermediate_result(eval_count, parameters, mean, std):
    counts.append(eval_count)
    values.append(mean)

# initital points
#random_init = np.random.random(var_form.num_parameters)
#fixed_init = result.optimal_point

# Setup the VQE algorithm
#vqe = VQE(ansatz=var_form, optimizer=optimizer, initial_point=random_init, quantum_instance=q_instance, callback=store_intermediate_result)
#vqe = VQE(ansatz=var_form, optimizer=optimizer, initial_point=fixed_init, quantum_instance=q_instance, callback=store_intermediate_result)
vqe = VQE(ansatz=var_form, optimizer=optimizer, quantum_instance=q_instance, callback=store_intermediate_result)

# run the VQE with out Hamiltonian operator
# since the VQE is a ground state solver, the syntax is the same as before ;-)
result = vqe.compute_minimum_eigenvalue(qubitOp)
vqe_result = np.real(result.eigenvalue)
print("VQE Result:", vqe_result)

# Plot convergence plot
plt.figure(figsize=(12,10))
plt.plot(counts, values, label="SPSA")
plt.xlabel('Eval count')
plt.ylabel('Energy')
plt.title('Energy convergence')
plt.legend(loc='upper right')

import qiskit.tools.jupyter
%qiskit_version_table