using Oceananigans
using Oceananigans.Units
using CairoMakie

Lx = 2000.0kilometers
Ly = 1000.0kilometers
Lz = 4.0kilometers

Δx = 25.0kilometers
n = 30

filename = "testbed_just_water_velocity_WENO"

grid = RectilinearGrid(size=(80, 40, n), 
                            x=(0, Lx), y=(0, Ly), z=(-Lz, 0),
                            halo=(4,4,4),
                            topology=(Periodic, Bounded, Bounded))

ridge = @inline (x, y) -> -Lz + 1000.0meters * exp(-((x - Lx / 2)^2) / (100000kilometers^2))

ridge_grid = ImmersedBoundaryGrid(grid, GridFittedBottom(ridge))

coriolis = FPlane(f=-1e-4)

model = HydrostaticFreeSurfaceModel(;grid=ridge_grid, 
                                    coriolis,
                                    momentum_advection=WENO(order=5))

function uᵢ(x, y, z)
    return 0.1 * exp(z/500.0)
end

set!(model, u=uᵢ)

simulation = Simulation(model, Δt=20minutes, stop_time=3 * 30days)

#region Callbacks and Checkpointing
using Printf

wall_clock = Ref(time_ns())

function progress(sim)
    elapsed = 1e-9 * (time_ns() - wall_clock[])

    msg = @sprintf("iteration: %d, time: %s, wall time: %s, max|u|: %6.3e, m s⁻¹, max|v|: %6.3e, m s⁻¹, max|w|: %6.3e, m s⁻¹\n",
                   iteration(sim), prettytime(sim), prettytime(elapsed),
                   maximum(abs, sim.model.velocities.u),
                   maximum(abs, sim.model.velocities.v),
                   maximum(abs, sim.model.velocities.w))

    wall_clock[] = time_ns()

    @info msg

    return nothing
end

add_callback!(simulation, progress, name=:progress, IterationInterval(200))

# Write checkpoint files every year
checkpoint_time = 10days
simulation.output_writers[:checkpointer] = Checkpointer(model;
                                                        schedule = TimeInterval(checkpoint_time),
                                                        prefix = filename * "_checkpoint",
                                                        overwrite_existing = true)
#endregion

##

run!(simulation)