import taichi as ti
import time
import math

ti.init(arch=ti.cpu)

n_dim = 2
n_ptr = 8
dt    = 1e-3
steps = 50
ht    = dt * 0.5 # half of dt
max_num_particles = 1024*4

bottom_y      = 0.05
top_y         = 0.95
left_x        = 0.05
right_x       = 0.95

gravity          = ti.Vector.field(n_dim, ti.f32, shape=())
num_particles    = ti.field(ti.i32, shape=())
spring_stiffness = ti.field(ti.f32, shape=())
particle_mass    = ti.field(ti.f32, shape=())
damping          = ti.field(ti.f32, shape=())
paused           = ti.field(ti.i32, shape=())
vib_dis          = ti.field(ti.f32, shape=(3))


rest_length  = ti.field(ti.f32, shape=(max_num_particles, n_ptr))
num_ptr      = ti.field(ti.i32, shape=())
num_adj      = ti.field(ti.i32, shape=max_num_particles)
adj_ptr      = ti.field(ti.i32, shape=(max_num_particles, n_ptr))
position     = ti.Vector.field(n_dim, dtype=ti.f32, shape=max_num_particles) # position
velocity     = ti.Vector.field(n_dim, dtype=ti.f32, shape=max_num_particles) # velocity
new_velocity = ti.Vector.field(n_dim, dtype=ti.f32, shape=max_num_particles)
out_force    = ti.Vector.field(n_dim, dtype=ti.f32, shape=max_num_particles) 
is_fixed     = ti.field(dtype=ti.f32, shape=max_num_particles) 

# A @ new_velocity = b
A      = ti.Matrix.field(n_dim, n_dim, dtype=ti.f32, shape=(max_num_particles, n_ptr + 1))
b      = ti.Vector.field(n_dim, dtype=ti.f32, shape=max_num_particles)
r      = ti.Vector.field(n_dim, dtype=ti.f32, shape=max_num_particles)
new_dv = ti.Vector.field(n_dim, dtype=ti.f32, shape=max_num_particles)

resi_out = ti.field(ti.f32,shape=(2))

num_ptr[None]          = n_ptr
gravity[None]          = [0, 0] 
num_particles[None]    = 0
particle_mass[None]    = 1
damping[None]          = 0
spring_stiffness[None] = 10000
paused[None]           = False
vib_dis[0]             = 1000
vib_dis[2]             = 1e-6

@ti.func
def iterate():
    n = num_particles[None]
    for i in range(n):
        r[i] = b[i] * 1.0
        
    for i,m in ti.ndrange(n,num_ptr[None]) :
        j = adj_ptr[i,m]
        if j >= 0 :
            r[i] -= A[i, m] @ new_velocity[j]
                
    for i in range(n):
        ptr_i = num_ptr[None]
        new_dv[i] = A[i, ptr_i].inverse() @ r[i]
        
    for i in range(n):
        new_velocity[i] = new_dv[i] * 1.0


@ti.func
def resi() -> ti.f32:
    n = num_particles[None]
    res = 0.0
    
    for i in range(n):
        ptr_i = num_ptr[None]
        r[i] = b[i] * 1.0 - A[i,ptr_i] @ new_velocity[i]

    for i,m in ti.ndrange(n,num_ptr[None]) :
        j = adj_ptr[i,m]
        if j >= 0 :
            r[i] -= A[i, m] @ new_velocity[j]

    for i in range(n):
        res += r[i].x ** 2 
        res += r[i].y ** 2

    return res

@ti.kernel
def residual() -> ti.f32:
    return resi()

@ti.func
def solve_equation() :
    for no_loop in range(1) :
        for i in range(1) :
            iterate()
            if resi() < 1e-9 :
                break

@ti.func
def f_ij(i,j,ptr_j) :
    x_ij = position[j] - position[i] # from i to j
    x_d  = x_ij.normalized()
    x_n  = x_ij.norm()
    return spring_stiffness[None] * (x_n - rest_length[i, ptr_j]) * x_d

# (1 - l/xn) * (xj - xi)

# d_f/d_xi
# -[(1 - l/xn), 0] - x * l * [ x / xn^3, y / xn^3 ]
# -[0, (1 - l/xn)] - y * l * [ x / xn^3, y / xn^3 ]

# d_f/d_xj
# [(1 - l/xn), 0] + x * l * [ x / xn^3, y / xn^3 ]
# [0, (1 - l/xn)] + y * l * [ x / xn^3, y / xn^3 ]

# I * (1 - l/xn) + l/xn * x_o
# I + l/nx * (x_o - I)


@ti.func
def dfj_ij(i,ptr_j) : # df / dx
    j = adj_ptr[i,ptr_j]
    x_ij = position[j] - position[i] # from i to j
    x_d  = x_ij.normalized()
    x_n  = x_ij.norm()
    x_o  = x_d.outer_product(x_d)
    x_e = ti.Matrix([[1,0],[0,1]])
    l = rest_length[i,ptr_j]
    res =  x_e + l / x_n * (x_o - x_e)
    return spring_stiffness[None] * res

@ti.func
def substep_jacobi_semi():
    n = num_particles[None]
    n_p = num_ptr[None]
    ht2 = ht * ht

    mass = ti.Matrix([[particle_mass[None],0.0],[0.0,particle_mass[None]]])

    # init new velocity
    pm = particle_mass[None]
    for i in range(n) :
        new_f = gravity * pm + out_force[i]
        for m in ti.static(range(n_ptr)) :
            j = adj_ptr[i,m]
            if j >= 0 :
                new_f += f_ij(i,j,m)
        new_velocity[i] = velocity[i] + new_f / pm * dt
        b[i] = new_f * ht * 2 + mass @ velocity[i]


    # fill A b
    for i in range(n) :
        ptr_i = num_ptr[None]
        A[i,ptr_i] = mass
    for i,m in ti.ndrange(n,n_p) :
        ptr_i = num_ptr[None]
        j = adj_ptr[i,m]
        if j >= 0 :
            A[i,ptr_i] += ht2 * dfj_ij(i,m)
    for i,m in ti.ndrange(n,n_p) :
        ptr_i = num_ptr[None]
        j = adj_ptr[i,m]
        if j >= 0 :
            A[i,m] = - ht2 * dfj_ij(i,m)

    for i,m in ti.ndrange(n,n_p) :
        ptr_i = num_ptr[None]
        j = adj_ptr[i,m]
        if j >= 0 :
            b[i] -= ht2 * dfj_ij(i,m) @ velocity[i] 
    for i,m in ti.ndrange(n,n_p) :
        ptr_i = num_ptr[None]
        j = adj_ptr[i,m]
        if j >= 0 :
            b[i] += ht2 * dfj_ij(i,m) @ velocity[j] 

              
    # solve equation for new velocity
    solve_equation()

    if vib_dis[1] >= vib_dis[0] :
        vib_dis[1] = 0.0
        vib_dis[2] *= -1 
    vib_dis[1] += 1.0


    # Compute new position
    for i in range(n) :
        if is_fixed[i] == 2 :
            position[i].y += vib_dis[2]
            continue
        if is_fixed[i] == 1 :
            continue
        if True or position[i].y > bottom_y :
            position[i] += (velocity[i] + new_velocity[i]) * ht
            velocity[i] = new_velocity[i]
        # velocity[i] *= ti.exp(-dt * damping[None]) # damping

@ti.kernel
def step_jacobi():
    for no_loop in range(1) :
        for step in range(steps):
            substep_jacobi_semi()
            # substep_jacobi()
            pass
        
@ti.kernel
def new_particle(pos_x: ti.f32, pos_y: ti.f32): # Taichi doesn't support using Matrices as kernel arguments yet
    new_particle_id = num_particles[None]
    position[new_particle_id] = [pos_x, pos_y]
    velocity[new_particle_id] = [0, 0]
    num_particles[None] += 1

@ti.kernel
def conn_particle(pos_i: ti.i32, pos_j: ti.i32) :
    dist = (position[pos_i] - position[pos_j]).norm()
    ptr_j = num_adj[pos_i]
    rest_length[pos_i, ptr_j] = dist # 0.1
    adj_ptr[pos_i,ptr_j] = pos_j
    num_adj[pos_i] += 1

    ptr_i = num_adj[pos_j]
    rest_length[pos_j, ptr_i] = dist # 0.1
    adj_ptr[pos_j, ptr_i] = pos_i
    num_adj[pos_j] += 1


ratio_mul = 10

def init_0() :
    global ratio_mul

    ratio_mul = 100.0
    gravity[None]          = [0.0, 0.0] # -9.8]

    num_particles[None] = 0
    rest_length.fill(0)
    num_adj.fill(0)
    adj_ptr.fill(-1)
    n_xs = 40
    n_ys = 3
    dd = 0.02
    for i in range(n_xs) :
        for j in range(n_ys) :
            if i == 0 :
                is_fixed[i*n_ys+j] = 2
            if i == 1 and j == 1 :
                is_fixed[i*n_ys+j] = 0
            new_particle(left_x + i * dd, 0.5 + bottom_y + j * dd)
            if i > 0 :
                conn_particle(i*n_ys+j,(i-1)*n_ys+(j+0))
            if j > 0 :
                conn_particle(i*n_ys+j,(i-0)*n_ys+(j-1))
            if i > 0 and j > 0 :
                if (i+j) % 2 == 0 :
                    conn_particle(i*n_ys+j,(i-1)*n_ys+(j-1))
                else :
                    conn_particle((i-0)*n_ys+(j-1),(i-1)*n_ys+(j-0))

init_func_list = []
init_func_list.append(init_0)
init_state = 0

def init(st = 0) :
    global init_state

    init_state += st
    init_state %= len(init_func_list)

    vib_dis[1] = vib_dis[0] * 0.5

    return init_func_list[init_state]()

init()

gui = ti.GUI('Mass Spring System', res=(512, 512), background_color=0xaaaaaa)
gui.fps_limit = 24

def check_event() :
    for e in gui.get_events(ti.GUI.PRESS):
        if e.key in [ti.GUI.ESCAPE, ti.GUI.EXIT]:
            return False
        elif e.key == gui.SPACE:
            paused[None] = not paused[None]
        elif e.key == 'c':
            init()
        elif e.key == 'n':
            init(1)
        elif e.key == 's':
            if gui.is_pressed('Shift'):
                spring_stiffness[None] /= 1.1
            else:
                spring_stiffness[None] *= 1.1
        elif e.key == 'd':
            if gui.is_pressed('Shift'):
                damping[None] /= 1.1
            else:
                damping[None] *= 1.1
        elif e.key == 'f':
            if gui.is_pressed('Shift'):
                vib_dis[0] /= 1.1
            else:
                vib_dis[0] *= 1.1
    return True

def draw_gui() :
    X = position.to_numpy()
    Y = rest_length.to_numpy()
    Z = adj_ptr.to_numpy()
    # gui.circles(X[:num_particles[None]], color=0xffaa77, radius=1)
    
    gui.line(begin=(left_x, 0.0), end=(left_x, 1.0), color=0x0, radius=1)
    
    for i in range(num_particles[None]):
        for ptr_j,j in enumerate(Z[i]) :
            if j >= 0 and i > j :
                dist = ((X[i][0]-X[j][0])**2 + (X[i][1]-X[j][1])**2)**0.5
                ratio = (dist - Y[i][ptr_j]) / Y[i][ptr_j]
                ratio *= ratio_mul
                ratio = max(-1.0, ratio)
                ratio = min( 1.0, ratio)
                if ratio < 0 :
                    color = ti.rgb_to_hex((1.0,1.0+ratio,1.0+ratio))
                else :
                    color = ti.rgb_to_hex((1.0-ratio,1.0-ratio,1.0))
                
                gui.line(begin=X[i], end=X[j], radius=1.5, color=color)

    inv_freq = vib_dis[0] * (1.0 / 24.0 / steps)
    freq = 1.0 / inv_freq
    gui.text(content=f'C: clear all', pos=(left_x + 0.01, 0.99), color=0x0)
    gui.text(content=f'Space: pause', pos=(left_x + 0.01, 0.96), color=0x0)
    gui.text(content=f'S: Spring stiffness {spring_stiffness[None]:.1f}', pos=(left_x + 0.01, 0.93), color=0x0)
    gui.text(content=f'F: Freq {freq:.4f}Hz T={inv_freq:.4f}s', pos=(left_x + 0.01, 0.90), color=0x0)



def run() :
    while True:
        if check_event() == False :
            break
                    
        if not paused[None]:
            step_jacobi()
            # print(f'residual={residual():0.10f}')
        
        draw_gui()
        gui.show()


def export_png() : # export result as png file
    for i_frame in range(24*1) :
        if check_event() == False :
            break
                    
        if not paused[None]:
            step_jacobi()
        
        draw_gui()
        
        filename = f'frame_{i_frame:05d}.png'   # create filename with suffix png
        print(f'Frame {i_frame} is recorded in {filename}')
        gui.show(filename)  # export and show in GUI
    exit()

# export_png() # export result as png file
run()
