import Oceananigans: prognostic_state, restore_prognostic_state!
import Oceananigans.Grids: grid

struct ZeroField{T, N} <: AbstractField{Nothing, Nothing, Nothing, Nothing, T, N} end
struct OneField{T, N} <: AbstractField{Nothing, Nothing, Nothing, Nothing, T, N} end

ZeroField(T=Int) = ZeroField{T, 3}() # default 3D, integer 0
OneField(T=Int) = OneField{T, 3}() # default 3D, integer 1

@inline Base.getindex(::ZeroField{T, N}, ind...) where {N, T} = zero(T)
@inline Base.getindex(::OneField{T, N}, ind...) where {N, T} = one(T)

struct ConstantField{T, N} <: AbstractField{Nothing, Nothing, Nothing, Nothing, T, N}
    constant :: T
    ConstantField{N}(constant::T) where {T, N} = new{T, N}(constant)
end

# Default 3-dimensional field
ConstantField(constant) = ConstantField{3}(constant)

@inline Base.getindex(f::ConstantField, ind...) = f.constant

const CF = Union{ConstantField, ZeroField, OneField}

grid(::CF) = nothing

BoundaryConditions.fill_halo_regions!(::ZeroField, args...; kw...) = nothing
BoundaryConditions.fill_halo_regions!(::ConstantField, args...; kw...) = nothing

prognostic_state(f::Union{ZeroField, OneField, ConstantField}) = nothing
restore_prognostic_state!(restored::Union{ZeroField, OneField, ConstantField}, from) = restored
restore_prognostic_state!(restored::Union{ZeroField, OneField, ConstantField}, ::Nothing) = restored
