import numpy as np
import matplotlib.pyplot as plt

PI = np.pi

np.set_printoptions(precision=2, floatmode="maxprec", suppress=True)
cb = None
plot = None


Re = 170
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 = 0.8  # Pressure under-relaxation coefficient

t_m = 1
h_c = 0.8
l_c = 0.5

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

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

p = np.zeros(shape=shape, dtype=float)
p_star = np.random.rand(shape[0], shape[1])

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


def assert_rule2(value):
    assert value > -0.01, f'Coefficient must be positive: {value}' # > 0

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

t = 0.8

iteration = 0
while error > precision:
    iteration += 1
    # Inflow boundary condition
    for i in range(N):
        y = i * step
        u_star[i][0] = 1
        v_star[i][0] = 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
            assert_rule2(a_E)
            assert_rule2(a_W)
            assert_rule2(a_N)
            assert_rule2(a_S)

            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_star[i][j + 1] - p_star[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
            assert_rule2(a_E)
            assert_rule2(a_W)
            assert_rule2(a_N)
            assert_rule2(a_S)

            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_star[i][j] - p_star[i + 1][j])

    # 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

    # 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

    # Pressure correction
    p_prime = 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
            assert_rule2(a_E)
            assert_rule2(a_W)
            assert_rule2(a_N)
            assert_rule2(a_S)
            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_prime[i][j] = (a_E * p_prime[i][j+1] + a_W * p_prime[i][j-1] + a_N * p_prime[i-1][j] + a_S * p_prime[i+1][j] + b[i][j]) / a_P
    p = p_star + p_prime * alpha
    p_star = p

    # Velocity correction
    for i in range(1, N - 1):
        for j in range(1, M - 1):
            u[i][j] = u_star[i][j] + d_e[i][j] * (p_prime[i][j + 1] - p_prime[i][j])
            v[i][j] = v_star[i][j] + d_n[i][j] * (p_prime[i][j] - p_prime[i + 1][j])

    # Backwards-facing step boundary conditions enforce
    for i in range(int(h_c / step)):
        for j in range(int(l_c / step)):
            u[i][j] = 0
            v[i][j] = 0

    # Continuity residual as error measure
    new_error = 0
    for i in range(N):
        for j in range(M):
            new_error += abs(b[i][j])
    new_error /= N * M
    if abs(new_error - error) > 10 ** -20:
        error = new_error
    else:
        print('Error decrease too small, you probably wont do better')
        break

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

    if iteration % 5 == 0 or iteration == 1:
        print(f'Plotting, iteration: {iteration}')
        if cb:
            cb.remove()
        if plot:
            plot.remove()
        factor = np.sqrt(u ** 2 + v ** 2)
        u_normalized = u / factor
        v_normalized = v / factor

        density = 1
        plot = plt.quiver(
            x[::density, ::density],
            y[::density, ::density],
            u_normalized[::density, ::density],
            v_normalized[::density, ::density],
            p[::density, ::density],
            scale=30,
            cmap='plasma'
        )
        cb = plt.colorbar()
        plt.pause(0.0001)


plt.show()