using Base: @propagate_inbounds

using Oceananigans.Architectures: Architectures, AbstractArchitecture, architecture
using Oceananigans.Grids: rnode, rnodes

const ArchOrNothing = Union{AbstractArchitecture, Nothing}
const GridOrNothing = Union{AbstractGrid, Nothing}

"""
    AbstractField{LX, LY, LZ, G, T, N}

Abstract supertype for fields located at `(LX, LY, LZ)`
and defined on a grid `G` with eltype `T` and `N` dimensions.

Note: we need the parameter `T` to subtype AbstractArray.
"""
abstract type AbstractField{LX, LY, LZ, G <: GridOrNothing, T, N} <: AbstractArray{T, N} end

Base.IndexStyle(::AbstractField) = IndexCartesian()

#####
##### AbstractField functionality
#####

@inline Oceananigans.location(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = (LX, LY, LZ) # note no instantiation
@inline Oceananigans.instantiated_location(::AbstractField{LX, LY, LZ}) where {LX, LY, LZ} = (LX(), LY(), LZ())
Base.eltype(::AbstractField{<:Any, <:Any, <:Any, <:Any, T}) where T = T
Base.eltype(::Type{<:AbstractField{<:Any, <:Any, <:Any, <:Any, T}}) where T = T

"Returns the grid on which `f` is defined."
Grids.grid(f::AbstractField) = f.grid

"Returns the architecture of on which `f` is defined."
Architectures.architecture(f::AbstractField) = architecture(grid(f))
Architectures.child_architecture(f::AbstractField) = child_architecture(architecture(f))

"Returns the topology of a fields' `grid`."
@inline Grids.topology(f::AbstractField, args...) = topology(f.grid, args...)

"""
$(TYPEDSIGNATURES)

Returns the size of an `AbstractField{LX, LY, LZ}` located at `LX, LY, LZ`.
This is a 3-tuple of integers corresponding to the number of interior nodes
of `f` along `x, y, z`.
"""
Base.size(f::AbstractField) = size(f.grid, location(f))
Base.length(f::AbstractField) = prod(size(f))
Base.parent(f::AbstractField) = f

const Abstract3DField = AbstractField{<:Any, <:Any, <:Any, <:Any, <:Any, 3}
const Abstract4DField = AbstractField{<:Any, <:Any, <:Any, <:Any, <:Any, 4}

# TODO: to omit boundaries on Face fields, we have to return 2:N
# when topo=Bounded, and loc=Face
@inline axis(::Colon, N) = Base.OneTo(N)
@inline axis(index::UnitRange, N) = index
@inline axis(index::Base.OneTo, N) = index

@inline function Base.axes(f::Abstract3DField)
    Nx, Ny, Nz = size(f)
    ix, iy, iz = indices(f)

    ax = axis(ix, Nx)
    ay = axis(iy, Ny)
    az = axis(iz, Nz)

    return (ax, ay, az)
end

@inline function Base.axes(f::Abstract4DField)
    Nx, Ny, Nz, Nt = size(f)
    ix, iy, iz = indices(f)

    ax = axis(ix, Nx)
    ay = axis(iy, Ny)
    az = axis(iz, Nz)
    at = Base.OneTo(Nt)

    return (ax, ay, az, at)
end

"""
$(TYPEDSIGNATURES)

Returns a 3-tuple that gives the "total" size of a field including
both interior points and halo points.
"""
Grids.total_size(f::AbstractField) = total_size(f.grid, location(f))

interior(f::AbstractField) = f

#####
##### Coordinates of fields
#####

@propagate_inbounds  Grids.node(i, j, k, ψ::AbstractField) =  node(i, j, k, ψ.grid, instantiated_location(ψ)...)
@propagate_inbounds Grids.xnode(i, j, k, ψ::AbstractField) = xnode(i, j, k, ψ.grid, instantiated_location(ψ)...)
@propagate_inbounds Grids.ynode(i, j, k, ψ::AbstractField) = ynode(i, j, k, ψ.grid, instantiated_location(ψ)...)
@propagate_inbounds Grids.znode(i, j, k, ψ::AbstractField) = znode(i, j, k, ψ.grid, instantiated_location(ψ)...)
@propagate_inbounds Grids.rnode(i, j, k, ψ::AbstractField) = rnode(i, j, k, ψ.grid, instantiated_location(ψ)...)

Grids.xnodes(ψ::AbstractField; kwargs...) = xnodes(ψ.grid, instantiated_location(ψ)...; kwargs...)
Grids.ynodes(ψ::AbstractField; kwargs...) = ynodes(ψ.grid, instantiated_location(ψ)...; kwargs...)
Grids.znodes(ψ::AbstractField; kwargs...) = znodes(ψ.grid, instantiated_location(ψ)...; kwargs...)
Grids.rnodes(ψ::AbstractField; kwargs...) = rnodes(ψ.grid, instantiated_location(ψ)...; kwargs...)

Grids.nodes(ψ::AbstractField; kwargs...) = nodes(ψ.grid, instantiated_location(ψ); kwargs...)

#####
##### Some conveniences
#####

for f in (:+, :-)
    @eval Base.$f(ϕ::AbstractArray, ψ::AbstractField) = $f(ϕ, interior(ψ))
    @eval Base.$f(ϕ::AbstractField, ψ::AbstractArray) = $f(interior(ϕ), ψ)
end

const XReducedAF = AbstractField{Nothing}
const YReducedAF = AbstractField{<:Any, Nothing}
const ZReducedAF = AbstractField{<:Any, <:Any, Nothing}

const YZReducedAF = AbstractField{<:Any, Nothing, Nothing}
const XZReducedAF = AbstractField{Nothing, <:Any, Nothing}
const XYReducedAF = AbstractField{Nothing, Nothing, <:Any}

const XYZReducedAF = AbstractField{Nothing, Nothing, Nothing}

reduced_dimensions(field::AbstractField) = ()
reduced_dimensions(field::XReducedAF)    = tuple(1)
reduced_dimensions(field::YReducedAF)    = tuple(2)
reduced_dimensions(field::ZReducedAF)    = tuple(3)
reduced_dimensions(field::YZReducedAF)   = (2, 3)
reduced_dimensions(field::XZReducedAF)   = (1, 3)
reduced_dimensions(field::XYReducedAF)   = (1, 2)
reduced_dimensions(field::XYZReducedAF)  = (1, 2, 3)
