import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial import KDTree

dp = 0.05 # lattice pitch
dt = 0.001 # time step

L = 1 # length
W = 0.2 # width
K = 10 # stiffness between particles
b = 10 # sping damping
drag_damping = 1 #drag damping

g = np.array([0,-9.8]) #m/s

N = int((L // dp) * (W // dp)+1)
# print(N) # you may be wondering... floor division returns type float :0 not integer 

x = np.zeros((N,2)) # [x,y] pairs, N of 'em
v = np.zeros((N,2))
u = np.zeros((N,2))

m = 0.1*np.ones(N)
r = dp/3*np.ones(N)

# lattice = ti.Vector.field(2, float, shape=(int(L/dp),int(W/dp))) # [x,y] pairs, N of 'em
# v = ti.Vector.field(2, float, shape=(int(L/dp),int(W/dp))) # velocities
# u = ti.Vector.field(2, float, shape=(int(L/dp),int(W/dp))) # input forces

# m = ti.field(dtype=float, shape=(int(L/dp),int(W/dp))) # masses
# r = ti.field(dtype=float, shape=(int(L/dp),int(W/dp))) # particle radiuses

#now we populate. 
# for i in range(N):
#     lattice[i] = [dp*(i%int(L/dp)),dp*(i%int(W/dp))]

def initLattice():
    # lattice
    for i in range(N):
        x[i][0] = dp*(i%int(L/dp))
        x[i][1] = W*dp*int(i/int(L/dp))
        v[i] = [0,0]

def initU():
    pass

def substep():
    for i in range(N):
        #gravity
        v[i]+=g*dt
        #springs
        tree = KDTree(x)
        print(tree.query(x[i],k=2))
        # for i in range(N):
        #boundary conditions
        if x[i][0] == 0:
            v[i] = [0,0]
        #step
        x[i] +=v[i]*dt

initLattice()
substep()
# for i in range(100):
#     substep()

# plt.plot(x[:,0],x[:,1],'ko',markersize=10)
# plt.show()