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

# Parameters
gamma = 1.4
xmin, xmax = -1.0, 1.0
T = 0.25
N = 400            # number of physical cells
CFL = 0.5

# Create grid with ghost cells (index 0 and N+1)
dx = (xmax - xmin) / N
# physical cell centers for interior cells only
x = np.linspace(xmin + dx/2, xmax - dx/2, N)

# Total number of cells including ghost cells
n_cells = N + 2

# Initialize conservative variables U = [rho, rho*u, rho*E] with shape (3, n_cells)
U = np.zeros((3, n_cells))

# Set initial conditions for interior cells (indices 1 to N)
rho = np.where(x < 0.0, 1.0, 0.125)
u = np.zeros_like(x)
p = np.where(x < 0.0, 1.0, 0.1)
E = p/((gamma - 1)*rho) + 0.5*u**2

# Fill interior cells
U[0, 1:-1] = rho
U[1, 1:-1] = rho * u
U[2, 1:-1] = rho * E

def apply_reflective_bc(U):
    # Reflective BC: density and energy same, momentum reversed.
    # Left ghost: mirror interior cell at index 1
    U[0,0] = U[0,1]
    U[1,0] = - U[1,1]
    U[2,0] = U[2,1]
    # Right ghost: mirror interior cell at index -2
    U[0,-1] = U[0,-2]
    U[1,-1] = - U[1,-2]
    U[2,-1] = U[2,-2]

def compute_flux(U):
    # U shape: (3, n_cells)
    rho = U[0, :]
    mom = U[1, :]
    E_tot = U[2, :]
    u = mom / rho
    p = (gamma - 1) * (E_tot - 0.5 * rho * u**2)
    F = np.zeros_like(U)
    F[0, :] = mom
    F[1, :] = rho * u**2 + p
    F[2, :] = u * (E_tot + p)
    return F

# Main time loop (explicit finite difference using Lax-Friedrichs scheme)
t = 0.0
while t < T:
    apply_reflective_bc(U)
    F = compute_flux(U)
    # Determine maximum wave speed for CFL time step
    rho_in = U[0, 1:-1]
    u_in = U[1, 1:-1] / rho_in
    p_in = (gamma - 1) * (U[2, 1:-1] - 0.5 * rho_in * u_in**2)
    c = np.sqrt(gamma * p_in / rho_in)
    max_speed = np.max(np.abs(u_in) + c)
    dt = CFL * dx / max_speed
    if t + dt > T:
        dt = T - t

    U_new = U.copy()
    # Lax-Friedrichs update for interior cells (i=1 ... N)
    # U_new[:, i] = 0.5*(U[:,i+1] + U[:,i-1]) - dt/(2*dx) * (F[:,i+1] - F[:,i-1])
    U_new[:, 1:-1] = 0.5 * (U[:, 2:] + U[:, 0:-2]) - (dt/(2*dx)) * (F[:, 2:] - F[:, 0:-2])
    U = U_new
    t += dt

# Apply BC one final time
apply_reflective_bc(U)

# Extract final solution for interior cells
rho_final = U[0, 1:-1]
u_final = U[1, 1:-1] / rho_final
p_final = (gamma - 1) * (U[2, 1:-1] - 0.5 * rho_final * u_final**2)

# Save final 1D arrays to .npy files
np.save('/opt/CFD-Benchmark/PDE_Benchmark/results/prediction/o3-mini/prompts/rho_final_1D_Euler_Shock_Tube.npy', rho_final)
np.save('/opt/CFD-Benchmark/PDE_Benchmark/results/prediction/o3-mini/prompts/u_final_1D_Euler_Shock_Tube.npy', u_final)
np.save('/opt/CFD-Benchmark/PDE_Benchmark/results/prediction/o3-mini/prompts/p_final_1D_Euler_Shock_Tube.npy', p_final)