from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt

G = 9.81
NX = 400
X_MIN, X_MAX = 0.0, 1000.0
DX = (X_MAX - X_MIN) / NX
CFL = 0.45
T_END = 25.0

OUT = Path("public/demo-data/clawpack")
OUT.mkdir(parents=True, exist_ok=True)

x = np.linspace(X_MIN + 0.5 * DX, X_MAX - 0.5 * DX, NX)
h = np.where(x < 500.0, 6.0, 1.0)
u = np.zeros_like(h)
q = np.vstack([h, h * u])

def flux(q):
    h = np.maximum(q[0], 1.0e-8)
    hu = q[1]
    u = hu / h
    return np.vstack([hu, hu * u + 0.5 * G * h * h])

def rusanov_step(q, dt):
    h = np.maximum(q[0], 1.0e-8)
    u = q[1] / h
    a = np.abs(u) + np.sqrt(G * h)
    qL = q[:, :-1]
    qR = q[:, 1:]
    fL = flux(qL)
    fR = flux(qR)
    smax = np.maximum(a[:-1], a[1:])
    num_flux = 0.5 * (fL + fR) - 0.5 * smax * (qR - qL)
    qn = q.copy()
    qn[:, 1:-1] -= dt / DX * (num_flux[:, 1:] - num_flux[:, :-1])
    qn[:, 0] = qn[:, 1]
    qn[:, -1] = qn[:, -2]
    qn[0] = np.maximum(qn[0], 1.0e-6)
    return qn

snapshots = [(0.0, q.copy())]
t = 0.0
while t < T_END:
    h = np.maximum(q[0], 1.0e-8)
    u = q[1] / h
    dt = CFL * DX / np.max(np.abs(u) + np.sqrt(G * h))
    if t + dt > T_END:
        dt = T_END - t
    q = rusanov_step(q, dt)
    t += dt
    if len(snapshots) < 5 and t >= len(snapshots) * T_END / 4:
        snapshots.append((t, q.copy()))

plt.figure(figsize=(9, 5))
for ts, qs in snapshots:
    plt.plot(x, qs[0], label=f"t={ts:.1f}s")
plt.xlabel("x [m]")
plt.ylabel("water depth h [m]")
plt.title("1D shallow water dam-break demo")
plt.grid(True, alpha=0.3)
plt.legend()
plt.tight_layout()
plt.savefig(OUT / "clawpack_numpy_1d_dambreak.png", dpi=160)
print("saved:", OUT / "clawpack_numpy_1d_dambreak.png")