import time

import numpy as np
import taichi as ti

@ti.data_oriented
class Solver:
    def __init__(self, dt, dx, n, D, boundary_values, colormap, u_max, updates_per_batch) -> None:
        self.mesh_render_size = 5
        self.u_max = u_max
        self.n  = n
        self.dx = dx
        self.L  = dx * (n - 1)
        self.dt = dt
        self.D  = ti.field(dtype=float, shape=(n,n))
        self.c  = self.dt / self.dx**2
        self.populate_D(D)
        self.boundaries = ti.field(dtype=float, shape=(2,2))
        self.updates_per_batch = updates_per_batch
        for col in range(2):
            for row in range(2):
                self.boundaries[col, row] = np.array(boundary_values)[col, row]*self.u_max
        self.u = ti.field(dtype=float, shape=(n,n))
        self.q = ti.field(dtype=float, shape=(n,n))
        self.u_next = ti.field(dtype=float, shape=(n,n))

        # rendering
        self.colormap_field = ti.Vector.field(3, dtype=ti.f32, shape=len(colormap))
        for i, color in enumerate(colormap):
            self.colormap_field[i] = ti.Vector(color)

        self.vertices = ti.Vector.field(3, dtype=float, shape=(n*n))
        self.colors   = ti.Vector.field(3, dtype=float, shape=(n*n))
        self.indices  = ti.field(dtype=int, shape=(3*2*(self.n-1)**2))

    @ti.kernel
    def populate_D(self, D: float):
        self.D.fill(D)
        # # Crucifix thing
        # for col, row in self.D:
        #     if col > 5/7 * self.n and col < 6/7*self.n and row > 1/8*self.n and row < 7/8*self.n:
        #         self.D[col, row] = D / 500.0
        #     if col > 3/7 * self.n and col < 4/7*self.n and row > 1/8*self.n and row < 7/8*self.n:
        #         self.D[col, row] = D / 5000.0
        #     if col > 1/7 * self.n and col < 2/7*self.n and row > 1/8*self.n and row < 7/8*self.n:
        #         self.D[col, row] = D / 500.0
        #     if col > 2/7 * self.n and col < 3/7*self.n and row > 3/7*self.n and row < 4/7*self.n:
        #         self.D[col, row] = D / 5000.0
        #     if col > 4/7 * self.n and col < 5/7*self.n and row > 3/7*self.n and row < 4/7*self.n:
        #         self.D[col, row] = D / 5000.0
        for col, row in self.D:
            if col < 3/8 * self.n:
                self.D[col, row] = 0.5*D
            elif col < 4/8 * self.n:
                self.D[col, row] = 0.1*D
            elif col < 5/8 * self.n:
                self.D[col, row] = 0.01*D
            else:
                self.D[col, row] = 0.5*D
            if col > 3/8*self.n and col < 5/8*self.n and row > 10/21*self.n-2 and row < 11/21*self.n + 2:
                self.D[col, row] = D

    @ti.kernel
    def check_explicit_cfl(self):
        for i in ti.grouped(self.D):
            if self.D[i]*self.c > 0.25:
                print(f"Warning! CFL is quite high! {self.D[i]*self.c}")
            assert self.D[i]*self.c < 0.25, "The Courant condition is not satisified!"
            

    def explicit_batch(self):
        for i in range(self.updates_per_batch):
            self.explicit_step()

    @ti.kernel
    def explicit_step(self):
        """Step"""
        for col, row in self.u:
            if col == 0 or row == 0 or col == self.n-1 or row == self.n-1:
                self.handle_boundary_explicit(col, row)
            else:
                self.handle_internal_explicit(col, row)
        ti.sync()

        """Shift"""
        for node in ti.grouped(self.u_next):
            self.u[node] = self.u_next[node]
        ti.sync()

        """Compute U"""
        self.q.fill(0.0)
        for col, row in self.u:
            hor = 0.0
            ver = 0.0
            if col > 0:
                hor += (self.u[col, row] - self.u[col - 1, row]) * self.D[col, row]
            if col < self.n - 1:
                hor += (self.u[col + 1, row] - self.u[col, row]) * self.D[col, row]
            if row > 0:
                ver += (self.u[col, row] - self.u[col, row - 1]) * self.D[col, row]
            if row < self.n - 1:
                ver += (self.u[col, row + 1] - self.u[col, row]) * self.D[col, row]
            self.q[col, row] = ti.sqrt(0.25*(hor**2 + ver**2))*400000.0-0.05

    @ti.func
    def handle_boundary_explicit(self, col, row):
        if col == 0 or col == self.n - 1:
            boundary = self.boundaries[0, int(col / (self.n-1))]
            if boundary >= 0: # TODO: create a separate flag
                """Prescriptive"""
                self.u_next[col, row] = boundary
            else:
                """Adiabatic"""
                mult = 1.0
                alpha = self.D[col, row]*self.c
                right = 0.0
                left = 0.0
                up = 0.0
                down = 0.0
                if col > 0:
                    left  = self.u[col - 1, row]
                else:
                    right = self.u[col + 1, row]
                if row > 0:
                    down = self.u[col, row - 1]
                    mult = mult + 1
                if row < self.n-1:
                    up = self.u[col, row + 1]
                    mult = mult = mult + 1
                self.u_next[col, row] = (1 - mult*alpha) * self.u[col, row] + alpha * (left + right + up + down)
        else:
            boundary = self.boundaries[1, int(row / (self.n-1))]
            if boundary >= 0: # TODO: create a separate flag
                """Prescriptive"""
                self.u_next[col, row] = boundary
            else:
                """Adiabatic"""
                mult = 1.0
                alpha = self.D[col, row]*self.c
                right = 0.0
                left = 0.0
                up = 0.0
                down = 0.0
                if row > 0:
                    down = self.u[col, row - 1]
                else:
                    up = self.u[col, row + 1]
                if col > 0:
                    left = self.u[col - 1, row]
                    mult = mult + 1
                if col < self.n-1:
                    right = self.u[col + 1, row]
                    mult = mult = mult + 1
                self.u_next[col, row] = (1 - mult*alpha) * self.u[col, row] + alpha * (left + right + up + down)
    @ti.func 
    def handle_internal_explicit(self, col, row):
        D = self.D[col, row]
        left  = self.u[col - 1, row]
        right = self.u[col + 1, row]
        down  = self.u[col, row - 1]
        up    = self.u[col, row + 1]
        # alpha = self.D[col, row]*self.c
        # TODO: is the /2 necessary?
        alpha_l = self.c*(self.D[col - 1, row] + D)/2
        alpha_r = self.c*(self.D[col + 1, row] + D)/2
        alpha_d = self.c*(self.D[col, row - 1] + D)/2
        alpha_u = self.c*(self.D[col, row + 1] + D)/2
        alpha = alpha_l + alpha_r + alpha_d + alpha_u


        # self.u_next[col, row] = (1 - 4*alpha) * self.u[col, row] + alpha * (left + right + up + down)
        self.u_next[col, row] = (1 - alpha) * self.u[col, row] + alpha_l*left + alpha_r*right + alpha_d*down + alpha_u*up
        
    def benchmark_explicit(self, n_tests):
        print("Starting benchmark...")
        solver.explicit_step() # force compilation
        s = time.time()
        for _ in range(n_tests):
            solver.explicit_batch()
        e = time.time()
        b_time = e - s
        step_time = b_time / (n_tests * self.updates_per_batch)
        node_time = b_time / (n_tests * self.updates_per_batch * self.n)
        nodes_per_s = 1/node_time
        steps_in_year = 365*24*60*60/self.dt
        steps_in_week = 7*24*60*60/self.dt
        steps_in_day = 24*60*60/self.dt
        steps_in_hour = 60*60/self.dt
        hour_time = step_time*steps_in_hour
        day_time = step_time*steps_in_day
        week_time = step_time*steps_in_week
        year_time = step_time*steps_in_year
        print(f"\n--- Benchmark Results ---")
        print(f"dt: {self.dt:0.3f}s")
        print(f"dx: {self.dx*100:0.1f}cm")
        print(f"ar: {self.L*self.L}m2")
        print(f"{int(step_time*1e6)}ns/timestep")
        print(f"{nodes_per_s/1e6:0.3f}m nodes/s")
        print(f"{hour_time}s/hour")
        print(f"{day_time}s/day")
        print(f"{week_time}s/week")
        print(f"{year_time/60}min/year")






    @ti.kernel
    def init_mesh_vertices(self):
        for i,j in ti.ndrange(self.n, self.n):
            self.vertices[i + self.n*j].x = ti.cast(i/(self.n - 1)*self.mesh_render_size - self.mesh_render_size/2, ti.f32)
            self.vertices[i + self.n*j].y = ti.cast(0.0, ti.float32)
            self.vertices[i + self.n*j].z = ti.cast(j/(self.n - 1)*self.mesh_render_size - self.mesh_render_size/2, ti.f32)

    @ti.kernel
    def init_mesh_indices(self):
        for i, j in ti.ndrange(self.n-1, self.n-1):
                quad_id = (i * (self.n- 1)) + j
                # First triangle of the square
                self.indices[quad_id * 6 + 0] = i * self.n + j
                self.indices[quad_id * 6 + 1] = (i + 1) * self.n + j
                self.indices[quad_id * 6 + 2] = i * self.n + (j + 1)
                # Second triangle of the square
                self.indices[quad_id * 6 + 3] = (i + 1) * self.n + j + 1
                self.indices[quad_id * 6 + 4] = i * self.n + (j + 1)
                self.indices[quad_id * 6 + 5] = (i + 1) * self.n + j

    @ti.kernel
    def update_mesh_vertices(self):
        for i,j in self.u:
            h = ti.cast(self.u[i,j], ti.float32)/self.u_max
            h = ti.abs(ti.cast(self.q[i,j], ti.float32)/4)
            # self.vertices[i+self.n*j].y = h
            
            level = ti.max(ti.min(ti.floor(h*(self.colormap_field.shape[0]-1)),self.colormap_field.shape[0]-2), 0)
            colorphase = ti.cast(ti.min(ti.max(h*(self.colormap_field.shape[0]-1) - level, 0),1), ti.f32)
            level_idx = ti.cast(level, dtype=int)

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





if __name__ == '__main__':
    ti.init(arch=ti.cuda, default_fp=ti.f32)
    D = 0.00001 # [m2/s]
    dx = 0.01 # [m]
    dt = dx**2 / (4*D) # [s]
    p = 8
    n = 2**p

    boundary_values = [
        [1, 0],
        [-1, -1],
    ]

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

    u_max = 100

    solver = Solver(dt, dx, n, D, boundary_values, colormap, u_max, updates_per_batch=100)
    solver.check_explicit_cfl()
    solver.init_mesh_indices()
    solver.init_mesh_vertices()
    # solver.benchmark_explicit(n_tests=100)

    """
    Render setup
    """
    window = ti.ui.Window("2D Diffusion", (1000,1000))
    canvas = window.get_canvas()
    scene = ti.ui.Scene()
    camera = ti.ui.Camera()
    camera.position(0, 8, 0)
    camera.lookat(0,0,0)
    camera.up(0,1,0)
    camera.up(1,0,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, 2, 0), color=(1, 1, 1))

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

        solver.explicit_batch()
        # window.save_image(f"./week_5_fd_pde/images_4/{it:05d}.png")
        it += 1
        window.show()