import taichi as ti
import numpy as np


"""
You will need three arrays/fields, one for the lowr diagonal, one for the upper, and one for the main
The lower and upper diagonals should be padded with zeros at the start/end respectively
nb: for cyclic reduction to work, the matrix must have n=2^w-1 for some integer w; pad it out with trivial equations if needed
"""
# Pad the lower diag and upper diag with 0s at the start and end respectively
ld  = [0,  2, -1, 3,  8,  5, 1]
ud  = [3,  5, -1, 2, -2, -5, 0]
dd  = [2,  4, -2, 1, -3,  1, 2]



"""
NOTE: This uses globals 'passes' and 'n' to allow static for-loop unrolling!
"""
@ti.kernel
def solve_tridiagonal(
    ld: ti.template(), 
    dd: ti.template(), 
    ud: ti.template(), 
    rhs: ti.template(), 
    soln: ti.template(),
): 
    """
    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 = ld[pass_idx, i*2],     dd[pass_idx, i*2],     ud[pass_idx, i*2]
            a1, b1, c1 = ld[pass_idx, i*2 + 1], dd[pass_idx, i*2 + 1], ud[pass_idx, i*2 + 1]
            a2, b2, c2 = ld[pass_idx, i*2 + 2], dd[pass_idx, i*2 + 2], ud[pass_idx, i*2 + 2]

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

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

            rhs_new = rhs0+rhs1+rhs2
            a = alpha*a0
            b = alpha*c0 + beta*b1 + gamma*a2
            c = gamma*c2
            ld[pass_idx+1, i]  = a
            dd[pass_idx+1, i]        = b
            ud[pass_idx+1, i]  = c
            rhs[pass_idx+1, i] = rhs_new
        ti.sync()
    """
    Backwards substitution
    """
    center_idx = int((n-1)/2)
    soln[center_idx] =    rhs[passes-1,0] / dd[passes-1,0]
    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
            # TODO: Make sure that the zero padding on ld/ud successfully allows the following statement which skips conditional format for edges
            soln[soln_idx] = (
                rhs[pass_idx, reduced_mat_row_idx] 
                - ld[pass_idx, reduced_mat_row_idx]*soln[soln_idx-hop_size]
                - ud[pass_idx, reduced_mat_row_idx]*soln[soln_idx+hop_size]
                ) / dd[pass_idx,reduced_mat_row_idx]
        ti.sync()

@ti.kernel
def init_fields(
    ld: ti.template(), 
    dd: ti.template(), 
    ud: ti.template(), 
    rhs: ti.template(), 
):
    ud_v = 0.0
    ld_v = 0.0
    dd_v = 0.0
    for i in range(n):
        ld_v = 0 if i == 0 else ti.random()+1 # ld[i], 1
        ud_v = 0 if i == n-1 else ti.random()+1 #ud[i], 1
        dd_v = ti.random() + 3 #dd[i], D
        ld[0, i] = ld_v
        ud[0, i] = ud_v
        dd[0, i] = dd_v
        rhs[0,i] = (i-1)*ld_v + i*dd_v + (i+1)*ud_v #induces an answer of 0,1,2,3,...
        # if i == 0: # For solving a column of the inverse, set rhs to [1,0,...,0] etc
        #     rhs[0, i] = 1
        # else:
        #     rhs[0, i] = 0
        for j in ti.static(range(1, dd.shape[0], 1)):
            ld[j, i] = 0
            ud[j, i] = 0
            dd[j, i]       = 0
            rhs[j,i] = 0
    ti.sync()


if __name__ == '__main__':

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

    """
    We are using a randomly generated matrix, so let's just set n accordingly
    """
    n = 2**20-1
    passes = int(np.log2(n+1))
    print(f"Solving system will require {passes} forward passes and {passes} backward passes.")

    """
    Set up the fields
    """
    # Track solution reduction history through multidimensional field
    # TODO: use a better data structure so that the higher number of passes only need 2**(passes-pass_idx)-1 entries
    lower_diagonal      = ti.field(dtype=float, shape=(passes, n))
    upper_diagonal      = ti.field(dtype=float, shape=(passes, n))
    diagonal            = ti.field(dtype=float, shape=(passes, n))
    right_hand_side     = ti.field(dtype=float, shape=(passes, n))
    solution            = ti.field(dtype=float, shape=(n,)) 

    expected_answer = np.linspace(0,n-1,n)
    for i in range(20):
        init_fields(lower_diagonal, diagonal, upper_diagonal, right_hand_side)
        solve_tridiagonal(lower_diagonal, diagonal, upper_diagonal, right_hand_side, solution)
        print(f"MSE per index: {np.sum((expected_answer-solution.to_numpy())**2)/n}")
    # for i in range(n):
    #     print(solution[i])

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