import taichi as ti
import numpy as np
from utils import Sinusoid

# TODO: fix oob memory access for first and last equations

"""
Initial Conditions
"""
impulse_types = {
    "null": 0,
    "gaussian": 1,
    "splitting_gaussian": 2,
    "impulse": 3,
    "pulse_width": 4,
    "tri": 5,
    "fundamental": 6,
    "sine": 7,
    "harmonic": 8,
}
IMPULSE_MODE = impulse_types["gaussian"]
            
@ti.func
def gaussian(x_norm: float, mu: float, sigma: float) -> float:
    return 0.35/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: int, a: float):
    soln[loc] = a

"""
Diffusion Solver
#TODO: NOTE: This uses globals 'passes' and 'n' to allow static for-loop unrolling!
"""
@ti.kernel
def solve_tridiagonal(
    now: ti.template(),
    dt: float,
    pcr: ti.template(),
    soln: ti.template(),
    bcs: ti.template()
): 
    for _ in ti.static(range(UPDATES_PER_KERNEL)): # TODO: decide if this is actually faster or not
        """
        Copy current solution to RHS
        """
        for i in range(soln.shape[0]):
            pcr[0, i, 3] = soln[i] # rhs
            soln[i] = 0
        ti.sync()

        """Update BCs on RHS"""
        t = now[None]
        if bcs[0] == 0:
            pcr[0, 0, 3]   = daily_sine_l.z(t) + annual_sine_l.z(t) # rhs
        if bcs[1] == 0:
            pcr[0, n-1, 3] = daily_sine_r.z(t) + annual_sine_r.z(t) # rhs
        """
        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 i in range(eqns):
                a0, b0, c0 = pcr[pass_idx, i*2 + 0, 0], pcr[pass_idx, i*2 + 0, 1], pcr[pass_idx, i*2 + 0, 2]
                a1, b1, c1 = pcr[pass_idx, i*2 + 1, 0], pcr[pass_idx, i*2 + 1, 1], pcr[pass_idx, i*2 + 1, 2]
                a2, b2, c2 = pcr[pass_idx, i*2 + 2, 0], pcr[pass_idx, i*2 + 2, 1], pcr[pass_idx, i*2 + 2, 2]

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

                rhs0 = alpha * pcr[pass_idx, i*2 + 0, 3]
                rhs1 = beta  * pcr[pass_idx, i*2 + 1, 3]
                rhs2 = gamma * pcr[pass_idx, i*2 + 2, 3]

                rhs_new = rhs0+rhs1+rhs2
                a = alpha*a0
                b = alpha*c0 + beta*b1 + gamma*a2
                c = gamma*c2
                pcr[pass_idx+1, i, 0] = a       # LD
                pcr[pass_idx+1, i, 1] = b       # DD
                pcr[pass_idx+1, i, 2] = c       # UD
                pcr[pass_idx+1, i, 3] = rhs_new # RHS
            ti.sync()
        """
        Backwards substitution
        """
        center_idx = int((n-1)/2)
        soln[center_idx] = pcr[passes-1, 0, 3] / pcr[passes-1, 0, 1]
        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 unknown_counter in range(int(n_unknowns_in_pass)):
                reduced_mat_row_idx = unknown_counter*2
                soln_idx = hop_size-1 + reduced_mat_row_idx*hop_size
                rhs = pcr[pass_idx, reduced_mat_row_idx, 3]
                dd  = pcr[pass_idx, reduced_mat_row_idx, 1]
                ld  = 0.0
                ud  = 0.0
                if soln_idx-hop_size >= 0:
                    ld = pcr[pass_idx, reduced_mat_row_idx, 0] * soln[soln_idx-hop_size]
                if soln_idx+hop_size <= n-1:
                    ud = pcr[pass_idx, reduced_mat_row_idx, 2] * soln[soln_idx+hop_size]
                soln[soln_idx] = (rhs - ld - ud) / dd
            ti.sync()
        now[None] += dt

@ti.kernel
def assemble_tridiagonal(pcr: ti.template(), diff_field: ti.template(), bcs: ti.template(), dt: float, dx: float):
    for i in range(n):
        if i == 0:
            if bcs[0] == 0:
                """Prescriptive"""
                pcr[0, i, 0] = 0 # LD: zero pad the ld which doesn't exist
                pcr[0, i, 1] = 1 # DD: BC is given in rhs
                pcr[0, i, 2] = 0 # UD: BC is given in rhs
            else:
                """Adiabatic"""
                s = diff_field[i] * dt / (dx**2) 
                pcr[0, i, 0] = 0     # LD: non-existent
                pcr[0, i, 1] = 1 + s # DD
                pcr[0, i, 2] = -s    # UD
        elif i == n-1:
            if bcs[1] == 0:
                """Prescriptive"""
                pcr[0, i, 0] = 0 # LD: BC is given in rhs
                pcr[0, i, 1] = 1 # DD: BC is given in rhs
                pcr[0, i, 2] = 0 # UD: zero pad the ud which doesn't exist
            else:
                """Adiabatic"""
                s = diff_field[i] * dt / (dx**2) 
                pcr[0, i, 0] = -s    # LD
                pcr[0, i, 1] = 1 + s # DD
                pcr[0, i, 2] = 0     # UD: non-existent
        else:
            s = diff_field[i] * dt / (dx**2) 
            # s_outer = (diff_field[i+1] - diff_field[i-1]) * dt / (4 * (dx**2))
            s_next = diff_field[i+1] * dt / (dx**2) 
            """Centered"""
            # ld[0, i] = -(s-s_outer)
            # dd[0, i] = 1 + 2*s
            # ud[0, i] = -(s+s_outer)
            """Forward"""
            pcr[0, i, 0] = -s             # LD
            pcr[0, i, 1] = 1 + s + s_next # DD
            pcr[0, i, 2] = -s_next        # UD
    ti.sync()

    """Activate the SNode memory cells"""
    # for sys_id,eq_id,el_id in ti.ndrange(passes, n, 4):
    #     if sys_id > 0 and eq_id < 2**(passes - sys_id) - 1:
    #         pcr[sys_id, eq_id, el_id] = 0
    ti.sync()


@ti.kernel
def fill_diffusivity(D: float, diff_field: ti.template(), mat_field: ti.types.ndarray()):
    for i in diff_field:
        x_norm = i / (n-1)
        for j in range(mat_field.shape[0]):
            if x_norm <= mat_field[j, 0]:
                diff_field[i] = mat_field[j,1]*D
                break
    ti.sync()
    # ti.loop_config(serialize=True)
    # for i in range(1, diff_field.shape[0]):
    #     if diff_field[i] != diff_field[i-1]:
    #         diff_field[i] = (diff_field[i]+diff_field[i-1])/2
         

@ti.kernel
def apply_initial_conditions(soln: ti.template(), impulse_type: int, offset: float):
    mu = 0.5
    sigma = 0.10
    for i in range(n):
        if impulse_type != impulse_types["impulse"]:
            x_norm = i/(n-1)
            if impulse_type == impulse_types["gaussian"]:
                soln[i] = gaussian(x_norm, mu, sigma)
    ti.sync()
    if impulse_type == impulse_types["impulse"]:
        loc = int((n+1)/2) # center
        height = 0.35
        pt_impulse(soln, loc, height)
    ti.sync()
    for i in range(n):
        soln[i] += offset
    ti.sync()
        
"""
Geo/Rendering
"""
@ti.kernel
def update_geo(soln: ti.template(), circ_pts: ti.template(), line_pts: ti.template()):
    for i in range(n):
        p_x = i / (n-1)
        p_y = (soln[i] - Y_VIEW_OFFSET) / Y_VIEW_SCALE 

        circ_pts[i].x  = p_x
        circ_pts[i].y  = p_y 

        if (i != n-1):
            line_pts[2*i].x = p_x
            line_pts[2*i+1].x =  (i+1) / (n-1)
            line_pts[2*i].y = p_y
            line_pts[2*i+1].y = (soln[i+1] - Y_VIEW_OFFSET) / Y_VIEW_SCALE

@ti.kernel
def update_im(soln: ti.template(), im: ti.template()):
    for i,j in im:
        i_norm = i / 999
        sol_i = int((n-1)*i_norm)
        sol_norm = (soln[sol_i]-Y_VIEW_OFFSET)/Y_VIEW_SCALE
        im[i,j].r = sol_norm
        im[i,j].g = 0.5
        im[i,j].b = 1-sol_norm

@ti.kernel
def init_vert_lines(div_pts: ti.template(), mat_field: ti.types.ndarray()):
    for i in range(mat_field.shape[0]):
        div_pts[2*i].y = -10
        div_pts[2*i+1].y = 10
        div_pts[2*i].x = mat_field[i,0]
        div_pts[2*i+1].x = mat_field[i,0]


if __name__ == '__main__':
    import time

    """
    Taichi Init
    """
    # TODO: fix oob error on edges for final pass
    ti.init(arch=ti.cuda, default_fp=ti.f64)
    # ti.init(arch=ti.cuda, default_fp=ti.f64, kernel_profiler=True)

    """
    Material Defs
    K [W/mK], Cp [J/kgK], rho [kg/m3]

    J s^-1 m^-1 K^-1 J^-1 kg K kg^-1 m3
    J J^-1  K^-1 K kg kg^-1  m3 m^-1 s^-1 
    m2 m s^-1  
    """
    print(f"--- Material Library ---")
    mat_names = [
        "brick",
        "concrete-low-rho",
        "concrete-high-rho",
        "granite",
        "glass",
        "hardwood",
        "softwood",
        "pvc",
        "paper",
        "act",
        "particle-board-low-rho",
        "particle-board-high-rho",
        "gypsum",
        "fiberglass",
        "xps",
        "air"
    ]
    mat_defs = np.array([
        [0.70,  840, 1600], # Brick
        [0.40, 1000, 1200], # Concrete (light cast)
        [1.40,  840, 2100], # Concrete (dense cast)
        [2.50,  820, 2600], # Granite
        [0.80,  880, 2700], # Glass
        [0.16, 1250,  720], # Hardwood (oak)
        [0.12, 1350,  510], # Softwood (pine)
        [0.15, 1250, 1400], # PVC
        [0.04, 1300,  930], # Paper
        [0.06, 1340,  290], # Acoustic Tile
        [0.08, 1300,  590], # Particle Board (low density)
        [0.17, 1300, 1000], # Particle Board (high density)
        [0.26, 1089,  711], # Gypsum
        [0.04,  700,  150], # Fiberglass
        [0.03,  1200,  50], # XPS
        [0.025, 1000, 1.2], # Air (ambient)
    ])
    mat_diffs = mat_defs[:,0]/(mat_defs[:,1]*mat_defs[:,2])
    mat_ids = {
        name: i for i,name in enumerate(mat_names)
    }
    for name,mat_id in mat_ids.items():
        print(f"\nMaterial: {name}")
        print(f"k:   {mat_defs[mat_id, 0]:0.2f} [W/mK]")
        print(f"Cp:  { int(mat_defs[mat_id, 1]):04d} [J/kgK]")
        print(f"rho: { int(mat_defs[mat_id, 2]):04d} [kg/m3]")
        print(f"D:   {  mat_diffs[mat_id]:0.3e} [m2/s]")

    """
    Wall Definition
    """
    print(f"\n\n--- Assembly Definition ---")
    assembly_def = np.array([
        [0.050, mat_ids["brick"]             ], 
        [0.160, mat_ids["air"]               ], 
        [0.250, mat_ids["concrete-high-rho"] ], 
        [0.100, mat_ids["xps"]               ], 
        [0.020, mat_ids["gypsum"]            ]
    ])
    assembly_length = np.sum(assembly_def[:,0])
    material_ends   = np.cumsum(assembly_def[:,0]) / assembly_length
    assembly        = np.zeros(shape=assembly_def.shape)
    assembly[:,0]   = material_ends
    assembly[:,1]   = mat_diffs[assembly_def[:,1].astype(np.uint8)]
    materials       = ti.ndarray(dtype=float, shape=assembly_def.shape)
    materials.from_numpy(assembly)

    for i,(length, mat_id) in enumerate(assembly_def):
        print(f"\nMatId: {mat_id}")
        print(f"Len: {length}")
        print(f"End: {assembly[i, 0]}")
        print(f"D: {assembly[i, 1]:0.3e}")

    """
    Finite Difference Grid Setup
    """
    print(f"\n\n--- FD Scheme Config ---")
    p = 14
    n = 2**p - 1 #2**9 = 512, Needs to be 2^k-1 for tridiag solver, TODO: automatically pad out numbers
    L = assembly_length
    # L = 0.75 # [m]
    dx = L/(n-1) # [m]
    # dx = 1
    # L = (n-1)*dx
    D = 1 # Diffusivity scalar
    dt = 5*60 # s
    IMPULSE_TYPE = impulse_types["null"]
    now = ti.field(dtype=float, shape=())
    now[None] = 0
    UPDATES_PER_KERNEL = 1

    print(f"dx: {dx*1000:0.3f} [mm]")
    print(f"dt: {dt:0.3f} [s]")
    print(f"D_max dt / dx**2: {np.min(mat_diffs) * dt / (dx**2)}")

    """
    Tridiag setup
    """
    print(f"\n\n--- Tridiag Setup ---")
    passes = p#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
    diffusivity         = ti.field(dtype=float, shape=(n,))

    # Spatially sparse data structure
    cyclic_solver = ti.field(ti.f64)
    systems = ti.root.dense(ti.i, p)
    equations = systems.pointer(ti.j, n)
    coeffs = equations.dense(ti.k, 4) # LD: 0, DD: 1, UD: 2, RHS: 3
    coeffs.place(cyclic_solver)

    solution            = ti.field(dtype=float, shape=(n,)) 
    boundary_conditions = ti.field(dtype=int, shape=(2))
    # 0 = prescriptive, 1 = adiabatic
    boundary_conditions[0] = 0 
    boundary_conditions[1] = 1 

    """
    Geometry fields
    """
    line_points = ti.Vector.field(2,dtype=ti.f32, shape=(2*(n-1),))
    divider_points = ti.Vector.field(2,dtype=ti.f32, shape=(2*(materials.shape[0])))
    circ_points = ti.Vector.field(2,dtype=ti.f32, shape=(n,))
    Y_VIEW_SCALE = 40
    Y_VIEW_OFFSET = 270

    """
    Setup
    """
    fill_diffusivity(D, diffusivity, materials)
    assemble_tridiagonal(cyclic_solver, diffusivity, boundary_conditions, dt, dx)
    apply_initial_conditions(solution, IMPULSE_TYPE, offset=273)
    update_geo(solution, circ_points, line_points)
    daily_sine_l = Sinusoid(
        amp=3.5, 
        freq=1/(24*60*60),
        bias=0
    )
    annual_sine_l = Sinusoid(
        amp=13,#15, 
        freq=1/(365*24*60*60),
        bias=288
    )
    daily_sine_r = Sinusoid(
        amp=0, 
        freq=1/(24*60*60),
        bias=0
    )
    annual_sine_r = Sinusoid(
        amp=0, 
        freq=1/(365*24*60*60),
        bias=285
    )

    """
    Benchmark
    """
    # it = 0
    # solve_tridiagonal(now, dt, cyclic_solver, solution, boundary_conditions)
    # print("Starting benchmark!")
    # start = time.time()
    # while it<1000:
    #     solve_tridiagonal(now, dt, cyclic_solver, solution, boundary_conditions)
    #     it += 1
    # end = time.time()
    # print(f"{1000*(end-start)/(it*UPDATES_PER_KERNEL):0.3f}ms")

    # ti.profiler.memory_profiler.print_memory_profiler_info()
    # raise RuntimeError("STOP")

    """
    Render setup
    """
    window = ti.ui.Window('Diffusion 1D Implicit (Parallel)', res = (1000,1000), pos = (150, 150))
    canvas = window.get_canvas()
    image = ti.Vector.field(3,dtype=ti.f32, shape=(1000,1000))
    init_vert_lines(divider_points, materials)

    it = 0
    start = time.time()
    while window.running:
        update_im(solution, image)
        canvas.set_image(image)
        canvas.lines(line_points, 0.001, indices=None, color=(1,1,1))
        canvas.lines(divider_points, 0.001, indices=None, color=(0.5,1,0.5))
        update_geo(solution, circ_points, line_points)
        solve_tridiagonal(now, dt, cyclic_solver, solution, boundary_conditions)
        # it += 1
        # if it % (24*60*60 / dt) == 0:
        #     print(f"Completed day {int(it*dt / (24*60*60)-1):03d}")
        # if it == int((365 * 24 * 60 * 60) / dt):
        #     end = time.time()
        #     print(f"1 year took {end-start:0.3f}s ({(end-start)/it*1000:0.3f}ms/it)")
        window.show()

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