import numpy
from matplotlib import pyplot, cm
from mpl_toolkits.mplot3d import Axes3D


lineSingle = '------------------------------------------------'

print("Cavity Flow in 2D using Finite Difference Method\n")

print("Unsteady Term  : Forward Euler Scheme")
print("Convection Term: Backward Difference Scheme")
print("Diffusion Term : Central Difference Scheme\n")

print("Velocity of Cavity Lid: 1 m/s")

nx = 50  # Grid Points along X direction
ny = 50  # Grid Points along Y direction

iteration = input('Enter the number of Iterations to Solve: ')

if not iteration.isdigit():
    print("Please provide an integer\n")
else:
    iteration = int(iteration)

nt = 1000  # Number of Time Step

# Grid Spacing

dx = 2 / (nx - 1)
dy = 2 / (ny - 1)

x = numpy.linspace(0, 2, nx)
y = numpy.linspace(0, 2, ny)
X, Y = numpy.meshgrid(x, y)  # Generating a 2D Mesh

# fluid property

rho = 1  # Density
nu = 0.1  # Viscosity

dt = 0.001  # time step size

# initilization

u = numpy.zeros((ny, nx))
v = numpy.zeros((ny, nx))
p = numpy.zeros((ny, nx))
b = numpy.zeros((ny, nx))


# Solving RHS of poisson eqution in a separate function

def build_up_p(b, rho, dt, u, v, dx, dy):
    b[1:-1, 1:-1] = (
            rho * (1 / dt * ((u[1:-1, 2:] - u[1:-1, 0:-2]) / (2 * dx) + (v[2:, 1:-1] - v[:-2, 1:-1]) / (2 * dy))
                   - ((u[1:-1, 2:] - u[1:-1, 0:-2]) / (2 * dx)) ** 2 -
                   2 * ((u[2:, 1:-1] - u[:-2, 1:-1]) / (2 * dy) * (v[1:-1, 2:] - v[1:-1, :-2]) / (2 * dx))
                   - ((v[2:, 1:-1] - v[:-2, 1:-1]) / (2 * dy)) ** 2))

    return b


# Solving Poisson Equation for Pressure

def pressure_poisson(p, dx, dy, b):
    pn = numpy.empty_like(p)
    pn = p.copy()

    for q in range(iteration):
        pn = p.copy()
        p[1:-1, 1:-1] = (((pn[1:-1, 2:] + pn[1:-1, :-2]) * dy ** 2 + (pn[2:, 1:-1] + pn[:-2, 1:-1]) * dx ** 2) /
                         (2 * (dx ** 2 + dy ** 2)) - dx ** 2 * dy ** 2 / (2 * (dx ** 2 + dy ** 2)) * b[1:-1, 1:-1])

        # Boundary Condition
        p[:, -1] = p[:, -2]  # dp/dx = 0 at x = 2
        p[0, :] = p[1, :]  # dp/dy = 0 at y = 0
        p[:, 0] = p[:, 1]  # dp/dx = 0 at x = 0
        p[-1, :] = 0  # p = 0 at y = 0  pressure outlet, atmospheric presure

    return p


def cavity_flow(nt, u, v, dt, dx, dy, p, rho, nu):
    un = numpy.empty_like(u)
    vn = numpy.empty_like(v)
    b = numpy.zeros((ny, nx))

    for n in range(nt):
        un = u.copy()
        vn = v.copy()

        b = build_up_p(b, rho, dt, u, v, dx, dy)

        # Calling the Pressure Field

        p = pressure_poisson(p, dx, dy, b)

        # Solving X Momentum

        u[1:-1, 1:-1] = (un[1:-1, 1:-1] -
                         un[1:-1, 1:-1] * dt / dx *
                         (un[1:-1, 1:-1] - un[1:-1, :-2]) -
                         vn[1:-1, 1:-1] * dt / dy *
                         (un[1:-1, 1:-1] - un[:-2, 1:-1]) -
                         dt / (2 * rho * dx) * (p[1:-1, 2:] - p[1:-1, :-2]) +
                         nu * (dt / dx ** 2 *
                               (un[1:-1, 2:] - 2 * un[1:-1, 1:-1] + un[1:-1, :-2]) +
                               dt / dy ** 2 *
                               (un[2:, 1:-1] - 2 * un[1:-1, 1:-1] + un[:-2, 1:-1])))

        # Solving Y Momentum

        v[1:-1, 1:-1] = (vn[1:-1, 1:-1] -
                         un[1:-1, 1:-1] * dt / dx *
                         (vn[1:-1, 1:-1] - vn[1:-1, :-2]) -
                         vn[1:-1, 1:-1] * dt / dy *
                         (vn[1:-1, 1:-1] - vn[:-2, 1:-1]) -
                         dt / (2 * rho * dy) * (p[2:, 1:-1] - p[:-2, 1:-1]) +
                         nu * (dt / dx ** 2 *
                               (vn[1:-1, 2:] - 2 * vn[1:-1, 1:-1] + vn[1:-1, :-2]) +
                               dt / dy ** 2 *
                               (vn[2:, 1:-1] - 2 * vn[1:-1, 1:-1] + vn[:-2, 1:-1])))

        # Boundary Condition

        u[0, :] = 0
        u[:, 0] = 0
        u[:, -1] = 0
        u[-1, :] = 1  # at cavity lid velocity is 1
        v[0, :] = 0
        v[-1, :] = 0
        v[:, 0] = 0
        v[:, -1] = 0

    return u, v, p


print(lineSingle)
print("Calculating Numerical Solution......")
print(lineSingle)

u, v, p = cavity_flow(nt, u, v, dt, dx, dy, p, rho, nu)

print(lineSingle)
print("Iteration Completed!")
print(lineSingle)

fig = pyplot.figure(figsize=(11, 7), dpi=100)

# plotting pressure contour

print(lineSingle)
print("Plotting Pressure Field")
print(lineSingle)

pyplot.contourf(X, Y, u, alpha=0.5, cmap=cm.viridis)
pyplot.colorbar()
pyplot.contour(X, Y, p)
pyplot.xlabel('X')
pyplot.ylabel('Y')
pyplot.title('Pressure Contour')
pyplot.show()

# plotting vector field

print(lineSingle)
print("Plotting Velocity Vectors")
print(lineSingle)

pyplot.quiver(X[::2, ::2], Y[::2, ::2], u[::2, ::2], v[::2, ::2])
pyplot.xlabel('X')
pyplot.ylabel('Y')
pyplot.title('Velocity Vectors')
pyplot.show()

# plotting velocity contour

print(lineSingle)
print("Plotting Velocity Field")
print(lineSingle)

fig = pyplot.figure(figsize=(11, 7), dpi=100)
pyplot.contourf(X, Y, p, alpha=0.5, cmap=cm.viridis)
pyplot.colorbar()
pyplot.contour(X, Y, p, cmap=cm.viridis)
pyplot.streamplot(X, Y, u, v)
pyplot.xlabel('X')
pyplot.ylabel('Y')
pyplot.title('Velocity Contour')
pyplot.show()
