import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from tqdm import tqdm

def SWE_DamBreak_JA():
    dx = 10.0
    p = 0
    q = 1000
    N = int((q - p) / dx) + 1
    x = np.linspace(p, q, N)
    g = 9.81
    CFL = 0.8
    runningtime = 300
    t = 0
    pausetime = 0.01

    def bottom(x):
        return np.zeros_like(x)# make bottom flat 
    b_phys = bottom(x)
    plt.figure(figsize=(8, 4))
    plt.plot(x, b_phys, 'k', label='Bottom')
    plt.legend()
    plt.pause(2)
    plt.close()
    eta0 = np.where(x < 150, 15.0, 5.0)
    h = eta0 - b_phys    # water depth
    h = np.maximum(h, 1e-6)
    v = np.zeros(N)
    eta_initial = eta0.copy()
    h = np.concatenate(([h[0]], h, [h[-1]]))
    v = np.concatenate(([v[0]], v, [v[-1]]))
    U1 = h.copy() # water depth
    U2 = h * v # momentum
    b = np.concatenate(([b_phys[0]], b_phys, [b_phys[-1]]))