from sympy import sqrt, symbols, factorial, IndexedBase, Symbol

p = IndexedBase("p")
h = Symbol("h")

def taylor(a, b, c):
    n = 4
    derivatives = [0] * n
    steps = [0] * n

    result = p

    for order in range(1, n+1):
        for i in range(0, 3**order):
            for j in range(0, order):
                derivatives[j] = int((i / (3**(order-1-j))) % 3 + 1)
                step = [a, b, c][derivatives[j] - 1]
                steps[j] = step

            total_step = 1
            symbol_name = ""
            for j in range(0, order):
                total_step *= (steps[j] * h)
                component = derivatives[j] - 1
                symbol_name += ["x", "y", "z"][component]

            result += (1 / factorial(order)) * total_step * Symbol(symbol_name)
    return result;

def expand_laplacian(face, corner):
    # Note: Any linear combination of the face and corner coefficients will
    # produce a valid discretized Laplacian, except for when normalization == 0
    center = face * 4 + corner * 4
    normalization = face + 2 * corner

    expr = (
        - center * taylor(0, 0, 0)

        + face * taylor( 1,  0,  0)
        + face * taylor(-1,  0,  0)
        + face * taylor( 0, -1,  0)
        + face * taylor( 0,  1,  0)

        +  corner * taylor( 1,  1,  0)
        +  corner * taylor( 1, -1,  0)
        +  corner * taylor(-1,  1,  0)
        +  corner * taylor(-1, -1,  0)

    ) / (normalization*h*h)

    print("Laplacian:")
    print("\tCenter:        {:+d}".format(-center))
    print("\tFace:          {:+d}".format(face))
    print("\tCorner:        {:+d}".format(corner))
    print("\tNormalization: {:+d}".format(normalization))
    print("\tExpansion:     {}".format(expr.simplify()))
    print()

    return expr

#a = expand_laplacian(1, 0)
#b = expand_laplacian(0, 1)
#c = expand_laplacian(1, 1)
#d = expand_laplacian(2, 1)
#e = expand_laplacian(1, 2)

expr_colocated = (
    -4*taylor(0, 0, 0)
    +1*taylor(-2, 0, 0)
    +1*taylor(+2, 0, 0)
    +1*taylor(0, -2, 0)
    +1*taylor(0, +2, 0)
) / (4*h*h)

expr_wide = (
    -8*taylor(0, 0, 0)
    +1*taylor(-1, 0, 0)
    +1*taylor(+1, 0, 0)
    +1*taylor(0, -1, 0)
    +1*taylor(0, +1, 0)
    +1*taylor(-1,-1, 0)
    +1*taylor(+1,+1, 0)
    +1*taylor(+1, -1, 0)
    +1*taylor(-1, +1, 0)
) / (3*h*h)

expr_vertex = (
    -4*taylor(0, 0, 0)
    +1*taylor(-1,-1, 0)
    +1*taylor(+1,+1, 0)
    +1*taylor(+1, -1, 0)
    +1*taylor(-1, +1, 0)
) / (2*h*h)

expr_mac = (
    -4*taylor(0, 0, 0)
    +1*taylor(-1, 0, 0)
    +1*taylor(+1, 0, 0)
    +1*taylor(0, -1, 0)
    +1*taylor(0, +1, 0)
) / (h*h)

print((expr_colocated - expr_mac).simplify())
print((expr_colocated - expr_wide).simplify())
print((expr_colocated - expr_vertex).simplify())