import numpy as np
import taichi as ti
import sounddevice as sd
from scipy.fft import fft, fftfreq
import matplotlib.pyplot as plt

ti.init(arch=ti.cpu)

n = 200
dt = 0.1 / n
substeps = int(1 / 60 // dt)


x = ti.Vector.field(2, dtype=float, shape=(n,))

v = ti.Vector.field(2, dtype=float, shape=(n,))

quad_size = 1.0 / n

spring_Y = 10
dashpot_damping = 10
drag_damping = 1

@ti.kernel
def initializeXV():
    for i in x:
        x[i] = [i * quad_size,  0.5]
        v[i] = [0,0]

@ti.kernel
def substep():
    #ti.loop_config(serialize=True) # Disable auto-parallelism in Taichi
    # for i in range(1,(n-2)):
    for i in x:
        if i!=0 and i!=(n-1):
            force = ti.Vector([0.0,0.0])
            l0 = ti.Vector([quad_size,0.0])
            fdiff = (x[i+1] - x[i])
            vfdiff = (v[i+1]-v[i])
            force += -spring_Y *(fdiff-l0)
            force += -vfdiff * dashpot_damping 
            bdiff = -x[i-1] + x[i]
            vbdiff = -v[i-1]+v[i]
            force += -spring_Y *(bdiff-l0)
            force += -vbdiff * dashpot_damping 
            #print(force)
            v[i] += force * dt
        elif i==0:
            force = ti.Vector([0.0,0.0])
            l0 = ti.Vector([quad_size,0.0])
            fdiff = (x[i+1] - x[i])
            vfdiff = (v[i+1]-v[i])
            force += -spring_Y *(fdiff-l0)
            force += -vfdiff * dashpot_damping 
        else:
            force = ti.Vector([0.0,0.0])
            l0 = ti.Vector([quad_size,0.0])
            bdiff = -x[i-1] + x[i]
            vbdiff = -v[i-1]+v[i]
            force += -spring_Y *(bdiff-l0)
            force += -vbdiff * dashpot_damping 
    
    for i in x:
        v[i] *= ti.exp(-drag_damping * dt)
    for i in x:
        x[i] += dt * v[i]
        if i ==0:
            x[i] == [0,0.5]
            v[i] == [0,0]
        elif i == (n-1):
            x[i] == [0,0.5]
            v[i] == [0,0]

            

gui = ti.GUI('click to add impulse, with damping', res=(640, 400))

initializeXV()

while gui.running:
    if gui.get_event(ti.GUI.PRESS,ti.GUI.RMB):
        mouse_x, mouse_y = gui.get_cursor_pos()
        #print(x[int(mouse_x*n)])
        x[int(mouse_x*n)] = [x[int(mouse_x*n)][0],mouse_y]

    for i in range(substeps):
        substep()
    pos = x.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()

tune = pos[:,1] #just the y values.
tunef = fft(tune)
xf = fftfreq(n,quad_size)[:n//2]
plt.plot(xf,2.0/n * np.abs(tunef[0:n//2]))
plt.show()