import numpy as np
import taichi as ti

ti.init(arch=ti.cpu)

N=500
D = 1
dt = 0.1

# substeps = int(1 / 60 // dt)
#print(substeps)

L = 1

#print(D*dt/((L/N)**2))
C = dt
u = ti.Vector.field(2, dtype=float, shape=(N,))
uprev = ti.Vector.field(2, dtype=float, shape=(N,))

@ti.kernel
def init():
    for i in u:
        if i==int(N/2):
            u[i] = [i*L/N,1]
            uprev[i] = [i*L/N,1]
        else:
            u[i] = [i*L/N,0]
            uprev[i] = [i*L/N,0]

@ti.kernel
def nextStep():
    for i in u:
        if i == 0:
            u[i][1] = 0
        elif i == N-1:
            u[i][1] = 0
        else:
            u[i][1] = uprev[i][1] +C*(uprev[i+1][1]-2*uprev[i][1]+uprev[i-1][1])
        uprev[i] = u[i]
        #print(u[i])
        


gui = ti.GUI('title', res=(640, 400))

init()
pos = u.to_numpy()#np.column_stack((x_range_resize,(u_next.to_numpy()+0.5)))
#print(pos)
gui.lines(begin=pos[:(N-2)],end=pos[1:(N-1)],radius=2,color=0x068587)
gui.show()

while gui.running:
    nextStep()
    pos = u.to_numpy()#np.column_stack((x_range_resize,(u_next.to_numpy()+0.5)))
    #print(pos)
    gui.lines(begin=pos[:(N-2)],end=pos[1:(N-1)],radius=2,color=0x068587)
    gui.show()