import numpy as np
from numpy import linalg as LA
from scipy.linalg import expm, sinm, cosm
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import math
from scipy import stats
%matplotlib inline

from IPython.display import Image, display, Math, Latex
sns.set(color_codes=True)

#number of vertices
n = 4

#Define adjacency matrix A_Cn
A = np.zeros((n, n))
for i in range(n):
    j1 = (i - 1)%n
    j2 = (i + 1)%n
    A[i][j1] = 1
    A[i][j2] = 1

#Define our initial state Psi_a
psi_a = np.zeros(n)
psi_a[3] = 1

#Define the time t >= 0
t = math.pi/2

#Exponentiate or hamiltonian
U_t = expm(1j*t*A)
U_mt = expm(1j*(-t)*A)

#Compute Psi_t
psi_t = U_t @ psi_a

#Compute the probabilities
prob_t = abs(psi_t)**2

M_t = U_t*U_mt 
M_t = np.around(M_t, decimals = 3)
M_t

x = M_t[:, 0].real
plt.bar(range(len(x)), x, tick_label=[0, 1, 2, 3])
plt.xlabel('Vertices')
plt.ylabel('Probability')