#!/usr/bin/env python3
import numpy as np

# Domain and simulation parameters
nx = 41       # number of grid points in x-direction
ny = 41       # number of grid points in y-direction
lx = 2.0      # domain length in x-direction
ly = 2.0      # domain length in y-direction
dx = lx / (nx - 1)
dy = ly / (ny - 1)

rho = 1.0     # density
nu = 0.1      # kinematic viscosity

t_end = 10.0  # final time
dt = 0.01     # time step size
nt = int(t_end / dt)  # number of time steps

nit = 50      # number of iterations for pressure Poisson equation

# Create grid
x = np.linspace(0, lx, nx)
y = np.linspace(0, ly, ny)

# Initialize variables: velocity components and pressure
u = np.zeros((ny, nx))
v = np.zeros((ny, nx))
p = np.zeros((ny, nx))
b = np.zeros((ny, nx))

def build_up_b(b, u, v, dx, dy, dt, rho):
    # compute the source term for pressure Poisson equation
    b[1:-1,1:-1] = (rho * (1/dt *
                    ((u[1:-1,2:] - u[1:-1,0:-2])/(2*dx) + (v[2:,1:-1] - v[0:-2,1:-1])/(2*dy))
                    - ((u[1:-1,2:] - u[1:-1,0:-2])/(2*dx))**2
                    - 2 * ((u[2:,1:-1] - u[0:-2,1:-1])/(2*dy) *
                           (v[1:-1,2:] - v[1:-1,0:-2])/(2*dx))
                    - ((v[2:,1:-1] - v[0:-2,1:-1])/(2*dy))**2))
    return b

def pressure_poisson(p, dx, dy, b, nit):
    pn = np.empty_like(p)
    for q in range(nit):
        pn[:] = p[:]
        p[1:-1,1:-1] = (((pn[1:-1,2:] + pn[1:-1,0:-2]) * dy**2 +
                         (pn[2:,1:-1] + pn[0:-2,1:-1]) * dx**2 -
                         b[1:-1,1:-1] * dx**2 * dy**2)
                         / (2*(dx**2+dy**2)))
        # Pressure boundary conditions:
        # dp/dx = 0 at x = 0 and x = lx
        p[:,0] = p[:,1]
        p[:,-1] = p[:,-2]
        # dp/dy = 0 at y = 0
        p[0,:] = p[1,:]
        # p = 0 at y = ly
        p[-1,:] = 0
    return p

# Time-stepping loop
for n in range(nt):
    un = u.copy()
    vn = v.copy()
    
    # Build source term for pressure Poisson
    b = build_up_b(b, un, vn, dx, dy, dt, rho)
    
    # Solve for pressure field using pressure Poisson equation
    p = pressure_poisson(p, dx, dy, b, nit)
    
    # Update velocity field using finite difference approximations
    # Interior points update for u
    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,0:-2]) -
             vn[1:-1,1:-1]*dt/dy*(un[1:-1,1:-1] - un[0:-2,1:-1]) -
             dt/(2*rho*dx)*(p[1:-1,2:] - p[1:-1,0:-2]) +
             nu*dt*((un[1:-1,2:] - 2*un[1:-1,1:-1] + un[1:-1,0:-2])/(dx**2) +
                     (un[2:,1:-1] - 2*un[1:-1,1:-1] + un[0:-2,1:-1])/(dy**2)))
    
    # Interior points update for v
    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,0:-2]) -
             vn[1:-1,1:-1]*dt/dy*(vn[1:-1,1:-1] - vn[0:-2,1:-1]) -
             dt/(2*rho*dy)*(p[2:,1:-1] - p[0:-2,1:-1]) +
             nu*dt*((vn[1:-1,2:] - 2*vn[1:-1,1:-1] + vn[1:-1,0:-2])/(dx**2) +
                     (vn[2:,1:-1] - 2*vn[1:-1,1:-1] + vn[0:-2,1:-1])/(dy**2)))
    
    # Apply velocity boundary conditions:
    # For u:
    u[0,:] = 0          # bottom wall no-slip
    u[-1,:] = 1         # top wall (lid-driven): u = 1
    u[:,0] = 0          # left wall no-slip
    u[:,-1] = 0         # right wall no-slip
    # For v:
    v[0,:] = 0          # bottom wall no-slip
    v[-1,:] = 0         # top wall no-slip
    v[:,0] = 0          # left wall no-slip
    v[:,-1] = 0         # right wall no-slip

# Save the final solution arrays as .npy files
np.save('/opt/CFD-Benchmark/PDE_Benchmark/results/prediction/o3-mini/prompts/u_2D_Navier_Stokes_Cavity.npy', u)
np.save('/opt/CFD-Benchmark/PDE_Benchmark/results/prediction/o3-mini/prompts/v_2D_Navier_Stokes_Cavity.npy', v)
np.save('/opt/CFD-Benchmark/PDE_Benchmark/results/prediction/o3-mini/prompts/p_2D_Navier_Stokes_Cavity.npy', p)