import numpy as np
import matplotlib.pyplot as plt

PI = np.pi

np.set_printoptions(precision=2, floatmode="maxprec", suppress=True)
figure, axes = plt.subplots()
cb = None


Re = 100
nu = 1 / Re
domain_size = (1, 2)
step = 0.05
N = int(domain_size[0] / step)
M = int(domain_size[1] / step)


shape = (N, M)
alpha = 1  # Coefficient

t_m = 1
h_c = 0.3
l_c = 0.6

# Staggered vars
u = np.zeros(shape=shape, dtype=float)
u_star = np.zeros(shape=shape, dtype=float)
u_new = np.zeros(shape=shape, dtype=float)

v = np.zeros(shape=shape, dtype=float)
v_star = np.zeros(shape=shape, dtype=float)
v_new = np.zeros(shape=shape, dtype=float)

p = np.zeros(shape=shape, dtype=float)
p_star = np.zeros(shape=shape, dtype=float)
p_new = np.zeros(shape=shape, dtype=float)

d_e = np.zeros(shape=shape, dtype=float)
d_n = np.zeros(shape=shape, dtype=float)
b = np.zeros(shape=shape, dtype=float)


def u_boundary(t, y):
    def f(t):
        return 1 if t_m < t else 0.5 * (np.sin(0.5 * PI * (2 * t / t_m - 1)) + 1)
    return 6 * f(t) * (y - h_c) * (1 - y) / (1 - h_c)**2


# Loop
error = 1
precision = 10 ** -7

t = 0.35

iteration = 0
while error > precision:
    iteration += 1
    # Inflow boundary condition
    for i in range(N):
        y = domain_size[0] - i * step
        u_star[i][0] = u_boundary(t, y)
        v_star[i][0] = 0

    # Sides boundary conditions
    for j in range(M):
        v_star[0][j] = 0
        v_star[N - 1][j] = 0
        u_star[0][j] = 0
        u_star[N - 1][j] = 0

    # x-momentum
    for i in range(1, N - 1):
        for j in range(1, M - 1):
            u_E = 0.5 * (u[i][j] + u[i][j + 1])
            u_W = 0.5 * (u[i][j] + u[i][j - 1])
            v_N = 0.5 * (v[i - 1][j] + v[i - 1][j + 1])
            v_S = 0.5 * (v[i][j] + v[i][j + 1])

            a_E = -0.5 * u_E * step + nu
            a_W = +0.5 * u_W * step + nu
            a_N = -0.5 * v_N * step + nu
            a_S = +0.5 * v_S * step + nu

            a_e = 0.5 * step * (u_E - u_W + v_N - v_S) + 4 * nu
            A_e = -step

            d_e[i][j] = A_e / a_e

            u_star[i][j] = (a_E * u[i][j + 1] + a_W * u[i][j - 1] + a_N * u[i - 1][j] + a_S * u[i + 1][j]) / a_e
            + d_e[i][j] * (p[i][j + 1] - p[i][j])

    # y-momentum
    for i in range(1, N - 1):
        for j in range(1, M - 1):
            u_E = 0.5 * (u[i][j] + u[i + 1][j])
            u_W = 0.5 * (u[i][j - 1] + u[i + 1][j - 1])
            v_N = 0.5 * (v[i - 1][j] + v[i][j])
            v_S = 0.5 * (v[i][j] + v[i + 1][j])

            a_E = -0.5 * u_E * step + nu
            a_W = +0.5 * u_W * step + nu
            a_N = -0.5 * v_N * step + nu
            a_S = +0.5 * v_S * step + nu

            a_n = 0.5 * step * (u_E - u_W + v_N - v_S) + 4 * nu
            A_n = -step

            d_n[i][j] = A_n / a_n

            v_star[i][j] = (a_E * v[i][j + 1] + a_W * v[i][j - 1] + a_N * v[i - 1][j] + a_S * v[i + 1][j]) / a_n
            + d_n[i][j] * (p[i][j] - p[i + 1][j])

    # Backwards-facing step boundary conditions (same as sides)
    for i in range(int(h_c / step)):
        for j in range(int(l_c / step)):
            u_star[i][j] = 0
            v_star[i][j] = 0
            p_new[i][j] = 100

    # Pressure correction
    p_c = np.zeros(shape=shape, dtype=float)
    for i in range(1, N - 1):
        for j in range(1, M - 1):
            a_E = -d_e[i][j] * step
            a_W = -d_e[i][j-1] * step
            a_N = -d_n[i-1][j] * step
            a_S = -d_n[i][j] * step
            a_P = a_E + a_W + a_N + a_S
            b[i][j] = step * (-(u_star[i][j] - u_star[i][j-1]) + (v_star[i][j] - v_star[i-1][j]))

            p_c[i][j] = (a_E * p_c[i][j+1] + a_W * p_c[i][j-1] + a_N * p_c[i-1][j] + a_S * p_c[i+1][j] + b[i][j]) / a_P
    p_new = p + p_c * alpha

    # Pressure boundaries
    for i in range(N - 1):
        p_new[i][0] = p_new[i][1]
        p_new[i][M - 1] = p_new[i][M - 2]
    for j in range(M - 1):
        p_new[0][j] = p_new[1][j]
        p_new[N - 1][j] = p_new[N - 2][j]

    # Velocity correction
    for i in range(1, N - 1):
        for j in range(1, M - 1):
            u_new[i][j] = u_star[i][j] + alpha * d_e[i][j] * (p_c[i + 1][j] - p_c[i][j])
            v_new[i][j] = v_star[i][j] + alpha * d_n[i][j] * (p_c[i][j] - p_c[i + 1][j])

    # Continuity residual as error measure
    error = 0
    for i in range(N):
        for j in range(M):
            error += abs(b[i][j])

    u = u_new
    v = v_new
    p = p_new

    # Plotting
    print(error)
    x, y = np.meshgrid(
        np.linspace(0, domain_size[1], shape[1]),
        np.linspace(0, domain_size[0], shape[0]),
    )

    if iteration % 100 == 0:
        pcolormesh = axes.pcolormesh(x, y, p, cmap='PuBuGn')
        if (cb):
            cb.remove()
        cb = plt.colorbar(pcolormesh)

        factor = np.sqrt(u ** 2 + v ** 2)
        u_normalized = u / factor
        v_normalized = v / factor
        plt.quiver(x, y, u_normalized, v_normalized, scale=30)
        plt.pause(0.0001)


plt.show()