import ClimateMachine.BalanceLaws:
# declaration
    vars_state,
# initialization
    init_state_prognostic!,
    init_state_auxiliary!,
    nodal_init_state_auxiliary!,
# rhs computation
    compute_gradient_argument!,
    compute_gradient_flux!,
    flux_first_order!,
    flux_second_order!,
    source!,
# boundary conditions
    boundary_conditions,
    boundary_state!

import ClimateMachine.DGMethods.NumericalFluxes:
    CentralNumericalFluxGradient,
    CentralNumericalFluxSecondOrder,
    NumericalFluxFirstOrder,
    NumericalFluxSecondOrder,
    RusanovNumericalFlux,
    numerical_boundary_flux_second_order!,
    numerical_flux_second_order!,
    numerical_boundary_flux_first_order!,
    numerical_flux_first_order!,
    normal_boundary_flux_second_order!

"""
    Boundary condition modifications for the coupler
"""
@inline boundary_conditions(model::ModelSetup) = model.boundary_conditions

@inline function boundary_state!(numerical_flux, bc::FluidBC, model::ModelSetup, args...)
    # We need to apply boundary conditions for state variables. 
    # This depends on first, second, and high-order
    # operations, hence the diffusion model dependence.
    # TODO!: make work for higher-order diffusion
    diffusion = model.physics.diffusion # defaults to `nothing`
    calc_boundary_state!(numerical_flux, bc.ρu, model, diffusion, args...)
    calc_boundary_state!(numerical_flux, bc.ρθ, model, diffusion, args...)
end

#calc_boundary_state!(numerical_flux, bc, model, diffusion, args...) = nothing

# first-order flux and gradient flux follow the original Impenetrable() / Freeslip() conditions. 

# second order fluxes require additional treatment

function numerical_boundary_flux_second_order!(
    numerical_flux::Union{PenaltyNumFluxDiffusive},
    bctype::FluidBC,
    balance_law::ModelSetup,
    args...,
) where {S, D, A, HD}

    numerical_boundary_flux_second_order!( # this overwrites the total flux for ρθ (using methods below)
        numerical_flux,
        bctype.ρθ,
        balance_law,
        args...,
    )

    numerical_boundary_flux_second_order!( # this overwrites the total flux for ρu (using methods below)
        numerical_flux,
        bctype.ρu,
        balance_law,
        args...,
    )

end

numerical_boundary_flux_second_order!(::PenaltyNumFluxDiffusive, ::Impenetrable{FreeSlip}, _...) = nothing # no contribution to fluxᵀn.ρu from nF_diffusive

function numerical_boundary_flux_second_order!(
    numerical_flux::Union{PenaltyNumFluxDiffusive},
    bctype::CoupledPrimaryBoundary,
    balance_law::ModelSetup,
    fluxᵀn::Vars{S},
    normal_vector::SVector,
    state_prognostic⁻::Vars{S},
    state_gradient_flux⁻::Vars{D},
    state_hyperdiffusive⁻::Vars{HD},
    state_auxiliary⁻::Vars{A},
    state_prognostic⁺::Vars{S},
    state_gradient_flux⁺::Vars{D},
    state_hyperdiffusive⁺::Vars{HD},
    state_auxiliary⁺::Vars{A},
    t,
    state1⁻::Vars{S},
    diff1⁻::Vars{D},
    aux1⁻::Vars{A},
) where {S, D, A, HD}
    F_tot = calculate_land_sfc_fluxes(balance_law, state_prognostic⁻, state_auxiliary⁻, t)  # W/m^2
    fluxᵀn.ρθ += -F_tot / balance_law.parameters.cp_d

end

function numerical_boundary_flux_second_order!(
    numerical_flux::Union{PenaltyNumFluxDiffusive},
    bctype::Insulating,
    balance_law::ModelSetup,
    fluxᵀn::Vars{S},
    _...,
) where {S}
    FT = eltype(fluxᵀn)
    fluxᵀn.ρθ += FT(0)
end

"""
function numerical_flux_second_order!(::PenaltyNumFluxDiffusive, 
  Penalty flux formulation of second order numerical flux. This formulation
  computes the CentralNumericalFluxSecondOrder term first (which is just the average
  of the + and - fluxes and an edge), and then adds a "penalty" flux that relaxes
  the edge state + and - toward each other.
"""
function numerical_flux_second_order!(
    ::PenaltyNumFluxDiffusive,
    bl::ModelSetup,
    fluxᵀn::Vars{S},
    n::SVector,
    state⁻::Vars{S},
    diff⁻::Vars{D},
    hyperdiff⁻::Vars{HD},
    aux⁻::Vars{A},
    state⁺::Vars{S},
    diff⁺::Vars{D},
    hyperdiff⁺::Vars{HD},
    aux⁺::Vars{A},
    t,
) where {S, HD, D, A}

    numerical_flux_second_order!(
        CentralNumericalFluxSecondOrder(),
        bl,
        fluxᵀn,
        n,
        state⁻,
        diff⁻,
        hyperdiff⁻,
        aux⁻,
        state⁺,
        diff⁺,
        hyperdiff⁺,
        aux⁺,
        t,
    )

    Fᵀn = parent(fluxᵀn)
    FT = eltype(Fᵀn)
    tau = FT(0.0) # decide if want to use this penalty
    Fᵀn .+= tau * (parent(state⁻) - parent(state⁺))
end

numerical_flux_second_order!(::PenaltyNumFluxDiffusive, bl::SlabLandModelSetup, _...) = nothing
