using Oceananigans.TimeSteppers: _rk3_substep_field!, stage_Δt, compute_flux_bc_tendencies!
import Oceananigans.TimeSteppers: rk3_substep!

"""
$(TYPEDSIGNATURES)

Perform a single RK3 substep for `NonhydrostaticModel` with pressure correction.
Dispatches to `pressure_correction_rk3_substep!` which advances velocities and tracers
using the RK3 coefficients, then applies a pressure correction to enforce incompressibility.
"""
rk3_substep!(model::NonhydrostaticModel, Δt, γ, ζ, callbacks) =
    pressure_correction_rk3_substep!(model, Δt, γ, ζ, callbacks)

"""
$(TYPEDSIGNATURES)

Implement a single RK3 substep with pressure correction for `NonhydrostaticModel`.

The substep advances the state as

    U += Δt * (γⁿ * Gⁿ + ζⁿ * G⁻)

where:
- `γⁿ` is the coefficient for the current tendency
- `ζⁿ` is the coefficient for the previous tendency (or `nothing` for the first substep)
- The effective substep size is `Δτ = Δt * (γⁿ + ζⁿ)`

After advancing velocities, a pressure Poisson equation is solved and velocities
are corrected to satisfy the incompressibility constraint.
"""
function pressure_correction_rk3_substep!(model, Δt, γⁿ, ζⁿ, callbacks)
    grid = model.grid
    FT        = eltype(grid)
    kernel_Δt = convert(FT, Δt)
    Δτ        = convert(FT, stage_Δt(Δt, γⁿ, ζⁿ))

    compute_flux_bc_tendencies!(model)
    model_fields = prognostic_fields(model)

    # Prognostic variables stepping
    for (i, name) in enumerate(keys(model_fields))
        field = model_fields[name]
        exclude_periphery = i < 4 # We assume that the first 3 fields are velocity / momentum variables
        kernel_args = (field, kernel_Δt, γⁿ, ζⁿ, model.timestepper.Gⁿ[name], model.timestepper.G⁻[name])
        launch!(architecture(grid), grid, :xyz, _rk3_substep_field!, kernel_args...; exclude_periphery)

        implicit_step!(field,
                       model.timestepper.implicit_solver,
                       model.closure,
                       model.closure_fields,
                       Val(i-3), # We assume that the first 3 fields are velocity / momentum variables
                       model.clock,
                       fields(model),
                       Δτ,
                       model.advection,
                       model.velocities)
    end

    compute_pressure_correction!(model, Δτ)
    make_pressure_correction!(model, Δτ)

    return nothing
end
