module TurbulenceClosures

export
    VerticalScalarDiffusivity,
    HorizontalScalarDiffusivity,
    HorizontalDivergenceScalarDiffusivity,
    ScalarDiffusivity,
    VerticalScalarBiharmonicDiffusivity,
    HorizontalScalarBiharmonicDiffusivity,
    HorizontalDivergenceScalarBiharmonicDiffusivity,
    ScalarBiharmonicDiffusivity,
    TwoDimensionalLeith,
    SmagorinskyLilly,
    DynamicSmagorinsky,
    LillyCoefficient,
    DynamicCoefficient,
    LagrangianAveraging,
    AnisotropicMinimumDissipation,
    ConvectiveAdjustmentVerticalDiffusivity,
    RiBasedVerticalDiffusivity,
    IsopycnalSkewSymmetricDiffusivity,
    CATKEVerticalDiffusivity,
    TKEDissipationVerticalDiffusivity,
    FluxTapering,

    ExplicitTimeDiscretization,
    VerticallyImplicitTimeDiscretization,

    build_closure_fields,
    compute_closure_fields!,
    step_closure_prognostics!,
    initialize_closure_fields!,

    viscosity, diffusivity,

    ∇_dot_qᶜ,
    ∂ⱼ_τ₁ⱼ,
    ∂ⱼ_τ₂ⱼ,
    ∂ⱼ_τ₃ⱼ,

    cell_diffusion_timescale,
    closure_required_tracers

using Adapt: Adapt, adapt
using DocStringExtensions: TYPEDSIGNATURES
using KernelAbstractions: @index, @kernel

using Oceananigans: Oceananigans, fields
using Oceananigans.Architectures: Architectures, on_architecture
using Oceananigans.TimeSteppers: TimeSteppers,
                                 AbstractTimeDiscretization,
                                 ExplicitTimeDiscretization,
                                 VerticallyImplicitTimeDiscretization
using Oceananigans.Grids: AbstractGrid, Bounded, Center, Face, znode
using Oceananigans.Operators: Operators,
    Ax_qᶜᶜᶜ, Ax_qᶠᶜᶜ, Ax_qᶠᶜᶠ, Ax_qᶠᶠᶜ, Ay_qᶜᶜᶜ, Ay_qᶜᶠᶜ, Ay_qᶜᶠᶠ, Ay_qᶠᶠᶜ, Az_qᶜᶜᶜ, Az_qᶜᶜᶠ, Az_qᶜᶠᶠ, Az_qᶠᶜᶠ, Az⁻¹ᶜᶜᶜ, Az⁻¹ᶜᶜᶠ, Az⁻¹ᶜᶠᶜ, Az⁻¹ᶠᶜᶜ, Az⁻¹ᶠᶠᶜ,
    V⁻¹ᶜᶜᶜ, V⁻¹ᶜᶜᶠ, V⁻¹ᶜᶠᶜ, V⁻¹ᶠᶜᶜ,
    Δx_qᶜᶜᶜ, Δx_qᶜᶠᶜ, Δx_qᶜᶠᶠ, Δx_qᶠᶜᶜ, Δxᶜᶜᶜ, Δy_qᶜᶜᶜ, Δy_qᶜᶠᶜ, Δy_qᶠᶜᶜ, Δy_qᶠᶜᶠ, Δyᶜᶜᶜ, Δzᶜᶜᶜ, Δz⁻¹ᶜᶠᶜ, Δz⁻¹ᶠᶜᶜ,
    δxᶜᵃᵃ, δxᶜᶜᶜ, δxᶜᶠᶜ, δxᶠᵃᵃ, δxᶠᶜᶜ, δxᶠᶠᶜ, δyᵃᶜᵃ, δyᵃᶠᵃ, δyᶜᶜᶜ, δyᶜᶠᶜ, δyᶠᶜᶜ, δzᵃᵃᶜ, δzᵃᵃᶠ,
    ζ₃ᶠᶠᶜ,
    ℑxyzᶠᶠᶜ, ℑxyᶜᶜᵃ, ℑxyᶜᶠᵃ, ℑxyᶠᶜᵃ, ℑxzᶜᵃᶜ, ℑxzᶜᵃᶠ, ℑxzᶠᵃᶜ, ℑxᶜᵃᵃ, ℑxᶠᵃᵃ, ℑyzᵃᶜᶜ, ℑyzᵃᶜᶠ, ℑyzᵃᶠᶜ, ℑyᵃᶜᵃ, ℑyᵃᶠᵃ, ℑzᵃᵃᶠ,
    ∂xᵣᶜᶜᶜ, ∂xᵣᶜᶠᶜ, ∂xᵣᶠᶜᶜ, ∂xᵣᶠᶜᶠ, ∂xᵣᶠᶠᶜ, ∂yᵣᶜᶜᶜ, ∂yᵣᶜᶠᶜ, ∂yᵣᶜᶠᶠ, ∂yᵣᶠᶜᶜ, ∂yᵣᶠᶠᶜ, ∂zᶜᶜᶜ, ∂zᶜᶜᶠ, ∂zᶜᶠᶠ, ∂zᶠᶜᶠ, ∂²zᶜᶜᶠ, ∂²zᶜᶠᶜ, ∂²zᶠᶜᶜ, ∂³zᶜᶜᶠ,
    ∂x_zᶠᶜᶜ, ∂x_zᶠᶜᶠ, ∂x_zᶜᶠᶜ, ∂x_zᶜᶜᶠ,
    ∂y_zᶜᶠᶜ, ∂y_zᶜᶠᶠ, ∂y_zᶠᶜᶜ, ∂y_zᶜᶜᶠ,
    ∇²hᶜᶜᶜ, ∇²hᶜᶠᶜ, ∇²hᶠᶜᶜ, ∇²ᶜᶜᶜ, ∇²ᶜᶜᶠ, ∇²ᶜᶠᶜ, ∇²ᶠᶜᶜ
using Oceananigans.BoundaryConditions: FieldBoundaryConditions, fill_halo_regions!
using Oceananigans.Utils: Utils, launch!, prettysummary, with_tracers
using Oceananigans.Fields: Field, CenterField, FunctionField, ZFaceField
using Oceananigans.ImmersedBoundaries: AbstractGridFittedBottom, ImmersedBoundaryGrid

import Oceananigans.Grids: required_halo_size_x, required_halo_size_y, required_halo_size_z

const VerticallyBoundedGrid{FT} = AbstractGrid{FT, <:Any, <:Any, <:Bounded}

#####
##### Abstract types
#####

"""
    AbstractTurbulenceClosure

Abstract supertype for turbulence closures.
"""
abstract type AbstractTurbulenceClosure{TimeDiscretization, RequiredHalo} end

# Fallbacks
validate_closure(closure) = closure
closure_summary(closure) = summary(closure)
Utils.with_tracers(tracers, closure::AbstractTurbulenceClosure) = closure
compute_closure_fields!(K, closure::AbstractTurbulenceClosure, args...; kwargs...) = nothing

import Oceananigans.TimeSteppers: step_closure_prognostics!
step_closure_prognostics!(K, closure::AbstractTurbulenceClosure, args...) = nothing
step_closure_prognostics!(K, closure::AbstractArray{<:AbstractTurbulenceClosure}, args...) = nothing

# Initialize closure fields when simulation starts (after velocities are properly set)
initialize_closure_fields!(K, closure::AbstractTurbulenceClosure, args...) = nothing
initialize_closure_fields!(K, closure::AbstractArray{<:AbstractTurbulenceClosure}, args...) = nothing
initialize_closure_fields!(K, ::Nothing, args...) = nothing

# Handle tuple of closures (iterate through each closure and its fields)
function initialize_closure_fields!(closure_fields_tuple::Tuple, closures::Tuple, args...)
    for (fields, closure) in zip(closure_fields_tuple, closures)
        initialize_closure_fields!(fields, closure, args...)
    end
    return nothing
end

# Tracer names that a closure requires (eg TKE-based closures)
# Fallbacks: by default closures do not require extra tracers.
closure_required_tracers(::Nothing) = ()
closure_required_tracers(::AbstractTurbulenceClosure) = ()
closure_required_tracers(::AbstractArray{<:AbstractTurbulenceClosure}) = ()
closure_required_tracers(closures::Tuple) = begin
    names = Symbol[]
    for c in closures
        for n in closure_required_tracers(c)
            if !(n in names)
                push!(names, n)
            end
        end
    end
    return Tuple(names)
end

# The required halo size to calculate diffusivities. Take care that if the diffusivity can
# be calculated from local information, still `B = 1`, because we need at least one additional
# point at each side to calculate viscous fluxes at the edge of the domain.
# If diffusivity itself requires one halo to be computed (e.g. κ = ℑxᶠᵃᵃ(i, j, k, grid, ℑxᶜᵃᵃ, T),
# or `AnisotropicMinimumDissipation` and `Smagorinsky`) then B = 2
@inline required_halo_size_x(::AbstractTurbulenceClosure{TD, B}) where {TD, B} = B
@inline required_halo_size_y(::AbstractTurbulenceClosure{TD, B}) where {TD, B} = B
@inline required_halo_size_z(::AbstractTurbulenceClosure{TD, B}) where {TD, B} = B

const ClosureKinda = Union{Nothing, AbstractTurbulenceClosure, AbstractArray{<:AbstractTurbulenceClosure}}
add_closure_specific_boundary_conditions(closure::ClosureKinda, bcs, args...) = bcs

# Interface for KE-based closures
function shear_production end
function buoyancy_flux end
function dissipation end
function hydrostatic_turbulent_kinetic_energy_tendency end
function buoyancy_force end # to be defined by models that want to use TurbulenceClosures
function buoyancy_tracers end # to be defined by models that want to use TurbulenceClosures

#####
##### Fallback: flux = 0
#####

for dir in (:x, :y, :z)
    diffusive_flux = Symbol(:diffusive_flux_, dir)
    viscous_flux_u = Symbol(:viscous_flux_u, dir)
    viscous_flux_v = Symbol(:viscous_flux_v, dir)
    viscous_flux_w = Symbol(:viscous_flux_w, dir)
    @eval begin
        @inline $diffusive_flux(i, j, k, grid, clo::AbstractTurbulenceClosure, args...) = zero(grid)
        @inline $viscous_flux_u(i, j, k, grid, clo::AbstractTurbulenceClosure, args...) = zero(grid)
        @inline $viscous_flux_v(i, j, k, grid, clo::AbstractTurbulenceClosure, args...) = zero(grid)
        @inline $viscous_flux_w(i, j, k, grid, clo::AbstractTurbulenceClosure, args...) = zero(grid)
    end
end

#####
##### The magic
#####

# Closure ensemble util
@inline getclosure(i, j, closure::AbstractMatrix{<:AbstractTurbulenceClosure}) = @inbounds closure[i, j]
@inline getclosure(i, j, closure::AbstractVector{<:AbstractTurbulenceClosure}) = @inbounds closure[i]
@inline getclosure(i, j, closure::AbstractTurbulenceClosure) = closure

@inline clip(x) = max(zero(x), x)

#####
##### Height, Depth and Bottom interfaces
#####

const c = Center()
const f = Face()

const AGFBIBG = ImmersedBoundaryGrid{<:Any, <:Any, <:Any, <:Any, <:Any, <:AbstractGridFittedBottom}

@inline z_top(i, j, grid) = znode(i, j, grid.Nz+1, grid, c, c, f)
@inline z_bottom(i, j, grid) = znode(i, j, 1, grid, c, c, f)
@inline z_bottom(i, j, ibg::AGFBIBG) = @inbounds ibg.immersed_boundary.bottom_height[i, j, 1]

@inline depthᶜᶜᶠ(i, j, k, grid) = clip(z_top(i, j, grid) - znode(i, j, k, grid, c, c, f))
@inline depthᶜᶜᶜ(i, j, k, grid) = clip(z_top(i, j, grid) - znode(i, j, k, grid, c, c, c))

@inline function height_above_bottomᶜᶜᶠ(i, j, k, grid)
    h = znode(i, j, k, grid, c, c, f) - z_bottom(i, j, grid)

    # Limit by thickness of cell below
    Δz = Δzᶜᶜᶜ(i, j, k-1, grid)
    return max(Δz, h)
end

@inline function height_above_bottomᶜᶜᶜ(i, j, k, grid)
    Δz = Δzᶜᶜᶜ(i, j, k, grid)
    h = znode(i, j, k, grid, c, c, c) - z_bottom(i, j, grid)
    return max(Δz/2, h)
end

@inline wall_vertical_distanceᶜᶜᶠ(i, j, k, grid) = min(depthᶜᶜᶠ(i, j, k, grid), height_above_bottomᶜᶜᶠ(i, j, k, grid))
@inline wall_vertical_distanceᶜᶜᶜ(i, j, k, grid) = min(depthᶜᶜᶜ(i, j, k, grid), height_above_bottomᶜᶜᶜ(i, j, k, grid))

include("discrete_diffusion_function.jl")
include("implicit_explicit_time_discretization.jl")
include("turbulence_closure_utils.jl")
include("closure_kernel_operators.jl")
include("velocity_tracer_gradients.jl")
include("abstract_scalar_diffusivity_closure.jl")
include("abstract_scalar_biharmonic_diffusivity_closure.jl")
include("closure_tuples.jl")
include("isopycnal_rotation_tensor_components.jl")
include("immersed_diffusive_fluxes.jl")

# Implicit closure terms (diffusion + linear terms)
include("vertically_implicit_diffusion_solver.jl")

# Implementations:
include("turbulence_closure_implementations/nothing_closure.jl")

# AbstractScalarDiffusivity closures:
include("turbulence_closure_implementations/scalar_diffusivity.jl")
include("turbulence_closure_implementations/scalar_biharmonic_diffusivity.jl")

# Dispatch on the type of the user-provided AMD model constant.
# Only numbers, arrays, and functions supported now.
@inline closure_coefficient(i, j, k, grid, C::Number) = C
@inline closure_coefficient(i, j, k, grid, C::AbstractArray) = @inbounds C[i, j, k]
@inline function closure_coefficient(i, j, k, grid, C)
    x, y, z = node(i, j, k, grid, c, c, c)
    return C(x, y, z)
end

include("turbulence_closure_implementations/anisotropic_minimum_dissipation.jl")
include("turbulence_closure_implementations/Smagorinskys/Smagorinskys.jl")
include("turbulence_closure_implementations/convective_adjustment_vertical_diffusivity.jl")
include("turbulence_closure_implementations/TKEBasedVerticalDiffusivities/TKEBasedVerticalDiffusivities.jl")
include("turbulence_closure_implementations/ri_based_vertical_diffusivity.jl")

# Special non-abstracted diffusivities:
# TODO: introduce abstract typing for these
include("turbulence_closure_implementations/isopycnal_skew_symmetric_diffusivity.jl")
include("turbulence_closure_implementations/isopycnal_skew_symmetric_diffusivity_with_triads.jl")
include("turbulence_closure_implementations/advective_skew_diffusion.jl")
include("turbulence_closure_implementations/leith_enstrophy_diffusivity.jl")

using .TKEBasedVerticalDiffusivities: CATKEVerticalDiffusivity, TKEDissipationVerticalDiffusivity
using .Smagorinskys: Smagorinsky, DynamicSmagorinsky, SmagorinskyLilly
using .Smagorinskys: LillyCoefficient, DynamicCoefficient, LagrangianAveraging

# Miscellaneous utilities
include("closure_fields.jl")
include("turbulence_closure_diagnostics.jl")

#####
##### Some value judgements here
#####

end # module
