#File "qtipwalk.py", by KWR for CSE610C, Topics in Quantum Computation, Fall 2021
#Quantum walks on the 3-node "Q-tip" graph with edges (1,1),(1,2),(2,3),(3,3)
#Also has 3-cycle graph (commented-out)
#Usage: choose graph, edit coin matrix [[a b] [c d]] and number of iterations, and run

#Does exact integer counts before normalizing---but this overflows after iteration 64
#Separate floating-point version avoids this overflow.

import numpy as np
import math
import cmath
#from numpy import random

np.set_printoptions(precision=5, suppress=True, linewidth=160)

#J coin
a = 1
b = complex(0,1)
c = complex(0,1)
d = 1
#Hadamard coin
#a = 1
#b = 1
#c = 1
#d = -1
R = math.sqrt(2)
R2 = 2
normsq = R2
norm = R
A = np.array([[a,b,0,0,0,0],[0,0,c,d,0,0],[0,0,0,0,a,b],[c,d,0,0,0,0],[0,0,a,b,0,0],[0,0,0,0,c,d]])
#A = np.array([[0,0,0,0,a,b],[0,0,c,d,0,0],[a,b,0,0,0,0],[0,0,0,0,c,d],[0,0,a,b,0,0],[c,d,0,0,0,0]])  # 3-cycle
statevec = np.array([1,0,0,0,0,0])
print(1)
print(A)
print("\nnormalized\n")
print(A/norm)
print()
probs2 = np.array([(A[i,0]*np.conj(A[i,0])).real for i in range(6)])
probs = np.array([(probs2[2*i]+probs2[2*i+1])/normsq for i in range(3)])
statevec = A.dot(statevec)/R
print("Squares in G' are",probs2)
print("State is",statevec)
print()
print("Probabilities are",probs)
print()

A2t = A

for t in range(2,41,1):
   norm *= R
   normsq *= R2
   A2t = np.matmul(A2t,A)
   print(t)
   print(A2t)
   print()
   probs2 = np.array([(A2t[i,0]*np.conj(A2t[i,0])).real for i in range(6)])
   probs = np.array([(probs2[2*i]+probs2[2*i+1])/normsq for i in range(3)])
   statevec = A.dot(statevec)/R
   probschk2 = np.array([(statevec[i]*np.conj(statevec[i])).real for i in range(6)])
   probschk = np.array([(probschk2[2*i]+probschk2[2*i+1]) for i in range(3)])
   print("Squares in G' are",probs2)
   print("State is",statevec)
   print()
   print("G probabilities for t = ",t,"are",probs)
   print()
   print("Checked via float arithmetic:  ",probschk)
   print()