# Here's where the number of qubits are set. The Hamiltonian will be have a size of (N-2)x(N-2)
qubits = 8
N = (2**qubits)+2
print(f"The Hamiltonian will be a {N-2}x{N-2} matrix")

#lattice lower bound
a =  -6
#lattice upper bound
b =  6

# Import necessary libraries
import numpy as np
import matplotlib.pyplot as plt

# Define x grid and step size
x = np.linspace(a,b,N)
h = x[1]-x[0]

print(f"Discretized grid with {N} sites and step size {h:.4f}")
plt.plot(x,np.zeros(x.shape),'o', label="Grid points")
plt.legend(loc='upper right');

print(f"The extremal points will not be considered. That is why we have a {N-2}x{N-2} Hamiltonian but {N} points.")

# Define potential energy function. Change this to change the potential.
def Vpot(x):
    return (1/2)*x**2


# Create Kinetic Energy Matrix
T = np.zeros((N-2)**2).reshape(N-2,N-2)
for i in range(N-2):
    for j in range(N-2):
        if i==j:
            T[i,j]= -2
        elif np.abs(i-j)==1:
            T[i,j]=1
        else:
            T[i,j]=0

# Create Potential Energy Matrix
V = np.zeros((N-2)**2).reshape(N-2,N-2)
for i in range(N-2):
    for j in range(N-2):
        if i==j:
            V[i,j]= Vpot(x[i+1])
        else:
            V[i,j]=0
            
# Create Hamiltonian: a big matrix 😂
H = -T/(2*h**2) + V

Tp = np.zeros((N-2,N-2))  # create a kinetic matrix full of zeros 
T_id = np.diag_indices(Tp.shape[0])  # select the indices on the diagonal
T_idp = (T_id[0][:-1],(T_id[1][:-1]+1))  # select the indices above the diagonal y-x=1
T_idm = (T_id[0][:-1]+1,T_id[1][:-1])  # select the indices below the diagonal x-y=1
# fill the 3 "diagonals" with entries for the discretized derivative
Tp[T_id] = -2.0
Tp[T_idp] = Tp[T_idm] = 1.0
# check it is the same as the matrix created with for loops
assert(np.allclose(T,Tp))

Tp.shape

Vp = np.zeros((N-2,N-2))
V_id = np.diag_indices(Vp.shape[0])
Vp[V_id] = Vpot(x[1:-1])
assert(np.allclose(V,Vp))

Vp.shape

# this has 2 more sites
x.shape

# Create Hamiltonian: a big matrix 😂
Hp = -Tp/(2*h**2) + Vp
assert(np.allclose(H,Hp))

plt.imshow(H, cmap='jet')
plt.colorbar();

# Find eigenvalues and eigenvectors, then sort them in ascending order
val, vec = np.linalg.eig(H)
z = np.argsort(val)
z = z[0:4]
energies = val[z]
print("Finite Difference Energies for the first 4 states: ", energies)


# Compare finite difference calculation to theoretical calculation
nVec = np.linspace(0, len(energies) - 1, len(energies))
En = nVec + 0.5
print("Theoretical Energies: ", En)
percentError = ((En - energies) / (energies)) * 100
print("Percent Error: ", percentError)


# potential energy
# get the ground state
gs_vec = vec[:,z[0]]
# sandwich the potential energy
gs_vec.T.dot(Vp.dot(gs_vec))

En[0]/2

# Plot wavefunctions for first 4 lowest states
plt.figure(figsize=(12,10))
scalingFactor=1.5
for i in range(len(z)):
    y = []
    y = np.append(y,vec[:,z[i]])
    y = np.append(y,0)
    y = np.insert(y,0,0)
    plt.plot(x,scalingFactor*y+energies[i],lw=3, label="{} ".format(i))
plt.xlabel('x', size=14)
plt.ylabel(r'$\Psi$(x)',size=14,color='tab:red')
plt.legend(title="nth Quantum State")
plt.title('normalized wavefunctions for a harmonic oscillator using finite difference method',size=14)

# function to compute the ground state energy as a function of R and L
def gs_energy(R, L, potential=Vpot):
    # add boundary points
    # Define x grid and step size
    x = np.linspace(-R, R, L)
    h = x[1] - x[0]
    # print(R,L,h)
    # add boundary
    x = np.append(x, x[-1] + h)
    x = np.insert(x, 0, x[0] - h)
    # kinetic
    Tp = np.zeros((L, L))  # create a kinetic matrix full of zeros
    T_id = np.diag_indices(Tp.shape[0])  # select the indices on the diagonal
    T_idp = (
        T_id[0][:-1],
        (T_id[1][:-1] + 1),
    )  # select the indices above the diagonal y-x=1
    T_idm = (
        T_id[0][:-1] + 1,
        T_id[1][:-1],
    )  # select the indices below the diagonal x-y=1
    # fill the 3 "diagonals" with entries for the discretized derivative
    Tp[T_id] = -2.0
    Tp[T_idp] = Tp[T_idm] = 1.0
    # Potential
    Vp = np.zeros((L, L))
    V_id = np.diag_indices(Vp.shape[0])
    Vp[V_id] = potential(x[1:-1])
    # Hamiltonian
    Hp = -Tp / (2 * h**2) + Vp
    # Ground state
    val, vec = np.linalg.eig(Hp)
    z = np.argsort(val)
    # print(val[z[0]])
    delta_energy = np.fabs(val[z[0]] - 0.5)
    # get the ground state
    gs_vec = vec[:, z[0]]
    # sandwich the potential energy
    potential_energy = gs_vec.T.dot(Vp.dot(gs_vec))
    return np.array([h, R, L, delta_energy, potential_energy], dtype=np.float64)


gs_energy(11,1050)

from tqdm.notebook import tqdm

df = list()
# for r in tqdm([0.5,1.,1.5,2.,2.5,3.,3.5,4.,5.,7.,9.,11.], desc='R'):
#   for l in tqdm([50,150,250,350,450,550,650,850,1050], desc='L'):
for r in tqdm([4.,5.,7.,9.,11.], desc='R'):
  for l in tqdm([450,550,650,850,1050], desc='L'):
    df.append(gs_energy(r,l))
df = np.asarray(df)

import pandas as pd

df = pd.DataFrame(df, columns=["A","R","L","D","P"], dtype=np.float64)

df.head()

import seaborn as sns 

g = sns.relplot(data=df, x="L", y="D", hue="R", palette='tab20', kind='line', height=4, aspect=2)
g.set_axis_labels(r"$\Lambda$", r"$\Delta E_0$")
g.set(yscale='log', xscale='log')

g = sns.relplot(data=df, x="A", y="P", hue="L", palette='tab20', kind='line', height=4, aspect=2)
g.set_axis_labels(r"$a_{dig}$", r"$V_p$")
g.set(yscale='log', xscale='log')

a_dig = 0.7


def ir_cutoff(a_dig):
    return 1_000 * a_dig


def uv_cutoff(a_dig):
    return int(1 + 2 * ir_cutoff(a_dig) / a_dig)


R = ir_cutoff(a_dig)
L = uv_cutoff(a_dig)
print(f"Choosing a_dig={a_dig} with R={R} and L={L}")


# Define potential energy function. Change this to change the potential.
# M^2 = 1
def Vpot_quartic(x):
    return (1/2)*x**2 + (1/4)*x**4

# M^2 = -1
def Vpot_quartic_neg(x):
    return -(1/2)*x**2 + (1/4)*x**4

gs_energy(R,L,potential=Vpot_quartic)

gs_energy(R,L,potential=Vpot_quartic_neg)

#  T = 0.1
T = 0.1
beta = 1.0/T
print(f"Temperature = {T} \n (beta={beta})")

# Example: correction to the 0.5 ground state energy of the harmonic oscillator at zero temperature
np.fabs(0.5 - 0.5*1./np.tanh(beta/2))

def harmonic_energy(beta):
    return 0.5*1./np.tanh(beta/2)


# Example: harmonic oscillator potential energy at finite temperature
v_beta = np.sum([np.exp(-beta*val[i])*(vec[:,i].T.dot(Vp.dot(vec[:,i]))) for i in z])
# normalize using the partition function
z_beta = np.sum(np.exp(-beta*val))
v_beta /= z_beta
print(f"Potential energy at beta={beta} : {v_beta}")

# result "in the continuum": it is slightly higher than 0.25 because of thermal fluctuations
harmonic_energy(beta)/2

# function to compute the average potential energy at finite temperature as a function of R and L
def pot_energy(beta, R, L, potential=Vpot):
    # add boundary points
    # Define x grid and step size
    x = np.linspace(-R, R, L)
    h = x[1] - x[0]
    # print(R,L,h)
    # add boundary
    x = np.append(x, x[-1] + h)
    x = np.insert(x, 0, x[0] - h)
    # kinetic
    Tp = np.zeros((L, L))  # create a kinetic matrix full of zeros
    T_id = np.diag_indices(Tp.shape[0])  # select the indices on the diagonal
    T_idp = (
        T_id[0][:-1],
        (T_id[1][:-1] + 1),
    )  # select the indices above the diagonal y-x=1
    T_idm = (
        T_id[0][:-1] + 1,
        T_id[1][:-1],
    )  # select the indices below the diagonal x-y=1
    # fill the 3 "diagonals" with entries for the discretized derivative
    Tp[T_id] = -2.0
    Tp[T_idp] = Tp[T_idm] = 1.0
    # Potential
    Vp = np.zeros((L, L))
    V_id = np.diag_indices(Vp.shape[0])
    Vp[V_id] = potential(x[1:-1])
    # Hamiltonian
    Hp = -Tp / (2 * h**2) + Vp
    # Spectrum: use eigh because eig will not return all eigenvalues correctly (contains nan)
    val, vec = np.linalg.eigh(Hp)
    ei = np.argsort(val)
    # partition function
    z_beta = np.sum(np.exp(-beta * val))
    # potential energy
    v_beta = np.sum(
        [
            np.exp(-beta * val[i]) * (vec[:, i].T.dot(Vp.dot(vec[:, i]))) / z_beta
            for i in ei[::-1]
        ]
    )
    return np.array([h, R, L, z_beta, v_beta], dtype=np.float64)


pot_energy(1/T, a, N, potential=Vpot)

a_dig = 0.3
R = ir_cutoff(a_dig)
L = uv_cutoff(a_dig)
T = 0.1
beta = 1/T
print(f"Quadratic potential: {pot_energy(beta, R, L, potential=Vpot)}")
print(f"Quartic msq=1 potential: {pot_energy(beta, R, L, potential=Vpot_quartic)}")
print(f"Quartic msq=-1 potential: {pot_energy(beta, R, L, potential=Vpot_quartic_neg)}")

a_dig = 0.5
R = ir_cutoff(a_dig)
L = uv_cutoff(a_dig)
T = 0.1
beta = 1/T
print(f"Quadratic potential: {pot_energy(beta, R, L, potential=Vpot)}")
print(f"Quartic msq=1 potential: {pot_energy(beta, R, L, potential=Vpot_quartic)}")
print(f"Quartic msq=-1 potential: {pot_energy(beta, R, L, potential=Vpot_quartic_neg)}")

a_dig = 0.7
R = ir_cutoff(a_dig)
L = uv_cutoff(a_dig)
T = 0.1
beta = 1/T
print(f"Quadratic potential: {pot_energy(beta, R, L, potential=Vpot)}")
print(f"Quartic msq=1 potential: {pot_energy(beta, R, L, potential=Vpot_quartic)}")
print(f"Quartic msq=-1 potential: {pot_energy(beta, R, L, potential=Vpot_quartic_neg)}")

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))
print(qubitOp)
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'], entanglement="full", reps=3)
display(var_form.draw(output='mpl',fold=100))

from qiskit.algorithms import NumPyEigensolver
# exactly diagonalize the system using numpy routines
solver = NumPyEigensolver(k=4)
exact_solution = solver.compute_eigenvalues(qubitOp)
print("Exact Result of qubit hamiltonian:", np.real(exact_solution.eigenvalues))
print("Exact Result of discrete hamiltonian:", energies[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

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

warnings.filterwarnings("ignore")

rngseed = 0
algorithm_globals.random_seed = rngseed

if HAS_GPU:
  # with GPU
  backend = Aer.get_backend("statevector_simulator", method="statevector_gpu")
else:
  # without GPU
  backend = Aer.get_backend("statevector_simulator", max_parallel_threads=1, max_parallel_experiments=0)
q_instance = QuantumInstance(
    backend, seed_transpiler=rngseed, seed_simulator=rngseed
)

optimizer = SLSQP(maxiter=600)

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

# Run the VQE
vqe = VQE(
    ansatz=var_form,
    optimizer=optimizer,
    quantum_instance=q_instance,
    callback=store_intermediate_result,
)
ret = vqe.compute_minimum_eigenvalue(qubitOp)
vqe_result = np.real(ret.eigenvalue)
print("VQE Result:", vqe_result)

# Plot convergence plot
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')

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

rngseed = 0
algorithm_globals.random_seed = rngseed

if HAS_GPU:
  backendQasm = Aer.get_backend('qasm_simulator', method="statevector_gpu")
else:
  backendQasm = Aer.get_backend('qasm_simulator', method="statevector")
  
q_instance = QuantumInstance(backend=backendQasm, shots=1024, seed_transpiler=rngseed, seed_simulator=rngseed)
optimizer = SPSA(maxiter=600)

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

# Run the VQE
vqe = VQE(
    ansatz=var_form,
    optimizer=optimizer,
    quantum_instance=q_instance,
    callback=store_intermediate_result,
)
ret = vqe.compute_minimum_eigenvalue(qubitOp)
vqe_result = np.real(ret.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 for various optimizers')
plt.legend(loc='upper right')

import qiskit.tools.jupyter
%qiskit_version_table