import sys
import numpy as np
N = 2  # cutoff for Fock space. The number of qubits used will be K = log2(N) 
annOp = np.array(np.diagflat(np.sqrt(np.linspace(1,N-1,N-1)),k=1))
with np.printoptions(precision=3, suppress=True, linewidth=120, threshold=sys.maxsize): # print array lines up to character 120 and floats using 3 digits
  print(annOp)

iden = np.identity(N)
with np.printoptions(precision=3, suppress=True, linewidth=120, threshold=sys.maxsize): # print array lines up to character 120 and floats using 3 digits
  print(iden)

Nmat = 6  # number of bosonic dof as matrices: SU(2) -> 2^2-1=6 generators
bosonList = [annOp]
for bosons in range(0,Nmat-1):
  bosonList.append(iden)
with np.printoptions(precision=3, suppress=True, linewidth=120, threshold=sys.maxsize):
  for i in bosonList:
    print(f"{i}\n")

# This for loop takes the appropriate Kronecker products for each boson.
for i in range(0,Nmat):
  for j in range(0,Nmat-1):
    # For the nth boson, the nth Kronecker product is with the annihilation operator.
    if j == i-1 and i != 0:
        bosonList[i] = np.kron(bosonList[i], annOp)
    # Else, the nth Kronecker product is with the identity matrix.
    else:
        bosonList[i] = np.kron(bosonList[i], iden)
# the following will crash the system
# with np.printoptions(precision=3, suppress=True, linewidth=120, threshold=sys.maxsize):
#   for i in bosonList:
#     print(f"{i}\n")

[x.shape for x in bosonList]

import matplotlib.pyplot as plt

fig, ax = plt.subplots(ncols=6, figsize=(18, 6))
fig.subplots_adjust(left=0.02, bottom=0.06, right=0.95, top=0.94, wspace=0.5, hspace=1)
for i,axi in enumerate(np.ravel(ax)):
  pos = axi.imshow(bosonList[i], cmap='jet')
  #fig.colorbar(pos, ax=axi)
  axi.set_title(f"Boson {i+1}")
fig.suptitle('Operators', fontsize=16)
fig.subplots_adjust(top=0.9)

def bosonHamiltonians(numBosons, bosonMatrixDim, coupling=0.1):
    # This function takes the number of bosons (numBosons) and the size of the bosons (bosonMatrixDim)
    # as arguments. If the size of the boson is 2 x 2, for example, then bosonMatrixDim = 2.
    # coupling is lambda/2 and lambda is the 'tHooft coupling
    
    # Create the n x n annhiliation operator. Here, we create a list with all the integers from
    # the square root of 1 to the square root of n-1, where n x n is the size of the boson.
    # This list is injected as the upper diagonal in an array of zeros.
    annOP = np.array(np.diagflat(np.sqrt(np.linspace(1,bosonMatrixDim-1,bosonMatrixDim-1)),k=1))
    
    # Create the n x n identity matrix.
    iden = np.identity(bosonMatrixDim)
    
    # Create a list which holds the six bosons. bosonList[0] is the first boson. bosonList[5] is the 6th boson.
    # For the first boson, the Kronecker product starts with the annihilation operator.
    bosonList = [annOP]
    
    # This for loop creates the list of bosons before the Kronecker products are taken. The first one was 
    # already created in the previous line.
    for bosons in range(0,numBosons-1):
        bosonList.append(iden)
    
    # This for loop takes the appropriate Kronecker products for each boson.
    for i in range(0,numBosons):
        for j in range(0,numBosons-1):
            
            # For the nth boson, the nth Kronecker product is with the annihilation operator.
            if j == i-1 and i != 0:
                bosonList[i] = np.kron(bosonList[i], annOP)
            # Else, the nth Kronecker product is with the identity matrix.
            else:
                bosonList[i] = np.kron(bosonList[i], iden)

    # Create the position operators. Normalized as in the notes
    x = []
    for r in range(0, numBosons):
        x.append((1/np.sqrt(2))*(bosonList[r] + np.transpose(np.conjugate(bosonList[r]))))
    
    # Create the simple quadratic Hamiltonian.
    H2MM = 0

    for i in range(0,numBosons):
        # The @ symbol is a shorthand for matrix multiplication. It's equivalent to using np.matmul().
        H2MM = H2MM + (np.transpose(np.conjugate(bosonList[i])) @ bosonList[i])

    H2MM = H2MM + 0.5*numBosons*np.identity(bosonMatrixDim**(numBosons))

    # Create the full quartic SU(2) Hamiltonian.
    x_sq = []

    for i in x:
        x_sq.append(i @ i)

    H4MM1 = (H2MM + coupling*((x_sq[2] @ x_sq[3]) 
                              + (x_sq[2] @ x_sq[4]) 
                              + (x_sq[1] @ x_sq[3]) 
                              + (x_sq[1] @ x_sq[5])
                              + (x_sq[0] @ x_sq[4])
                              + (x_sq[0] @ x_sq[5])
                              - 2*((x[0] @ x[2]) @ (x[3] @ x[5]))
                              - 2*((x[0] @ x[1]) @ (x[3] @ x[4]))
                              - 2*((x[1] @ x[2]) @ (x[4] @ x[5]))))
    return H2MM, H4MM1

# Set the number of bosons here. For the proper full quartic SU(2) Hamiltonian, this has to be 6.
numBosons = 6

# Set the size of the bosons here. For an n x n boson, bosonMatrixDim = n.
# 5 will crash on laptop -> out of memory > 10GB
# it also crashes this colab with 12Gb of RAM ->  tries 6 of these tcmalloc: large alloc 1953128448 bytes 
# THIS IS ONLY BECAUSE WE ARE USING NUMPY IN THE FUNCTION! THIS CAN BE AVOIDED USING QISKIT OR QUTIP DIRECTLY TO CREATE SPARSE REPRESENTATIONS
bosonMatrixDim = 4

# the function actually creates intermediate matrices x and x2 so the memory footprint is even higher...for N=5 it requires ~50GB of RAM
H2MM, H4MM1 = bosonHamiltonians(numBosons, bosonMatrixDim, coupling=0.1)  # we can change the coupling by adding coupling=0.5 to the args

# Set the Hamiltonian that you want to run the VQE algorithm with here.
hamiltonian = H4MM1

# Use numpy to find the eigenvalues of the simple quartic Hamiltonian
# this can take a long time for large matrices
val,vec=np.linalg.eig(hamiltonian)
z = np.argsort(val)
z = z[0:len(hamiltonian)]
energies=(val[z])

if hamiltonian.all() == H4MM1.all():
    # Sometimes there are negligible imaginary parts (e.g. order of 10^-16), so I take the real parts only.
    print('Full Quartic SU2 Eigenvalues:\n\n', np.real(energies))
else:
    print('Simple Quartic Eigenvalues:\n\n', energies)
    

from qiskit.algorithms import VQE, NumPyEigensolver
from qiskit import Aer

# Convert the Hamiltonian Matrix to a qubit operator.
import warnings
import time

from qiskit.opflow import MatrixOp

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

qubitOp = MatrixOp(primitive=hamiltonian)

print("Size of Hamiltonian:", np.shape(hamiltonian))
print(qubitOp)

end_time = time.time()

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

print("Total number of qubits: ",qubitOp.num_qubits)

print(type(qubitOp).__name__)

# Create the variational form.
from qiskit.circuit.library import EfficientSU2

var_form = EfficientSU2(qubitOp.num_qubits, su2_gates=['ry'], entanglement="full", reps=1)
var_form.draw(output='mpl',fold=100)
# var_form.draw(output='mpl',fold=100, style={"fontsize":10, "displaycolor": {"ry": ('#9DC3E6','#000000')}},filename="../figures/var_form_bmn2_Ry_depth1.pdf")

np.log2(hamiltonian.shape[0])

# This is a callback function that will allow us to store information about the optimizer while the VQE is running.
# This is called when instantiating an object of the VQE class.

def store_intermediate_result(eval_count, parameters, mean, std):
    counts.append(eval_count)
    values.append(mean)

from qiskit.algorithms.optimizers import SLSQP
from qiskit.utils import algorithm_globals, QuantumInstance

warnings.filterwarnings("ignore")

rngseed = 0
algorithm_globals.random_seed = rngseed
backend = Aer.get_backend(
    "statevector_simulator", max_parallel_threads=6, max_parallel_experiments=0
)
q_instance = QuantumInstance(
    backend, seed_transpiler=rngseed, seed_simulator=rngseed
)


counts = []
values = []

optim = SLSQP(maxiter=1000)

# Setup the VQE algorithm
vqe = VQE(
    ansatz=var_form,
    optimizer=optim,
    quantum_instance=q_instance,
    callback=store_intermediate_result,
)
# run the VQE with out Hamiltonian operator
result = vqe.compute_minimum_eigenvalue(qubitOp)
vqe_result = np.real(result.eigenvalue)
print(f"VQE gs energy: {vqe_result}")

# Plot all of the optimizers in a single convergence plot. If you see that the optimizer isn't converging,
# increase the maximum iterations or maximum function evaluations in the optimizer list above.

labelList = ["SLSQP"]

# Limit the range of the plot in order to make the convergence of each optimizer more visible.
plt.plot(counts, values)

plt.axhline(y=np.real(energies[0]),xmin=0,xmax=counts[-1],color='k',linestyle='--',label="Exact Energy")
plt.xlabel('Eval count')
plt.ylabel('Energy')
plt.title('Energy convergence for various optimizers')
plt.legend(loc='upper right', labels=labelList)

from qiskit.algorithms.optimizers import SPSA
from qiskit.utils import algorithm_globals, QuantumInstance

warnings.filterwarnings("ignore")

rngseed = 0
algorithm_globals.random_seed = rngseed
backend = Aer.get_backend("qasm_simulator")
q_instance = QuantumInstance(
    backend, shots=1024, seed_transpiler=rngseed, seed_simulator=rngseed
)

counts = []
values = []

optim = SPSA(maxiter=1000)

# Setup the VQE algorithm
vqe = VQE(
    ansatz=var_form,
    optimizer=optim,
    quantum_instance=q_instance,
    callback=store_intermediate_result,
)
# run the VQE with out Hamiltonian operator
result = vqe.compute_minimum_eigenvalue(qubitOp)
vqe_result = np.real(result.eigenvalue)
print(f"VQE gs energy: {vqe_result}")

# Plot convergence plot for the SPSA optimizer
plt.figure(figsize=(15,15))
plt.plot(counts, values, label="SPSA")
plt.plot(counts,[np.real(energies[0])]*len(counts), 'k--',label="Exact Energy")
plt.xlabel('Eval count')
plt.ylabel('Energy')
plt.title('Energy convergence of SPSA optimizer (Qasm Simulation)')
plt.legend(loc='upper right')

from scipy.sparse import diags
from scipy.sparse import identity
from scipy.sparse import kron
from scipy.sparse i

def build_operators(L: int, N_bos: int) -> list:
    """Generate all the annihilation operators needed to build the hamiltonian

    Args:
        L (int): the cutoff of the single site Fock space
        N_bos (int): the number of bosonic sites

    Returns:
        list: a list of annihilation operators, length=N_bos
    """
    # The annihilation operator for the single boson
    a_b = diags(np.sqrt(np.linspace(1, L - 1, L - 1)), offsets=1)
    # The identity operator of the Fock space of a single boson
    i_b = identity(L)
    # Bosonic Hilbert space
    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

    return a_b_list


def build_gauge_casimir(L: int, N_bos: int) -> list:
    """Generate the gauge generators operators

    Args:
        L (int): the single site cutoff of the Fock space
        N_bos (int): the number of bosonic sites

    Returns:
        list : 3 generators (for SU(2))
    """
    # generate the annihilation operators
    bosons = build_operators(L, N_bos)
    # define the generator list for SU(2)
    g_list = [0] * 3
    g_list[0] = 1j * (
        bosons[1].conjugate().transpose() * bosons[2]
        - bosons[2].conjugate().transpose() * bosons[1]
        + bosons[4].conjugate().transpose() * bosons[5]
        - bosons[5].conjugate().transpose() * bosons[4]
    )
    g_list[1] = 1j * (
        bosons[2].conjugate().transpose() * bosons[0]
        - bosons[0].conjugate().transpose() * bosons[2]
        + bosons[5].conjugate().transpose() * bosons[3]
        - bosons[3].conjugate().transpose() * bosons[5]
    )
    g_list[2] = 1j * (
        bosons[0].conjugate().transpose() * bosons[1]
        - bosons[1].conjugate().transpose() * bosons[0]
        + bosons[3].conjugate().transpose() * bosons[4]
        - bosons[4].conjugate().transpose() * bosons[3]
    )

    return g_list[0] * g_list[0] + g_list[1] * g_list[1] + g_list[2] * g_list[2]


g2 = build_gauge_casimir(bosonMatrixDim,numBosons)

# result is the output of the VQE
bra = result.eigenstate.conjugate().transpose()
ket = result.eigenstate
op = g2
expect = bra.dot(op.dot(ket))
print(f"<0| g2 |0>: {expect.real}")

counts = []
values = []

optim = SLSQP(maxiter=1000)

# Setup the VQE algorithm
vqe = VQE(
    ansatz=var_form,
    optimizer=optim,
    quantum_instance=q_instance,
    callback=store_intermediate_result,
)
# run the VQE with our Hamiltonian operator
result = vqe.compute_minimum_eigenvalue(qubitOp, aux_operators=[MatrixOp(g2)])
vqe_result = np.real(result.eigenvalue)
print(f"VQE gs energy: {vqe_result}")

print(result)

print("<0| g2 |0> from VQE: ",result.aux_operator_eigenvalues[0,0].real)
bra = result.eigenstate.conjugate().transpose()
ket = result.eigenstate
op = g2
expect = bra.dot(op.dot(ket))
print(f"<0| g2 |0> from matrix-vector: {expect.real}")

from scipy.sparse.linalg import eigsh

eigv, eigk = eigsh(H4MM1, 1, which="SA", return_eigenvectors=True, tol=0)

# eigk[:,0] is the ground state vector
bra = eigk[:,0].conjugate().transpose()
ket = eigk[:,0]
op = H4MM1
expect = bra.dot(op.dot(ket))
print(f"<0| H |0> from matrix-vector: {expect.real}")
print(f"eigenvalue of H: {eigv[0]}")

# eigk[:,0] is the ground state vector
bra = eigk[:,0].conjugate().transpose()
ket = eigk[:,0]
op = g2
expect = bra.dot(op.dot(ket))
print(f"<0| G2 |0> from matrix-vector: {expect.real}")

from numpy.linalg import eigh 

eigv, eigk = eigh(H4MM1)
print(eigk.shape)
print(eigv[0])

_, eigk = eigh(H4MM1)
for i in np.arange(3):
    bra = eigk[:,i].conjugate().transpose()
    ket = eigk[:,i]
    expect = bra @ g2 @ ket
    print(f"<{i}| G2 |{i}> from matrix-vector: {expect.real}")