import numpy as np
from scipy.sparse import diags
L = 8  # cutoff for Fock space
a = diags(np.sqrt(np.linspace(1,L-1,L-1)),offsets=1)

print(a.toarray())

from qiskit.opflow import MatrixOp

qubitOp = MatrixOp(primitive=a)
print(qubitOp.num_qubits)

type(qubitOp)

print(qubitOp)

a_pauli = qubitOp.to_pauli_op()
print(a_pauli)

type(a_pauli)

a_pauli.to_matrix()

from scipy.sparse import identity
iden = identity(L)

qubitOp = MatrixOp(primitive=iden)
print(qubitOp.num_qubits)

i_pauli = qubitOp.to_pauli_op()
print(i_pauli)

i_pauli.to_matrix()

Nmat = 6  # BMN2 for SU(2) has 6 bosonic matrices

from qiskit.opflow import ListOp, TensoredOp

boson = ListOp([i_pauli]*6)
print(boson)

print(boson[0])

type(boson[0])

# generically speaking, we construct the list of bosons and then take the outer product
a_list = []  # this will contain a1...a6 as a list of ListOp
for i in np.arange(0,Nmat):  # loop over all operators
    operator_list = [i_pauli] * Nmat  # only the identity repeated Nmat times
    operator_list[i] = a_pauli  # the i^th element is now the annihilation operator for a single boson
    a_list.append(ListOp(operator_list))


print(a_list[0])

type(a_list[0])

# here we create a list of operators which are the tensor products...
a_tensor = [TensoredOp(x) for x in a_list]


print(a_tensor[0])

print(a_tensor[0].to_pauli_op())

a_tensor[0].num_qubits

sparse = a_tensor[0].to_spmatrix()
type(sparse)

sparse.shape

# from qiskit.utils import algorithm_globals
# algorithm_globals.massive=True  # needed to big matrices...
# full = a_tensor[0].to_matrix()
# len(full)

# full.shape

i_tensor = TensoredOp(ListOp([i_pauli] * Nmat))

type(i_tensor)

print(i_tensor)

i_tensor.to_pauli_op()

from qiskit.opflow import I

I^Nmat

i_tensor.num_qubits

# full = i_tensor.to_matrix()
# len(full)

print(i_tensor.to_spmatrix())

# for each boson they are constructed using a and adag
x_tensor = [1/np.sqrt(2)*(~x + x) for x in a_tensor]

print(x_tensor[0])

print(x_tensor[0].to_pauli_op())

print(x_tensor[0].to_spmatrix())

from qiskit.opflow import SummedOp

H_zero = 0.5*Nmat*i_tensor

print(H_zero)

H_zero_pauli = 0.5*Nmat*(I^Nmat)

print(H_zero.to_pauli_op())

print(H_zero_pauli)

H_zero.to_spmatrix()

H_zero.to_pauli_op()

### Harmonic oscillator
# this should be summed over all the bosons (Nmat)
H_list = [H_zero]
for op in a_tensor:
    H_list.append((~op @ op))
H_osc = SummedOp(H_list)


print(H_osc)

type(H_osc)

H_osc_pauli = H_osc.to_pauli_op()

type(H_osc_pauli)

H_osc_pauli.num_qubits

H_osc.to_spmatrix()

quartic1 = SummedOp([x_tensor[2]@x_tensor[2]@x_tensor[3]@x_tensor[3], \
                    x_tensor[2]@x_tensor[2]@x_tensor[4]@x_tensor[4], \
                    x_tensor[1]@x_tensor[1]@x_tensor[3]@x_tensor[3], \
                    x_tensor[1]@x_tensor[1]@x_tensor[5]@x_tensor[5], \
                    x_tensor[0]@x_tensor[0]@x_tensor[4]@x_tensor[4], \
                    x_tensor[0]@x_tensor[0]@x_tensor[5]@x_tensor[5]])
                    

print(quartic1.to_pauli_op())

quartic2 = SummedOp([x_tensor[0]@x_tensor[2]@x_tensor[3]@x_tensor[5],\
                    x_tensor[0]@x_tensor[1]@x_tensor[3]@x_tensor[4],\
                    x_tensor[1]@x_tensor[2]@x_tensor[4]@x_tensor[5]], coeff = -2.0)

print(quartic2.to_pauli_op())

type(quartic2)

### Quartic Interaction
V =  quartic1 + quartic2

V.num_qubits

V_pauli = V.to_pauli_op()

g2N = 0.2
N = 2

### Full Hamiltonian
H = H_osc + g2N / N * V

type(H)

H.to_pauli_op()

type(H.to_pauli_op())

H.num_qubits

testH = H.to_matrix_op()
testH

Hpauli = H.to_pauli_op()

# the line below is also fine...
# Hpauli1 = H_osc.to_pauli_op() + g2N / N * V.to_pauli_op()
# Hpauli

print(Hpauli)

type(Hpauli)

print(f"Number of terms in H: {len(Hpauli.oplist)}")

# write expression to disk
with open('pauliH_L8_bos2.txt','w') as f:
    print(Hpauli,file=f)

import sys 

print(sys.getsizeof(Hpauli))

from pympler import asizeof

print(f"Size of Hpauli in MB: {asizeof.asizeof(Hpauli)/1024/1024}")
# print(asizeof.asized(Hpauli, detail=1).format())

from qiskit.algorithms import NumPyMinimumEigensolver

npme = NumPyMinimumEigensolver()

# result = npme.compute_minimum_eigenvalue(operator=H_osc)
# result = npme.compute_minimum_eigenvalue(operator=H)
result = npme.compute_minimum_eigenvalue(operator=Hpauli)
# result = npme.compute_minimum_eigenvalue(operator=testH)

print(result)

ref_value = result.eigenvalue.real
print(f'Ground state energy value: {ref_value:.5f}')

from scipy.sparse import kron 

def build_operators(L: int, N: int) -> list:
    """Generate all the annihilation operators needed to build the hamiltonian
    Args:
        L (int): the cutoff of the single site Fock space
        N (int): the number of colors of gauge group SU(N)
    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
    N_bos = int(2 * (N ** 2 - 1))  # 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
    return a_b_list

a_b_list = build_operators(L,N)
# Build the Hamiltonian
# Start piece by piece
x_list = []
# only use the bosonic operators
for op in a_b_list:
    x_list.append(1 / np.sqrt(2) * (op.conjugate().transpose() + op))
# Free Hamiltonian
H_k = 0
for a in a_b_list:
    H_k = H_k + a.conjugate().transpose() * a
# vacuum energy
H_k = H_k + 0.5 * Nmat * identity(L ** Nmat)

am = MatrixOp(a_b_list[0])

# this takes a really long time even for L=4
print(am.to_pauli_op())

xm = MatrixOp(x_list[0])

print(xm.to_pauli_op())

Hm = MatrixOp(H_k)
result = npme.compute_minimum_eigenvalue(operator=Hm)

ref_value = result.eigenvalue.real
print(f'Ground state energy value: {ref_value:.5f}')

v_b1 = ( 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]
    )
v_b2 = -2.0 * ( x_list[0] * x_list[2] * x_list[3] * x_list[5]
    + x_list[0] * x_list[1] * x_list[3] * x_list[4]
    + x_list[1] * x_list[2] * x_list[4] * x_list[5]
    )
V_b = v_b1 + v_b2


vm = MatrixOp(v_b1)

vm.to_pauli_op()

vm = MatrixOp(v_b2)

vm.to_pauli_op()

# full hamiltonian
H_full = H_k + g2N / N * V_b

Hm = MatrixOp(H_full)
Hm.to_pauli_op()

result = npme.compute_minimum_eigenvalue(operator=Hm)
ref_value = result.eigenvalue.real
print(f'Ground state energy value: {ref_value:.5f}')

result = npme.compute_minimum_eigenvalue(operator=Hm.to_pauli_op())
ref_value = result.eigenvalue.real
print(f'Ground state energy value: {ref_value:.5f}')

print(result)

