"""
    HydrostaticBoussinesqModel <: BalanceLaw

A `BalanceLaw` for ocean modeling.

write out the equations here

ρₒ = reference density of sea water
cʰ = maximum horizontal wave speed
cᶻ = maximum vertical wave speed
αᵀ = thermal expansitivity coefficient
νʰ = horizontal viscosity
νᶻ = vertical viscosity
κʰ = horizontal diffusivity
κᶻ = vertical diffusivity
fₒ = first coriolis parameter (constant term)
β  = second coriolis parameter (linear term)

# Usage

    HydrostaticBoussinesqModel(problem)

"""
struct HydrostaticBoussinesqModel{C, PS, P, MA, TA, F, FT, I} <: BalanceLaw
    param_set::PS
    problem::P
    coupling::C
    momentum_advection::MA
    tracer_advection::TA
    forcing::F
    state_filter::I
    ρₒ::FT
    cʰ::FT
    cᶻ::FT
    αᵀ::FT
    νʰ::FT
    νᶻ::FT
    κʰ::FT
    κᶻ::FT
    κᶜ::FT
    fₒ::FT
    β::FT
    function HydrostaticBoussinesqModel{FT}(
        param_set::PS,
        problem::P;
        coupling::C = Uncoupled(),
        momentum_advection::MA = nothing,
        tracer_advection::TA = NonLinearAdvectionTerm(),
        forcing::F = Forcing(),
        state_filter::I = nothing,
        ρₒ = FT(1000),  # kg / m^3
        cʰ = FT(0),     # m/s
        cᶻ = FT(0),     # m/s
        αᵀ = FT(2e-4),  # (m/s)^2 / K
        νʰ = FT(5e3),   # m^2 / s
        νᶻ = FT(5e-3),  # m^2 / s
        κʰ = FT(1e3),   # m^2 / s # horizontal diffusivity
        κᶻ = FT(1e-4),  # m^2 / s # background vertical diffusivity
        κᶜ = FT(1e-1),  # m^2 / s # diffusivity for convective adjustment
        fₒ = FT(1e-4),  # Hz
        β = FT(1e-11), # Hz / m
    ) where {FT <: AbstractFloat, PS, P, C, MA, TA, F, I}
        return new{C, PS, P, MA, TA, F, FT, I}(
            param_set,
            problem,
            coupling,
            momentum_advection,
            tracer_advection,
            forcing,
            state_filter,
            ρₒ,
            cʰ,
            cᶻ,
            αᵀ,
            νʰ,
            νᶻ,
            κʰ,
            κᶻ,
            κᶜ,
            fₒ,
            β,
        )
    end
end

HBModel = HydrostaticBoussinesqModel

boundary_conditions(ocean::HBModel) = ocean.problem.boundary_conditions

@inline noforcing(args...) = 0

function Forcing(; u = noforcing, v = noforcing, η = noforcing, θ = noforcing)
    return (u = u, v = v, η = η, θ = θ)
end

"""
    vars_state(::HBModel, ::Prognostic)

prognostic variables evolved forward in time

u = (u,v) = (zonal velocity, meridional velocity)
η = sea surface height
θ = temperature
"""
function vars_state(m::HBModel, ::Prognostic, T)
    @vars begin
        u::SVector{2, T}
        η::T # real a 2-D variable TODO: should be 2D
        θ::T
    end
end

"""
    init_state_prognostic!(::HBModel)

sets the initial value for state variables
dispatches to ocean_init_state! which is defined in a problem file such as SimpleBoxProblem.jl
"""
function init_state_prognostic!(m::HBModel, Q::Vars, A::Vars, local_geometry, t)
    return ocean_init_state!(m, m.problem, Q, A, local_geometry, t)
end

"""
    vars_state(::HBModel, ::Auxiliary)
helper variables for computation

second half is because there is no dedicated integral kernels
these variables are used to compute vertical integrals
w = vertical velocity
wz0 = w at z = 0
pkin = bulk hydrostatic pressure contribution


first half of these are fields that are used for computation
y = north-south coordinate

"""
function vars_state(m::HBModel, ::Auxiliary, T)
    @vars begin
        y::T     # y-coordinate of the box
        w::T     # ∫(-∇⋅u)
        pkin::T  # ∫(-αᵀθ)
        wz0::T   # w at z=0
        uᵈ::SVector{2, T}    # velocity deviation from vertical mean
        ΔGᵘ::SVector{2, T}   # vertically averaged tendency
    end
end

function ocean_init_aux! end

"""
    init_state_auxiliary!(::HBModel)

sets the initial value for auxiliary variables (those that aren't related to vertical integrals)
dispatches to ocean_init_aux! which is defined in a problem file such as SimpleBoxProblem.jl
"""
function init_state_auxiliary!(
    m::HBModel,
    state_auxiliary::MPIStateArray,
    grid,
    direction,
)
    init_state_auxiliary!(
        m,
        (m, A, tmp, geom) -> ocean_init_aux!(m, m.problem, A, geom),
        state_auxiliary,
        grid,
        direction,
    )
end

"""
    vars_state(::HBModel, ::Gradient)

variables that you want to take a gradient of
these are just copies in our model
"""
function vars_state(m::HBModel, ::Gradient, T)
    @vars begin
        ∇u::SVector{2, T}
        ∇uᵈ::SVector{2, T}
        ∇θ::T
    end
end

"""
    compute_gradient_argument!(::HBModel)

copy u and θ to var_gradient
this computation is done pointwise at each nodal point

# arguments:
- `m`: model in this case HBModel
- `G`: array of gradient variables
- `Q`: array of state variables
- `A`: array of aux variables
- `t`: time, not used
"""
@inline function compute_gradient_argument!(m::HBModel, G::Vars, Q::Vars, A, t)
    G.∇θ = Q.θ

    velocity_gradient_argument!(m, m.coupling, G, Q, A, t)

    return nothing
end

@inline function velocity_gradient_argument!(
    m::HBModel,
    ::Uncoupled,
    G,
    Q,
    A,
    t,
)
    G.∇u = Q.u

    return nothing
end

"""
    vars_state(::HBModel, ::GradientFlux, FT)

the output of the gradient computations
multiplies ∇u by viscosity tensor and ∇θ by the diffusivity tensor
"""
function vars_state(m::HBModel, ::GradientFlux, T)
    @vars begin
        ∇ʰu::T
        ν∇u::SMatrix{3, 2, T, 6}
        κ∇θ::SVector{3, T}
    end
end

"""
    compute_gradient_flux!(::HBModel)

copy ∇u and ∇θ to var_diffusive
this computation is done pointwise at each nodal point

# arguments:
- `m`: model in this case HBModel
- `D`: array of diffusive variables
- `G`: array of gradient variables
- `Q`: array of state variables
- `A`: array of aux variables
- `t`: time, not used
"""
@inline function compute_gradient_flux!(
    m::HBModel,
    D::Vars,
    G::Grad,
    Q::Vars,
    A::Vars,
    t,
)
    # store ∇ʰu for continuity equation (convert gradient to divergence)
    D.∇ʰu = G.∇u[1, 1] + G.∇u[2, 2]

    velocity_gradient_flux!(m, m.coupling, D, G, Q, A, t)

    κ = diffusivity_tensor(m, G.∇θ[3])
    D.κ∇θ = -κ * G.∇θ

    return nothing
end

@inline function velocity_gradient_flux!(m::HBModel, ::Uncoupled, D, G, Q, A, t)
    ν = viscosity_tensor(m)
    D.ν∇u = -ν * G.∇u

    return nothing
end

"""
    viscosity_tensor(::HBModel)

uniform viscosity with different values for horizontal and vertical directions

# Arguments
- `m`: model object to dispatch on and get viscosity parameters
"""
@inline viscosity_tensor(m::HBModel) = Diagonal(@SVector [m.νʰ, m.νʰ, m.νᶻ])

"""
    diffusivity_tensor(::HBModel)

uniform diffusivity in the horizontal direction
applies convective adjustment in the vertical, bump by 1000 if ∂θ∂z < 0

# Arguments
- `m`: model object to dispatch on and get diffusivity parameters
- `∂θ∂z`: value of the derivative of temperature in the z-direction
"""
@inline function diffusivity_tensor(m::HBModel, ∂θ∂z)
    ∂θ∂z < 0 ? κ = m.κᶜ : κ = m.κᶻ

    return Diagonal(@SVector [m.κʰ, m.κʰ, κ])
end

"""
    vars_integral(::HBModel)

location to store integrands for bottom up integrals
∇hu = the horizontal divegence of u, e.g. dw/dz
"""
function vars_state(m::HBModel, ::UpwardIntegrals, T)
    @vars begin
        ∇ʰu::T
        αᵀθ::T
    end
end

"""
    integral_load_auxiliary_state!(::HBModel)

copy w to var_integral
this computation is done pointwise at each nodal point

arguments:
m -> model in this case HBModel
I -> array of integrand variables
Q -> array of state variables
A -> array of aux variables
"""
@inline function integral_load_auxiliary_state!(
    m::HBModel,
    I::Vars,
    Q::Vars,
    A::Vars,
)
    I.∇ʰu = A.w # borrow the w value from A...
    I.αᵀθ = -m.αᵀ * Q.θ # integral will be reversed below

    return nothing
end

"""
    integral_set_auxiliary_state!(::HBModel)

copy integral results back out to aux
this computation is done pointwise at each nodal point

arguments:
m -> model in this case HBModel
A -> array of aux variables
I -> array of integrand variables
"""
@inline function integral_set_auxiliary_state!(m::HBModel, A::Vars, I::Vars)
    A.w = I.∇ʰu
    A.pkin = I.αᵀθ

    return nothing
end

"""
    vars_reverse_integral(::HBModel)

location to store integrands for top down integrals
αᵀθ = density perturbation
"""
function vars_state(m::HBModel, ::DownwardIntegrals, T)
    @vars begin
        αᵀθ::T
    end
end

"""
    reverse_integral_load_auxiliary_state!(::HBModel)

copy αᵀθ to var_reverse_integral
this computation is done pointwise at each nodal point

arguments:
m -> model in this case HBModel
I -> array of integrand variables
A -> array of aux variables
"""
@inline function reverse_integral_load_auxiliary_state!(
    m::HBModel,
    I::Vars,
    Q::Vars,
    A::Vars,
)
    I.αᵀθ = A.pkin

    return nothing
end

"""
    reverse_integral_set_auxiliary_state!(::HBModel)

copy reverse integral results back out to aux
this computation is done pointwise at each nodal point

arguments:
m -> model in this case HBModel
A -> array of aux variables
I -> array of integrand variables
"""
@inline function reverse_integral_set_auxiliary_state!(
    m::HBModel,
    A::Vars,
    I::Vars,
)
    A.pkin = I.αᵀθ

    return nothing
end

"""
    flux_first_order!(::HBModel)

calculates the hyperbolic flux contribution to state variables
this computation is done pointwise at each nodal point

# arguments:
m -> model in this case HBModel
F -> array of fluxes for each state variable
Q -> array of state variables
A -> array of aux variables
t -> time, not used

# computations
∂ᵗu = ∇⋅(g*η + g∫αᵀθdz + v⋅u)
∂ᵗθ = ∇⋅(vθ) where v = (u,v,w)
"""
@inline function flux_first_order!(
    m::HBModel,
    F::Grad,
    Q::Vars,
    A::Vars,
    t::Real,
    direction,
)
    # ∇ʰ • (g η)
    hydrostatic_pressure!(m, m.coupling, F, Q, A, t)

    # ∇ʰ • (- ∫(αᵀ θ))
    kinematic_pressure!(m, F, Q, A, t)

    # ∇ʰ • (v ⊗ u)
    momentum_advection!(m, m.momentum_advection, F, Q, A, t)

    # ∇ • (u θ)
    tracer_advection!(m, m.tracer_advection, F, Q, A, t)

    return nothing
end

@inline function hydrostatic_pressure!(m::HBModel, ::Uncoupled, F, Q, A, t)
    η = Q.η
    Iʰ = @SMatrix [
        1 -0
        -0 1
        -0 -0
    ]

    F.u += grav(parameter_set(m)) * η * Iʰ

    return nothing
end

@inline function kinematic_pressure!(m::HBModel, F, Q, A, t)
    pkin = A.pkin
    Iʰ = @SMatrix [
        1 -0
        -0 1
        -0 -0
    ]
    F.u += grav(parameter_set(m)) * pkin * Iʰ

    return nothing
end

momentum_advection!(::HBModel, ::Nothing, _...) = nothing
@inline function momentum_advection!(
    ::HBModel,
    ::NonLinearAdvectionTerm,
    F,
    Q,
    A,
    t,
)
    u = Q.u
    @inbounds v = @SVector [Q.u[1], Q.u[2], A.w]

    F.u += v * u'

    return nothing
end

tracer_advection!(::HBModel, ::Nothing, _...) = nothing
@inline function tracer_advection!(
    ::HBModel,
    ::NonLinearAdvectionTerm,
    F,
    Q,
    A,
    t,
)
    θ = Q.θ
    @inbounds v = @SVector [Q.u[1], Q.u[2], A.w]

    F.θ += v * θ

    return nothing
end

"""
    flux_second_order!(::HBModel)

calculates the parabolic flux contribution to state variables
this computation is done pointwise at each nodal point

# arguments:
- `m`: model in this case HBModel
- `F`: array of fluxes for each state variable
- `Q`: array of state variables
- `D`: array of diff variables
- `A`: array of aux variables
- `t`: time, not used

# computations
∂ᵗu = -∇⋅(ν∇u)
∂ᵗθ = -∇⋅(κ∇θ)
"""
@inline function flux_second_order!(
    m::HBModel,
    F::Grad,
    Q::Vars,
    D::Vars,
    HD::Vars,
    A::Vars,
    t::Real,
)
    F.u += D.ν∇u
    F.θ += D.κ∇θ

    return nothing
end

"""
    source!(::HBModel)

Calculates the source term contribution to state variables.
This computation is done pointwise at each nodal point.

Arguments:
    m -> model in this case HBModel
    F -> array of fluxes for each state variable
    Q -> array of state variables
    A -> array of aux variables
    t -> time, not used

Computations:
    ∂ᵗu = -f × u
    ∂ᵗη = w|(z=0)
"""
@inline function source!(
    m::HBModel,
    S::Vars,
    Q::Vars,
    D::Vars,
    A::Vars,
    t::Real,
    direction,
)
    # explicit forcing for SSH
    wz0 = A.wz0
    S.η += wz0

    coriolis_force!(m, m.coupling, S, Q, A, t)

    # Arguments for forcing functions
    # args = y, t, u, v, w, η, θ
    args = tuple(A.y, t, Q.u[1], Q.u[2], A.w, Q.η, Q.θ)

    Su = m.forcing.u(args...)
    Sv = m.forcing.v(args...)

    S.u += @SVector [Su, Sv]
    S.η += m.forcing.η(args...)
    S.θ += m.forcing.θ(args...)

    return nothing
end

@inline function coriolis_force!(m::HBModel, ::Uncoupled, S, Q, A, t)
    # f × u
    f = coriolis_parameter(m, m.problem, A.y)
    u, v = Q.u # Horizontal components of velocity
    S.u -= @SVector [-f * v, f * u]

    return nothing
end

"""
    wavespeed(::HBModel)

calculates the wavespeed for rusanov flux
"""
@inline wavespeed(m::HBModel, n⁻, _...) = abs(SVector(m.cʰ, m.cʰ, m.cᶻ)' * n⁻)

"""
    update_penalty(::HBModel)
    set Δη = 0 when computing numerical fluxes
"""
# We want not have jump penalties on η (since not a flux variable)
function update_penalty!(
    ::RusanovNumericalFlux,
    ::HBModel,
    n⁻,
    λ,
    ΔQ::Vars,
    Q⁻,
    A⁻,
    Q⁺,
    A⁺,
    t,
)
    ΔQ.η = -0

    return nothing
end

filter_state!(Q, filter::Nothing, grid) = nothing
filter_state!(Q, filter, grid) = apply!(Q, UnitRange(1, size(Q, 2)), grid)

"""
    update_auxiliary_state!(::HBModel)

Applies the vertical filter to the zonal and meridional velocities to preserve numerical incompressibility
Applies an exponential filter to θ to anti-alias the non-linear advective term

Doesn't actually touch the aux variables any more, but we need a better filter interface than this anyways
"""
function update_auxiliary_state!(
    dg::DGModel,
    m::HBModel,
    Q::MPIStateArray,
    t::Real,
    elems::UnitRange,
)
    FT = eltype(Q)
    MD = dg.modeldata

    # `update_aux!` gets called twice, once for the real elements and once for
    # the ghost elements.  Only apply the filters to the real elems.
    if elems == dg.grid.topology.realelems
        # required to ensure that after integration velocity field is divergence free
        vert_filter = MD.vert_filter
        apply!(Q, (:u,), dg.grid, vert_filter, direction = VerticalDirection())

        exp_filter = MD.exp_filter
        apply!(Q, (:θ,), dg.grid, exp_filter, direction = VerticalDirection())

        filter_state!(Q, m.state_filter, dg.grid)
    end

    compute_flow_deviation!(dg, m, m.coupling, Q, t)

    return true
end

@inline compute_flow_deviation!(dg, ::HBModel, ::Uncoupled, _...) = nothing


"""
    update_auxiliary_state_gradient!(::HBModel)

    ∇hu to w for integration
    performs integration for w and pkin (should be moved to its own integral kernels)
    copies down w and wz0 because we don't have 2D structures
"""
function update_auxiliary_state_gradient!(
    dg::DGModel,
    m::HBModel,
    Q::MPIStateArray,
    t::Real,
    elems::UnitRange,
)
    FT = eltype(Q)
    A = dg.state_auxiliary
    D = dg.state_gradient_flux

    # load -∇ʰu as ∂ᶻw
    index_w = varsindex(vars_state(m, Auxiliary(), FT), :w)
    index_∇ʰu = varsindex(vars_state(m, GradientFlux(), FT), :∇ʰu)
    @views @. A.data[:, index_w, elems] = -D.data[:, index_∇ʰu, elems]

    # compute integrals for w and pkin
    indefinite_stack_integral!(dg, m, Q, A, t, elems) # bottom -> top
    reverse_indefinite_stack_integral!(dg, m, Q, A, t, elems) # top -> bottom

    # We are unable to use vars (ie A.w) for this because this operation will
    # return a SubArray, and adapt (used for broadcasting along reshaped arrays)
    # has a limited recursion depth for the types allowed.
    number_aux = number_states(m, Auxiliary())
    index_wz0 = varsindex(vars_state(m, Auxiliary(), FT), :wz0)
    info = basic_grid_info(dg)
    Nqh, Nqk = info.Nqh, info.Nqk
    nelemv, nelemh = info.nvertelem, info.nhorzelem
    nrealelemh = info.nhorzrealelem

    # project w(z=0) down the stack
    data = reshape(A.data, Nqh, Nqk, number_aux, nelemv, nelemh)
    flat_wz0 = @view data[:, end:end, index_w, end:end, 1:nrealelemh]
    boxy_wz0 = @view data[:, :, index_wz0, :, 1:nrealelemh]
    boxy_wz0 .= flat_wz0

    return true
end

"""
    boundary_state!(nf, bc, ::HBModel, args...)

applies boundary conditions for the hyperbolic fluxes
dispatches to a function in OceanBoundaryConditions.jl based on bytype defined by a problem such as SimpleBoxProblem.jl
"""
@inline boundary_state!(nf, bc, ocean::HBModel, args...) =
    _ocean_boundary_state!(nf, bc, ocean, args...)

#=
"""
    boundary_state!(nf, ::HBModel, args...)

applies boundary conditions for the hyperbolic fluxes
dispatches to a function in OceanBoundaryConditions.jl based on bytype defined by a problem such as SimpleBoxProblem.jl
"""
@inline function boundary_state!(nf, ocean::HBModel, args...)
    boundary_conditions = ocean.problem.boundary_conditions
    return ocean_boundary_state!(nf, boundary_conditions, ocean, args...)
end
=#

"""
    ocean_boundary_state!(nf, bc::OceanBC, ::HBModel, args...)

splits boundary condition application into velocity and temperature conditions
"""
@inline function ocean_boundary_state!(nf, bc::OceanBC, ocean::HBModel, args...)
    ocean_velocity_boundary_state!(nf, bc.velocity, ocean, args...)
    ocean_temperature_boundary_state!(nf, bc.temperature, ocean, args...)

    return nothing
end

"""
    ocean_boundary_state!(nf, boundaries::Tuple, ::HBModel,
                          Q⁺, A⁺, n, Q⁻, A⁻, bctype)

applies boundary conditions for the first-order and gradient fluxes
dispatches to a function in OceanBoundaryConditions.jl based on bytype defined by a problem such as SimpleBoxProblem.jl
"""
@generated function ocean_boundary_state!(
    nf::Union{NumericalFluxFirstOrder, NumericalFluxGradient},
    boundaries::Tuple,
    ocean,
    Q⁺,
    A⁺,
    n,
    Q⁻,
    A⁻,
    bctype,
    t,
    args...,
)
    N = fieldcount(boundaries)
    return quote
        Base.Cartesian.@nif(
            $(N + 1),
            i -> bctype == i, # conditionexpr
            i -> ocean_boundary_state!(
                nf,
                boundaries[i],
                ocean,
                Q⁺,
                A⁺,
                n,
                Q⁻,
                A⁻,
                t,
            ), # expr
            i -> error("Invalid boundary tag")
        ) # elseexpr
        return nothing
    end
end

"""
    ocean_boundary_state!(nf, boundaries::Tuple, ::HBModel,
                          Q⁺, A⁺, D⁺, n, Q⁻, A⁻, D⁻, bctype)

applies boundary conditions for the second-order fluxes
dispatches to a function in OceanBoundaryConditions.jl based on bytype defined by a problem such as SimpleBoxProblem.jl
"""
@generated function ocean_boundary_state!(
    nf::NumericalFluxSecondOrder,
    boundaries::Tuple,
    ocean,
    Q⁺,
    D⁺,
    HD⁺,
    A⁺,
    n,
    Q⁻,
    D⁻,
    HD⁻,
    A⁻,
    bctype,
    t,
    args...,
)
    N = fieldcount(boundaries)
    return quote
        Base.Cartesian.@nif(
            $(N + 1),
            i -> bctype == i, # conditionexpr
            i -> ocean_boundary_state!(
                nf,
                boundaries[i],
                ocean,
                Q⁺,
                D⁺,
                A⁺,
                n,
                Q⁻,
                D⁻,
                A⁻,
                t,
            ), # expr
            i -> error("Invalid boundary tag")
        ) # elseexpr
        return nothing
    end
end
