import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import taichi as ti

ti.init(arch=ti.cpu)


N = 200
dx = 1
dt = 0.1
D = 1
C = D*dt/(2*dx**2)

# u = np.random.random((N,N))#np.zeros((N,2))
# u_next = u.copy()

#NxN field of size 1 vectors
u = ti.Vector.field(1, dtype=float, shape=(N,N))
uprev = ti.Vector.field(1, dtype=float, shape=(N,N))

@ti.kernel
def init():
    for i,j in u:
        if i == 0 or j == 0 or i == (N-1) or j == (N-1):
            u[i,j] = [0]
        else:
            u[i,j] = [ti.random()]
        uprev[i,j] = u[i,j]


init()
# u_half = u.copy()#np.zeros((N,2))
# u_full = np.zeros((N,N))

# def getStep():
#     for i in range(1,N-2):
#         for j in range(1,N-2):
#             u_half[i,j] = u[i,j] + C*(u_half[i+1,j]-2*u_half[i,j]+u_half[i-1,j]+u[i+1,j]-2*u[i,j]+u[i-1,j])
#     for i in range(1,N-2):
#         for j in range(1,N-2):
#             u_full[i,j] = u_half[i,j] + C*(u_half[i+1,j]-2*u_half[i,j]+u_half[i-1,j]+u_full[i+1,j]-2*u[i,j]+u[i-1,j])
#     return u_half, u_full


#print(u.to_numpy)
un = u.to_numpy()
#print(np.shape(un))
print(un[:,:,0])

def update_frame(n):
    # u_half,u_full = getStep()
    # u = u_full
    un = u.to_numpy()
    return un


fig, axes = plt.subplots()
sns.set_theme(style="white")
cmap = sns.cubehelix_palette(start=0, light=1, as_cmap=True)

# sns.kdeplot(data=un[2],
#         x=un[0], y=un[1],
#         cmap=cmap, fill=True,
#         clip=(-2, 2), cut=10,
#         thresh=0, levels=15,
#         ax=axes,
#     )

sns.heatmap(un[:,:,0],
        cmap=cmap,
        ax=axes,
    )
axes.set_axis_off()

ani = animation.FuncAnimation(
    fig, update_frame, interval=20)

plt.show()