import taichi as ti
import numpy as np
from matplotlib.image import imsave

# ti.init(arch=ti.cuda, device_memory_GB=8)
ti.init(arch=ti.cuda)
BENCHMARK = False
# N_POINTS = 550000000
# MESH_RES = 50
# UPDATES_PER_BATCH = 1
N_POINTS = 2500000
MESH_RES = 100
UPDATES_PER_BATCH = 1

D = 0.05
positions = ti.Vector.field(2,dtype=float, shape=(N_POINTS,))
z = ti.field(dtype=float, shape=(MESH_RES, MESH_RES))
z_prev = ti.field(dtype=float, shape=(MESH_RES, MESH_RES))

mesh_render_size = 5 # width in viewport
vertices = ti.Vector.field(3,dtype=float, shape=(MESH_RES**2))
indices = ti.field(dtype=int, shape=(3*2*(MESH_RES-1)**2))
colors = ti.Vector.field(3,dtype=float, shape=(MESH_RES**2))


# Z_SCALE = 400000
Z_SCALE = 400
colormap = [
    [64,57,144],
    [112,198,162],
    [230, 241, 146],
    [253,219,127],
    [244,109,69],
    [169,23,69]
]

colormap_field = ti.Vector.field(3,dtype=float, shape = len(colormap))
for i in range(len(colormap)):
    colormap_field[i] = ti.Vector(colormap[i])

@ti.kernel
def apply_initial_conditions():
    for i in positions:
        # circular initial distro
        # positions[i].x = ti.random() * ti.cos(ti.random()*2*np.pi)
        # positions[i].y = ti.random() * ti.sin(ti.random()*2*np.pi)

        # point source
        # positions[i].x = 0
        # positions[i].y = 0

        # square 
        # positions[i].x = -1*ti.random() 
        # positions[i].y = -1*ti.random() 

        # two squares
        side = ti.round(ti.random())*2-1
        positions[i].x = side*ti.random() 
        positions[i].y = side*ti.random() 
        x_cell = ti.round((positions[i].x + 1)/2 * (MESH_RES -1),dtype=int)
        y_cell = ti.round((positions[i].y + 1)/2 * (MESH_RES -1),dtype=int)
        # z_prev[x_cell, y_cell] += 1 / color_height
    for i, j in z_prev:
        if i < MESH_RES/2 and j < MESH_RES/2:
            z_prev[i,j] = ((N_POINTS / 2) / (MESH_RES /2)**2) / Z_SCALE
        if i > MESH_RES/2 and j > MESH_RES/2:
            z_prev[i,j] = ((N_POINTS / 2) / (MESH_RES /2)**2) / Z_SCALE

@ti.kernel
def update_positions():
    for i in ti.grouped(z):
        z[i] = 0.0
    ti.sync()
    for i in positions:
        for j in range(UPDATES_PER_BATCH):
            direction = ti.random()*2*np.pi
            next_x = D*ti.cos(direction)
            next_y = D*ti.sin(direction)
            if positions[i].x+next_x> 1:
                positions[i].x = 2-(positions[i].x+next_x)
            elif positions[i].x+next_x < -1:
                positions[i].x = -2-(positions[i].x+next_x)
            else:
                positions[i].x += next_x
            if positions[i].y +next_y> 1:
                positions[i].y = 2-(positions[i].y+next_y)
            elif positions[i].y+next_y < -1:
                positions[i].y = -2-(positions[i].y+next_y)
            else:
                positions[i].y += next_y
            
            if j == UPDATES_PER_BATCH-1:
                x_cell = ti.round((positions[i].x + 1)/2 * (MESH_RES -1),dtype=int)
                y_cell = ti.round((positions[i].y + 1)/2 * (MESH_RES -1),dtype=int)
                z[x_cell, y_cell] += 1 / Z_SCALE
    ti.sync()




alpha = 1.0
@ti.kernel
def update_mesh_vertices():
    for i,j in z:
        h = z[i,j]*alpha + z_prev[i,j]*(1-alpha)
        # h = + z_prev[i,j]*(1-alpha)
        vertices[i+MESH_RES*j].y = h
        z_prev[i,j] = h
        
        # h = (z[i,j]+2)/4 
        level = ti.min(ti.floor(h*(colormap_field.shape[0]-1)), colormap_field.shape[0]-2)
        colorphase = h*(colormap_field.shape[0]-1) - level
        level_idx = ti.cast(level, dtype=int)

        colors[i+MESH_RES*j].x = (colormap_field[level_idx].x * (1-colorphase) + colorphase*colormap_field[level_idx+1].x)/255
        colors[i+MESH_RES*j].y = (colormap_field[level_idx].y * (1-colorphase) + colorphase*colormap_field[level_idx+1].y)/255
        colors[i+MESH_RES*j].z = (colormap_field[level_idx].z * (1-colorphase) + colorphase*colormap_field[level_idx+1].z)/255


@ti.kernel
def init_mesh_vertices():
    for i,j in z:
        vertices[i+MESH_RES*j].x = i/(MESH_RES-1)*mesh_render_size-mesh_render_size/2
        vertices[i+MESH_RES*j].y = 0
        vertices[i+MESH_RES*j].z = j/(MESH_RES-1)*mesh_render_size-mesh_render_size/2
@ti.kernel
def init_mesh_indices():
    for i, j in ti.ndrange(MESH_RES - 1, MESH_RES - 1):
            quad_id = (i * (MESH_RES - 1)) + j
            # First triangle of the square
            indices[quad_id * 6 + 0] = i * MESH_RES + j
            indices[quad_id * 6 + 1] = (i + 1) * MESH_RES + j
            indices[quad_id * 6 + 2] = i * MESH_RES + (j + 1)
            # Second triangle of the square
            indices[quad_id * 6 + 3] = (i + 1) * MESH_RES + j + 1
            indices[quad_id * 6 + 4] = i * MESH_RES + (j + 1)
            indices[quad_id * 6 + 5] = (i + 1) * MESH_RES + j

if __name__ == '__main__':
    window = ti.ui.Window("Brownian 2D Diffusion", (1000,1000))
    canvas = window.get_canvas()
    scene = ti.ui.Scene()
    camera = ti.ui.Camera()
    camera.position(8, 4, 0)
    camera.lookat(0,0,0)
    camera.up(0,1,0)
    init_mesh_indices()
    init_mesh_vertices()
    apply_initial_conditions()

    if BENCHMARK == True:
        import time
        print("Starting benchmark...")
        it = 0
        start = time.time()
        for i in range(100):
            update_positions()
            it += 1
        end = time.time()
        print(f"Benchmark took {(end-start)/(it*UPDATES_PER_BATCH)}s/dt")

    it = 0
    while window.running:
        camera.track_user_inputs(window, movement_speed=0.03, hold_key=ti.ui.RMB)
        scene.set_camera(camera)
        scene.ambient_light((0.8, 0.8, 0.8))
        scene.point_light(pos=(0.5, 1.5, 1.5), color=(1, 1, 1))
        # update_z()
        # t[None] += 0.01
        update_positions()

        update_mesh_vertices()
        scene.mesh(vertices, indices, per_vertex_color=colors)

        canvas.scene(scene)
        # window.save_image(f"brownian_heat_{int(N_POINTS/1000000)}e6_{MESH_RES}_{it:05d}.png")
        it+=1
        window.show()