import taichi as ti
import numpy as np

ti.init(arch=ti.cpu, default_fp=ti.f32)

p = 9
MESH_RES = 2**p
n = MESH_RES
u = ti.field(dtype=float, shape=(n, n)) 
alpha = 1.9

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

colormap = [
    [64,57,144],
    [100,100,100],
    [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 init_mesh_vertices():
    for i,j in ti.ndrange(n,n):
        vertices[i+MESH_RES*j].x = ti.cast(i/(MESH_RES-1)*mesh_render_size-mesh_render_size/2, ti.float32)
        vertices[i+MESH_RES*j].y = ti.cast(0.0, ti.float32)
        vertices[i+MESH_RES*j].z = ti.cast(j/(MESH_RES-1)*mesh_render_size-mesh_render_size/2, ti.float32)

@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

@ti.kernel
def update_mesh_vertices(height_map: ti.template(), colormap_field: ti.template(), vertices: ti.template(), colors: ti.template()):
    h_min = -1
    h_max = 1
    h_range = h_max - h_min
    for i,j in height_map:
        h = ti.cast((height_map[i,j]-h_min)/h_range, ti.float32)
        vertices[i+MESH_RES*j].y = ti.cast(height_map[i,j], ti.float32)
        
        level = ti.max(ti.min(ti.floor(h*(colormap_field.shape[0]-1)), colormap_field.shape[0]-2), 0)
        colorphase = ti.cast(ti.min(ti.max(h*(colormap_field.shape[0]-1) - level, 0),1), ti.float32)
        level_idx = ti.cast(level, dtype=int)

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



@ti.kernel
def set_bcs(soln: ti.template()):
    """
    BCs
    """
    # TODO: move outside of loop
    for i in range(n):
        soln[i, 0] = 0
        soln[0, i] = 0
        soln[n-1, i] = -1
        soln[i, n-1] = 1
    ti.sync()

@ti.kernel
def update_sor(soln: ti.template()):
    """TODO: better way to prevent parallelization without ti.static loop unrolling"""
    ti.loop_config(serialize=True)
    for i in range(n-2):
        row = i+1
        for j in range(n-2):
            col = j+1
            soln[col, row] = (1-alpha)*soln[col, row] + alpha / 4 * (soln[col, row+1] + soln[col+1, row] + soln[col-1, row] + soln[col, row-1])

if __name__ == "__main__":
    set_bcs(u)
    init_mesh_indices()
    init_mesh_vertices()
    """
    Render setup
    """
    window = ti.ui.Window("SOR 2D Diffusion", (1000,1000))
    canvas = window.get_canvas()
    scene = ti.ui.Scene()
    camera = ti.ui.Camera()
    camera.lookat(0,0,0)
    camera.up(0,1,0)
    camera_radius = 7
    camera_height = 4.5
    camera_speed = 0.001

    camera_t = 0
    it = 0
    while window.running:
        camera_t += camera_speed
        camera.position(camera_radius*ti.sin(camera_t), camera_height, camera_radius*ti.cos(camera_t))
        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_mesh_vertices(u, colormap_field, vertices, colors)
        scene.mesh(vertices, indices, per_vertex_color=colors)
        canvas.scene(scene)

        update_sor(u)

        it += 1
        window.show()