import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import imageio
import os

g = 9.81
rho = 1000.0
H0 = 10.0
eta0 = 0.0
tau_x0 = 0.05
use_wind = True
a_wind = tau_x0/(rho*H0) if use_wind else 0.0
nx = 254 # Grid
L = 190e3
dx = L/(nx-1)
x = np.linspace(0, L, nx)
CFL = 0.8
forecast_length_days = 10
t_end = forecast_length_days * 24 * 3600
save_interval = 3600.0
next_save = save_interval
data_dir = "swe_output"
os.makedirs(data_dir, exist_ok=True)
np.savetxt(f"{data_dir}/x_grid.dat", x)

h = np.ones(nx) * (H0 + eta0)
u = np.zeros(nx)
hu = h * u
snapshot_id = 1
eta = h - H0
np.savetxt(f"{data_dir}/eta_{snapshot_id:04d}.dat", eta) # Save initial conditions
np.savetxt(f"{data_dir}/u_{snapshot_id:04d}.dat", u)
print("Snapshot 0001 | time = 0.00 hours | initial condition")
snapshot_id += 1
def fill_reflective_ghosts(hg, hug): # Reflective BC
    hg[0]  = hg[1]
    hg[-1] = hg[-2]
    hug[0]  = -hug[1]
    hug[-1] = -hug[-2]
    return hg, hug