import numpy as np
import matplotlib.pyplot as plt

# Problem parameters
L = 2 * np.pi  # Domain length
T = 5.0        # Total simulation time
Nx = 200       # Number of spatial points 
Nt = 1000      # Number of time points

# Grid setup
dx = L / (Nx - 1)
dt = T / (Nt - 1)
x = np.linspace(0, L, Nx)
t = np.linspace(0, T, Nt)

# Initial condition
u = np.sin(x) + 0.5 * np.sin(0.5 * x)

# Lax-Wendroff method for solving advection equation
for n in range(1, Nt):
    # Create copy of previous time step
    u_old = u.copy()
    
    # Periodic boundary conditions
    u[0] = u_old[-2]
    u[-1] = u_old[1]
    
    # Lax-Wendroff scheme
    for i in range(1, Nx-1):
        u[i] = u_old[i] - 0.5 * dt/dx * u_old[i] * (u_old[i+1] - u_old[i-1]) + \
               0.5 * (dt/dx)**2 * u_old[i] * (u_old[i+1] - 2*u_old[i] + u_old[i-1])

# Save final solution
np.save('/opt/CFD-Benchmark/PDE_Benchmark/results/prediction/sonnet-35/prompts/u_1D_Nonlinear_Convection.npy', u)