from pathlib import Path
import numpy as np
from clawpack import pyclaw
from clawpack import riemann

def main():
    grav = 9.81
    solver = pyclaw.ClawSolver1D(riemann.shallow_roe_with_efix_1D)
    solver.kernel_language = "Fortran"
    solver.bc_lower[0] = pyclaw.BC.extrap
    solver.bc_upper[0] = pyclaw.BC.extrap

    x = pyclaw.Dimension(0.0, 1000.0, 400, name="x")
    domain = pyclaw.Domain(x)
    state = pyclaw.State(domain, 2)
    state.problem_data["grav"] = grav

    xc = state.grid.x.centers
    h = np.where(xc < 500.0, 6.0, 1.0)
    state.q[0, :] = h
    state.q[1, :] = 0.0

    controller = pyclaw.Controller()
    controller.solver = solver
    controller.solution = pyclaw.Solution(state, domain)
    controller.tfinal = 25.0
    controller.num_output_times = 5
    controller.outdir = str(Path("_output_pyclaw_1d"))
    controller.keep_copy = True
    controller.run()

if __name__ == "__main__":
    main()