import numpy as np
import matplotlib.pyplot as plt

# Problem parameters
c = 1.0  # convection speed
epsilon = 5e-4  # damping factor
x_start, x_end = -5, 5
t_start, t_end = 0, 10

# Discretization parameters
nx = 800  # spatial points
nt = 4000  # time steps

# Grid setup
dx = (x_end - x_start) / (nx - 1)
dt = (t_end - t_start) / nt
x = np.linspace(x_start, x_end, nx)

# Initialize solution array
u = np.exp(-x**2)  # initial condition

# Finite difference method (Lax-Wendroff scheme)
for _ in range(nt):
    # Compute fluxes with periodic boundary conditions
    u_next = np.zeros_like(u)
    
    for i in range(1, nx-1):
        # Lax-Wendroff scheme with diffusion term
        u_next[i] = u[i] - 0.5 * c * dt/dx * (u[i+1] - u[i-1]) + \
                    0.5 * (c * dt/dx)**2 * (u[i+1] - 2*u[i] + u[i-1]) + \
                    epsilon * dt/dx**2 * (u[i+1] - 2*u[i] + u[i-1])
    
    # Periodic boundary conditions
    u_next[0] = u_next[-2]
    u_next[-1] = u_next[1]
    
    u = u_next

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