# this is the generic numerical simulation code for the theory part of the 
# following figures:
# 1. Main Figure 2(a-b) theory
# 2. Main Figure 3(a-b) theory
# 3. Supp Figure 5(a) theory
# 4. Supp Figure 6(a-c) theory
# One should manually tune the parameters in this simulation code for all the 
# above figures


#this is the code for the 2D HOTI
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.cm as cm           # import colormap stuff!
import matplotlib as mpl
from mpl_toolkits.mplot3d import Axes3D


#print((44.558+44.517)/4+(45.064+45.0234)/4-(29.6974+29.7176)/2)

delta_c = -0.175
delta_f = -0.175

amp = 1.0
print("coarse_ratio:",max((1+delta_c)/(1-delta_c),(1-delta_c)/(1+delta_c))**2)
print("fine_ratio:",max((1+delta_f)/(1-delta_f),(1-delta_f)/(1+delta_f))**2)
print(6/(max((1+delta_c)/(1-delta_c),(1-delta_c)/(1+delta_c))**2))
total_step = 40 # fix this to be an EVEN number!
N = total_step

turn_off_step = 1 + 24
#extra_step = 6

half_scope = 16



def BS(U,L):
        
    return (U+L)/np.sqrt(2), (U-L)/np.sqrt(2)

def inverse_BS(U,L):
    return (U-L)/np.sqrt(2), (U+L)/np.sqrt(2)


def small_delay(U,L):
    return np.roll(U,-1,axis=1),np.roll(L,1,axis=1)

def large_delay(U,L):
    return np.roll(U,-1,axis=0),np.roll(L,1,axis=0)

# A = [[0,1,0],
#      [0,1,1],
#      [0,1,0]]
# print(np.roll(A,-1,axis=1))


def fine_modulation(U,L,d_f,f_s):
    for i in range(len(U)):
      for j in range(len(U[0])):
        if j<= N + f_s :
            U[i][j] = U[i][j]*np.exp(d_f)#*(1+d_f)
        else:
            U[i][j] = U[i][j]*np.exp(-d_f)#*(1-d_f)
    for i in range(len(L)):
      for j in range(len(U[0])):
        if j<= N + f_s:
            L[i][j] = L[i][j]*np.exp(-d_f)#*(1-d_f)
        else :
            L[i][j] = L[i][j]*np.exp(d_f)#*(1+d_f)
    return U,L


##def fine_modulation_2(U,L,d_f):
##    for i in range(len(U)):
##      for j in range(len(U[0])):
##        if j<N and j%2==0:
##            U[i][j] = U[i][j]*(1+d_f)
##        elif j%2==0:
##            U[i][j] = U[i][j]*(1-d_f)
##        else:
##            U[i][j] = U[i][j]*(1.0) 
##        
##    for i in range(len(L)):
##      for j in range(len(L[0])):
##        if j<N and j%2==0:
##            L[i][j] = L[i][j]*(1-d_f)
##        elif j%2==0:
##            L[i][j] = L[i][j]*(1+d_f)
##        else:
##            L[i][j] = L[i][j]*(1.0)
##    return U,L


def coarse_modulation(U,L,d_c,c_s):
    for i in range(len(U)):
        if i<N + c_s:
            U[i] = U[i]*np.exp(d_c)#*(1+d_c)
        else :
            U[i] = U[i]*np.exp(-d_c)#*(1-d_c)
    for i in range(len(L)):
        if i<N + c_s:
            L[i] = L[i]*np.exp(-d_c)#*(1-d_c)
        else:
            L[i] = L[i]*np.exp(d_c)#*(1+d_c)
    return U,L


##def coarse_modulation_2(U,L,d_c):
##    for i in range(len(U)):
##        if i<N and i%2==0:
##            U[i] = U[i]*(1+d_c)
##        elif i%2==0 :
##            U[i] = U[i]*(1-d_c)
##        else:
##            U[i] = U[i]*(1.0)
##    for i in range(len(L)):
##        if i<N and i%2==0:
##            L[i] = L[i]*(1-d_c)
##        elif i%2==0:
##            L[i] = L[i]*(1+d_c)
##        else:
##            U[i] = U[i]*(1.0)
##    return U,L


# def one_step(U,L):
#     U1,L1 = BS_operation(U, L)
#     U2,L2 = fine_modulation(U1, L1)
#     U3,L3 = small_delay(U2,L2)
#     U4,L4 = BS_operation(U3, L3)
#     U5,L5 = coarse_modulation(U4, L4)
#     U6,L6 = large_delay(U5,L5)
#     return U6,L6




def walk_N_step(U,L,d_c,d_f,n,c_s,f_s):
    for i in range(n):
        U,L = BS(U, L)
        U,L = fine_modulation(U , L, d_f , f_s)
        U,L = small_delay(U , L)
        U,L = BS(U, L)
        U,L = coarse_modulation(U, L, d_c , c_s)
        U,L = large_delay(U,L)
        #U,L = U*amp, L*amp
    return U,L


def initialize(N):   #define two (2N+1)*(2N+1) matrix, upper and lower channel
    U = [[0.0 for i in range(2*N)] for i in range(2*N)]
    L = [[0.0 for i in range(2*N)] for i in range(2*N)]
    L[N][N] = 1.0
    #U[N-1][N] = 0.0
    return np.array(U),np.array(L)

def trajectory(N):
    U,L = initialize(N)
    tra_U = []
    tra_L = []
    for i in range(N):
        U,L = walk_N_step(U, L, 1)
        tra_U.append(U)
        tra_L.append(L)
    return np.array(tra_U),np.array(tra_L)   



'''
test_U = np.array([[1.0 for i in range(2*N)] for i in range(2*N)])
test_L = np.array([[1.0 for i in range(2*N)] for i in range(2*N)])

output_U, output_L = fine_modulation(test_U,test_L,delta_c)
oo_U, oo_L = coarse_modulation(output_U, output_L,delta_f)
plt.imshow(output_U)
'''





#======================conducting experiments==============================#
U0,L0 = initialize(total_step)


# U1,L1 = walk_N_step(U0, L0, delta_c, delta_f, 8+1, 0, 0)
# U2,L2 = walk_N_step(U1, L1, delta_c, delta_f, 3 , -10, 0)
# U3,L3 = walk_N_step(U2,L2, delta_c, delta_f, 18+1-11, -10, 0)
# U4,L4 = walk_N_step(U3,L3, delta_c, delta_f, 28+1-18, -10, 0)


U1,L1 = walk_N_step(U0, L0, delta_c, delta_f, 1, 0, 0)
U2,L2 = walk_N_step(U1, L1, delta_c, delta_f, 4 , 0, 0)
U3,L3 = walk_N_step(U2,L2, delta_c, delta_f, 4, 0, 0)
U4,L4 = walk_N_step(U3,L3, delta_c, delta_f, 4, 0, 0)


#===================plotting the simulation results========================#
final_U1 = np.flipud(np.fliplr(U1))
final_L1 = np.flipud(np.fliplr(L1))
final_U2 = np.flipud(np.fliplr(U3))
final_L2 = np.flipud(np.fliplr(L3))
final_U3 = np.flipud(np.fliplr(U4))
final_L3 = np.flipud(np.fliplr(L4))


shot1 = np.abs((final_U1[total_step-half_scope:total_step+half_scope\
                    ,total_step-half_scope:total_step+half_scope])**2)

shot2 = np.abs((final_U2[total_step-half_scope:total_step+half_scope\
                        ,total_step-half_scope:total_step+half_scope])**2)

shot3 = np.abs((final_U3[total_step-half_scope:total_step+half_scope\
                        ,total_step-half_scope:total_step+half_scope])**2)
   
    
    

fig, (ax1, ax2, ax3) = plt.subplots(ncols=3, figsize=(17.5, 3.7))
x_label_list = ['-15', '-0', '+15']
y_label_list = ['+15', '-0', '-15']
font_axis = {'family' :'','weight' : 'normal','size' : 15}

#ax = fig.add_subplot(1, 2, 1)
neg1 = ax1.imshow(shot1,cmap=plt.cm.coolwarm)#, vmin=min_val, vmax=max_val)
for axis in ['top','bottom','left','right']:
    ax1.spines[axis].set_linewidth(2)
ax1.tick_params(axis = "both",direction = "out",length= 10 ,width =3 ,\
               labelsize = 22)
ax1.set_xticks([1,16,31])
ax1.set_xticklabels(x_label_list)
ax1.set_yticks([1,16,31])
ax1.set_yticklabels(y_label_list)


#ax2 = fig.add_subplot(1, 2, 2)
neg2 = ax2.imshow(shot2,cmap=plt.cm.coolwarm)#, vmin=min_val, vmax=max_val)
for axis in ['top','bottom','left','right']:
    ax2.spines[axis].set_linewidth(2)
ax2.tick_params(axis = "both",direction = "out",length= 10 ,width =3 ,\
               labelsize = 22)
ax2.set_xticks([1,16,31])
ax2.set_xticklabels(x_label_list)
ax2.set_yticks([1,16,31])
ax2.set_yticklabels(y_label_list)


neg3 = ax3.imshow(shot3,cmap=plt.cm.coolwarm)#, vmin=min_val, vmax=max_val)
for axis in ['top','bottom','left','right']:
    ax3.spines[axis].set_linewidth(2)
ax3.tick_params(axis = "both",direction = "out",length= 10 ,width =3 ,\
               labelsize = 22)
ax3.set_xticks([1,16,31])
ax3.set_xticklabels(x_label_list)
ax3.set_yticks([1,16,31])
ax3.set_yticklabels(y_label_list)

ax1.title.set_text("step" + " " + str(8))
ax2.title.set_text("step" + " " + str(18))
ax3.title.set_text("step" + " " + str(28))


cbar = fig.colorbar(neg3, ticks=[np.min(shot3), (np.min(shot3)+np.max(shot3))/2,
                                 np.max(shot3)],aspect=16)
cbar.ax.set_yticklabels(['0', '0.5', '1'],fontsize = 10)
#cbar.ax.tick_params(axis = "both",direction = "out",length= 8 ,width =3 ,\
#               labelsize = 10)

#plt.imshow(shot1)
plt.show()
#plt.colorbar(location='right')


#=======================end figure paper plotting====================#


'''
test_U,test_L = trajectory(total_step)
abs_U = np.abs(test_U)
normalized = []
for ele in abs_U:
    normalized.append(ele/np.linalg.norm(ele))

plt.imshow(abs_U[31])

'''

'''
# Change the Size of Graph using
# Figsize
fig = plt.figure(figsize=(10, 10))
 
# Generating a 3D sine wave
ax = plt.axes(projection='3d')
 
# Creating array points using
# numpy
x = []
y = []
z = []
c = []
for i in range(2*total_step+1):
    for j in range(2*total_step+1):
        for k in range(total_step):
            if abs_U[k][j][i]>=0.4:
               x.append(i)
               y.append(j)
               z.append(k)
               c.append(normalized[k][j][i])
 
# To create a scatter graph
ax.scatter(x, y, z,c = c,cmap = "Greens")
 
# trun off/on axis
plt.axis('off')
 
# show the graph
plt.show()
'''