import taichi as ti
import numpy as np

ti.init(arch=ti.cuda, default_fp=ti.f64)

from utils import Sinusoid

p = 8
MESH_RES = 2**p -1
n = MESH_RES

"""
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],
    [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 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():
    for i,j in solution:
        h = ti.cast(solution[i,j], ti.float32)
        vertices[i+MESH_RES*j].y = h
        
        level = ti.min(ti.floor(h*(colormap_field.shape[0]-1)), colormap_field.shape[0]-2)
        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)



"""
NOTE: This uses globals 'passes' and 'n' to allow static for-loop unrolling!
"""
@ti.kernel
def solve_tridiagonal(
    now: ti.template(),
    dt: float,
    ld: ti.template(), 
    dd: ti.template(), 
    ud: ti.template(), 
    rhs: ti.template(), 
    soln: ti.template(),
    orientation: int,
): 
    for _ in ti.static(range(UPDATES_PER_KERNEL)): # TODO: decide if this is actually faster or not
        """
        Copy current solution to rhs
        """
        for col, row in soln:
            # if col == 0 or row == 0 or col == n-1 or row == n-1:
            #     pass
            if  (row== 0 or  row == n-1) and orientation == 0:
                pass
            elif (col == 0 or col == n-1) and orientation == 1:
                pass
            else:
                left = 0.0
                right = 0.0
                if orientation == 0:
                    if col > 0:
                        left = s*soln[col-1, row]
                    if col < n-1:
                        right = s*soln[col+1, row]
                    rhs[0, col, row] = (1-2*s)*soln[col, row] + left + right
                else:
                    if row > 0:
                        left = s*soln[col, row-1]
                    if row < n-1:
                        right = s*soln[col, row+1]
                    rhs[0, col, row] = (1-2*s)*soln[col, row] + left + right
            soln[col, row] = 0.0
        ti.sync()
        update_boundaries(rhs, now[None], orientation)
        """
        Forward reduction
        """
        for pass_idx in ti.static(range(passes-1)):
            eqns = int(0.5**(pass_idx+1)*(n+1))-1
            # print(f"Beginning to update solutions for dataset {pass_idx+1}, which will contain {eqns} eqns once complete")
            for line, i in ti.ndrange(n, eqns): # TODO: using n here assumes square grid! that's okay tho
                a0, b0, c0, a1, b1, c1, a2, b2, c2, = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0
                if orientation == 0:
                    row0 = i*2
                    row1 = i*2+1
                    row2 = i*2+2
                    col = line
                    a0, b0, c0 = ld[pass_idx, col, row0], dd[pass_idx, col, row0], ud[pass_idx, col, row0]
                    a1, b1, c1 = ld[pass_idx, col, row1], dd[pass_idx, col, row1], ud[pass_idx, col, row1]
                    a2, b2, c2 = ld[pass_idx, col, row2], dd[pass_idx, col, row2], ud[pass_idx, col, row2]
                else:
                    col0 = i*2
                    col1 = i*2+1
                    col2 = i*2+2
                    row = line
                    a0, b0, c0 = ld[pass_idx, col0, row], dd[pass_idx, col0, row], ud[pass_idx, col0, row]
                    a1, b1, c1 = ld[pass_idx, col1, row], dd[pass_idx, col1, row], ud[pass_idx, col1, row]
                    a2, b2, c2 = ld[pass_idx, col2, row], dd[pass_idx, col2, row], ud[pass_idx, col2, row]

                beta = 1.0
                alpha = -beta*a1/b0
                gamma = -beta*c1/b2

                rhs0, rhs1, rhs2 = 0.0, 0.0, 0.0
                if orientation == 0:
                    row0 = i*2
                    row1 = i*2+1
                    row2 = i*2+2
                    col = line
                    rhs0 = alpha * rhs[pass_idx, col, row0]
                    rhs1 = beta  * rhs[pass_idx, col, row1]
                    rhs2 = gamma * rhs[pass_idx, col, row2]
                else:
                    col0 = i*2
                    col1 = i*2+1
                    col2 = i*2+2
                    row = line
                    rhs0 = alpha * rhs[pass_idx, col0, row]
                    rhs1 = beta  * rhs[pass_idx, col1, row]
                    rhs2 = gamma * rhs[pass_idx, col2, row]

                rhs_new = rhs0+rhs1+rhs2
                a = alpha*a0
                b = alpha*c0 + beta*b1 + gamma*a2
                c = gamma*c2
                if orientation == 0:
                    col = line
                    ld[pass_idx+1, col, i]  = a
                    dd[pass_idx+1, col, i]  = b
                    ud[pass_idx+1, col, i]  = c
                    rhs[pass_idx+1, col, i] = rhs_new
                else:
                    row = line
                    ld[pass_idx+1, i, row]  = a
                    dd[pass_idx+1, i, row]  = b
                    ud[pass_idx+1, i, row]  = c
                    rhs[pass_idx+1, i, row] = rhs_new
            ti.sync()
        """
        Backwards substitution
        """
        center_idx = int((n-1)/2)
        for line in range(n):
            if orientation == 0:
                col = line
                soln[col, center_idx] =    rhs[passes-1, col, 0] / dd[passes-1, col, 0]
            else:
                row = line
                soln[center_idx, row] =    rhs[passes-1, 0, row] / dd[passes-1, 0, row]
        ti.sync()
        for pass_counter in ti.static(range(passes-1)):
            pass_idx = passes-2-pass_counter
            n_eqs_in_pass = 2**(pass_counter+2)-1
            n_unknowns_in_pass = (n_eqs_in_pass+1)/2
            hop_size = 2**pass_idx
            for line, unknown_counter in ti.ndrange(n, int(n_unknowns_in_pass)):
                reduced_mat_row_idx = unknown_counter*2
                soln_idx = hop_size-1 + reduced_mat_row_idx*hop_size
                if orientation==0:
                    col = line
                    left = 0.0
                    right = 0.0
                    if soln_idx - hop_size >= 0:
                        left  = ld[pass_idx, col, reduced_mat_row_idx] * soln[col, soln_idx-hop_size]
                    if soln_idx + hop_size <= n-1:
                        right = ud[pass_idx, col, reduced_mat_row_idx] * soln[col, soln_idx+hop_size]
                    soln[col, soln_idx] = ( rhs[pass_idx, col, reduced_mat_row_idx] - left - right ) / dd[pass_idx, col, reduced_mat_row_idx]
                else:
                    row = line
                    left = 0.0
                    right = 0.0
                    if soln_idx - hop_size >= 0:
                        left  = ld[pass_idx, reduced_mat_row_idx, row] * soln[soln_idx-hop_size, row]
                    if soln_idx + hop_size <= n-1:
                        right = ud[pass_idx, reduced_mat_row_idx, row] * soln[soln_idx+hop_size, row]
                    soln[soln_idx, row] = ( rhs[pass_idx, reduced_mat_row_idx, row] - left - right ) / dd[pass_idx, reduced_mat_row_idx, row]
            ti.sync()
        now[None] += dt
        ti.sync()

@ti.kernel
def assemble_tridiagonal(ld: ti.template(), dd: ti.template(), ud: ti.template(), s: float, orientation: int):
    # TODO: we can recycle the same field!
    for col,row in ti.ndrange(n, n):
        tester = 0
        if orientation == 0:
            tester = row
        else:
            tester = col
        if tester == 0:
            ld[0, col, row] = 0 # zero pad the ld which doesn't exist
            dd[0, col, row] = 1 # BC is given in rhs
            ud[0, col, row] = 0
        elif tester == n-1:
            ld[0, col, row] = 0 
            dd[0, col, row] = 1 # BC is given in rhs
            ud[0, col, row] = 0 # zero pad the ud which doesn't exist
        else:
            ld[0, col, row] = -s
            dd[0, col, row] = 1 + 2*s
            ud[0, col, row] = -s
    ti.sync()

"""
Initial Conditions
"""
impulse_types = {
    "null": 0,
    "gaussian": 1,
    "splitting_gaussian": 2,
    "impulse": 3,
    "pulse_width": 4,
    "tri": 5,
    "fundamental": 6,
    "sine": 7,
    "harmonic": 8,
    "noise": 9,
}
            
@ti.func
def gaussian(x_norm: float, mu: float, sigma: float) -> float:
    return 1/ti.sqrt(2*np.pi*(sigma**2)) * ti.exp((-(x_norm- mu)**2)/(2*(sigma**2)))

@ti.func
def pt_impulse(soln: ti.template(), loc_x: int, loc_y: int, a: float):
    soln[loc_x, loc_y] = a



@ti.kernel
def apply_initial_conditions(soln: ti.template(), impulse_type: int):
    mu = 0
    sigma = 0.05
    for col, row in ti.ndrange(n, n):
        if impulse_type != impulse_types["impulse"]:
            x_norm = col/(n-1)
            y_norm = row/(n-1)
            if impulse_type == impulse_types["gaussian"]:
                soln[col, row] = gaussian(ti.sqrt((x_norm-0.5)**2 + (y_norm-0.5)**2), mu, sigma)
            if impulse_type == impulse_types["noise"]:
                soln[col, row] = 3*ti.random()
    ti.sync()
    if impulse_type == impulse_types["impulse"]:
        loc = int((n+1)/2) # center
        height = 2
        square_size = int(0.125*n)
        for i, j in ti.ndrange(square_size, square_size):
            pt_impulse(soln, loc-i, loc-j, height)
            pt_impulse(soln, loc+i, loc+j, height)
            pt_impulse(soln, loc-i, loc+j, height)
            pt_impulse(soln, loc+i, loc-j, height)
    ti.sync()
        
@ti.func
def update_boundaries(rhs: ti.template(), t: float, orientation: int):
    # rhs[0, 0] = daily_sine.z(t) + annual_sine.z(t)
    for i in range(n): # TODO: don't reference n here
        # set the top row to zero
        rhs[0, i, 0] = 0
        # Set the bottom row to zero
        rhs[0, i, n-1] = 0
        # set the left col to zero
        rhs[0,  0, i] = 0
        # Set the right col to zero
        rhs[0, n-1, i] = 0


energy = ti.field(dtype=float, shape=())
@ti.kernel
def total_energy(soln: ti.template()):
    energy[None] = 0
    for i in ti.grouped(soln):
        energy[None] += soln[i] / n
    ti.sync()

if __name__ == '__main__':
    import time

    """
    Taichi Init
    """
    # ti.init(arch=ti.cuda, default_fp=ti.f64, kernel_profiler=True)

    """
    Finite Difference Grid Setup
    """
    # n = 2**20 - 1 #2**9 = 512, Needs to be 2^k-1 for tridiag solver, TODO: automatically pad out numbers
    dx = 1
    D = 20#10
    L = (n-1)*dx
    dt = 0.05
    s = D* dt / (2 * (dx**2))
    IMPULSE_TYPE = impulse_types["noise"]
    now = ti.field(dtype=float, shape=())
    now[None] = 0
    UPDATES_PER_KERNEL = 1

    """
    Tridiag setup
    """
    passes = int(np.log2(n+1))
    print(f"Solving system will require {passes} forward passes and {passes} backward passes.")

    """
    Set up the fields
    """
    # TODO: use a better data structure so that the higher number of passes only need 2**(passes-pass_idx)-1 entries
    lower_diagonal_rows      = ti.field(dtype=float, shape=(passes, n, n))
    upper_diagonal_rows      = ti.field(dtype=float, shape=(passes, n, n))
    diagonal_rows            = ti.field(dtype=float, shape=(passes, n, n))
    lower_diagonal_cols      = ti.field(dtype=float, shape=(passes, n, n))
    upper_diagonal_cols      = ti.field(dtype=float, shape=(passes, n, n))
    diagonal_cols            = ti.field(dtype=float, shape=(passes, n, n))
    right_hand_side          = ti.field(dtype=float, shape=(passes, n, n))
    solution                 = ti.field(dtype=float, shape=(n, n)) 

    """
    Geometry fields
    """
    Y_VIEW_SCALE = 0.4
    Y_VIEW_OFFSET = 0.5

    """
    Setup
    """
    assemble_tridiagonal(lower_diagonal_rows, diagonal_rows, upper_diagonal_rows, s, orientation=0)
    assemble_tridiagonal(lower_diagonal_cols, diagonal_cols, upper_diagonal_cols, s, orientation=1)
    apply_initial_conditions(solution, IMPULSE_TYPE)
    init_mesh_vertices()
    init_mesh_indices()
    update_mesh_vertices()

    """
    Render setup
    """
    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.position(0, 8, 0)
    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.track_user_inputs(window, movement_speed=0.03, hold_key=ti.ui.RMB)
        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()
        scene.mesh(vertices, indices, per_vertex_color=colors)
        canvas.scene(scene)

        solve_tridiagonal(now, dt, lower_diagonal_cols, diagonal_cols, upper_diagonal_cols, right_hand_side, solution, orientation=1)
        ti.sync()
        solve_tridiagonal(now, dt, lower_diagonal_rows, diagonal_rows, upper_diagonal_rows, right_hand_side, solution, orientation=0)
        ti.sync()

        # if it % 100 == 0:
        #     total_energy(solution)
        #     ti.sync()
        #     print(f"Current energy is: {energy[None]}")
        it += 1
        window.show()

    # ti.profiler.print_scoped_profiler_info()
    # ti.profiler.print_kernel_profiler_info()  # The default mode: 'count'