using Oceananigans.Grids: AbstractGrid, hack_cosd, hack_sind, λnode

"""
Notes:
This file defines grid lengths, areas, and volumes for staggered structured grids.
Each "reference cell" is associated with an index i, j, k.
The "location" of each reference cell is roughly the geometric centroid of the reference cell.
On the staggered grid, there are 7 cells additional to the "reference cell"
that are staggered with respect to the reference cell in x, y, and/or z.

The staggering is denoted by the locations "Center" and "Face":
  - "Center" is shared with the reference cell;
  - "Face" lies between reference cell centers, roughly at the interface between
    reference cells.

The three-dimensional location of an object is defined by a 3-tuple of locations, and
denoted by a triplet of superscripts. For example, an object `φ` whose cell is located at
(Center, Center, Face) is denoted `φᶜᶜᶠ`. `ᶜᶜᶠ` is Centered in `x`, `Centered` in `y`, and on
reference cell interfaces in `z` (this is where the vertical velocity is located, for example).

The operators in this file fall into three categories:

1. Operators needed for an algorithm valid on rectilinear grids with
   at most a stretched vertical dimension and regular horizontal dimensions.
2. Operators needed for an algorithm on a grid that is curvilinear in the horizontal
   at rectilinear (possibly stretched) in the vertical.
"""

#####
#####
##### Spacings!!
#####
#####
# This metaprogramming loop defines all possible combinations of locations for spacings.
# The general 2D and 3D spacings are reconducted to their one - dimensional counterparts.
# Grids that do not have a specific one - dimensional spacing for a given location need to
# extend these functions (for example, LatitudeLongitudeGrid).

# Calling a non existing function (for example Δxᶜᵃᶜ on an OrthogonalSphericalShellGrid) will throw an error because
# the associated one - dimensional function is not defined.
for L1 in (:ᶜ, :ᶠ), L2 in (:ᶜ, :ᶠ)
    Δxˡᵃᵃ = Symbol(:Δx, L1, :ᵃ, :ᵃ)
    Δyᵃˡᵃ = Symbol(:Δy, :ᵃ, L1, :ᵃ)
    Δzᵃᵃˡ = Symbol(:Δz, :ᵃ, :ᵃ, L1)
    Δλˡᵃᵃ = Symbol(:Δλ, L1, :ᵃ, :ᵃ)
    Δφᵃˡᵃ = Symbol(:Δφ, :ᵃ, L1, :ᵃ)
    Δrᵃᵃˡ = Symbol(:Δr, :ᵃ, :ᵃ, L1)

    Δxˡˡᵃ = Symbol(:Δx, L1, L2, :ᵃ)
    Δyˡˡᵃ = Symbol(:Δy, L2, L1, :ᵃ)
    Δzˡᵃˡ = Symbol(:Δz, L2, :ᵃ, L1)
    Δλˡˡᵃ = Symbol(:Δλ, L1, L2, :ᵃ)
    Δφˡˡᵃ = Symbol(:Δφ, L2, L1, :ᵃ)
    Δrˡᵃˡ = Symbol(:Δr, L2, :ᵃ, L1)

    Δxˡᵃˡ = Symbol(:Δx, L1, :ᵃ, L2)
    Δyᵃˡˡ = Symbol(:Δy, :ᵃ, L1, L2)
    Δzᵃˡˡ = Symbol(:Δz, :ᵃ, L2, L1)
    Δλˡᵃˡ = Symbol(:Δλ, L1, :ᵃ, L2)
    Δφᵃˡˡ = Symbol(:Δφ, :ᵃ, L1, L2)
    Δrᵃˡˡ = Symbol(:Δr, :ᵃ, L2, L1)

    @eval @inline $Δxˡˡᵃ(i, j, k, grid) = $Δxˡᵃᵃ(i, j, k, grid)
    @eval @inline $Δxˡᵃˡ(i, j, k, grid) = $Δxˡᵃᵃ(i, j, k, grid)

    @eval @inline $Δyˡˡᵃ(i, j, k, grid) = $Δyᵃˡᵃ(i, j, k, grid)
    @eval @inline $Δyᵃˡˡ(i, j, k, grid) = $Δyᵃˡᵃ(i, j, k, grid)

    @eval @inline $Δzˡᵃˡ(i, j, k, grid) = $Δzᵃᵃˡ(i, j, k, grid)
    @eval @inline $Δzᵃˡˡ(i, j, k, grid) = $Δzᵃᵃˡ(i, j, k, grid)

    @eval @inline $Δλˡˡᵃ(i, j, k, grid) = $Δλˡᵃᵃ(i, j, k, grid)
    @eval @inline $Δλˡᵃˡ(i, j, k, grid) = $Δλˡᵃᵃ(i, j, k, grid)

    @eval @inline $Δφˡˡᵃ(i, j, k, grid) = $Δφᵃˡᵃ(i, j, k, grid)
    @eval @inline $Δφᵃˡˡ(i, j, k, grid) = $Δφᵃˡᵃ(i, j, k, grid)

    @eval @inline $Δrˡᵃˡ(i, j, k, grid) = $Δrᵃᵃˡ(i, j, k, grid)
    @eval @inline $Δrᵃˡˡ(i, j, k, grid) = $Δrᵃᵃˡ(i, j, k, grid)

    for L3 in (:ᶜ, :ᶠ)
        Δxˡˡˡ = Symbol(:Δx, L1, L2, L3)
        Δyˡˡˡ = Symbol(:Δy, L2, L1, L3)
        Δzˡˡˡ = Symbol(:Δz, L2, L3, L1)
        Δλˡˡˡ = Symbol(:Δλ, L1, L2, L3)
        Δφˡˡˡ = Symbol(:Δφ, L2, L1, L3)
        Δrˡˡˡ = Symbol(:Δr, L2, L3, L1)

        @eval @inline $Δxˡˡˡ(i, j, k, grid) = $Δxˡˡᵃ(i, j, k, grid)
        @eval @inline $Δyˡˡˡ(i, j, k, grid) = $Δyˡˡᵃ(i, j, k, grid)
        @eval @inline $Δzˡˡˡ(i, j, k, grid) = $Δzˡᵃˡ(i, j, k, grid)
        @eval @inline $Δλˡˡˡ(i, j, k, grid) = $Δλˡˡᵃ(i, j, k, grid)
        @eval @inline $Δφˡˡˡ(i, j, k, grid) = $Δφˡˡᵃ(i, j, k, grid)
        @eval @inline $Δrˡˡˡ(i, j, k, grid) = $Δrˡᵃˡ(i, j, k, grid)
    end
end

#####
##### One - dimensional Vertical spacing (same for all grids)
#####

@inline getspacing(k, Δz::Number) = Δz
@inline getspacing(k, Δz::AbstractVector) = @inbounds Δz[k]

# 3D AbstractArray getspacing: dispatch on Δz type to handle regular (Number) vs stretched (Vector)
@inline function getspacing_3d(i, j, k, Δ::Number)
    Base.fill(Δ, (length(Base.axes(i, 1)),
                   length(Base.axes(j, 2)),
                   length(Base.axes(k, 3))))
end

@inline function getspacing_3d(i, j, k, Δ::AbstractVector)
    z = @inbounds Δ[k]
    bc = Base.Broadcast.Broadcasted(
        Base.Broadcast.BroadcastStyle(typeof(z)),
        Base.identity,
        (Base.Broadcast.Extruded(z, (false, false, true), (1, 1, 1)),),
        (Base.axes(i, 1), Base.axes(j, 2), Base.axes(k, 3)))
    return Base.Broadcast.materialize(bc)
end

@inline Δrᵃᵃᶜ(i, j, k, grid) = getspacing(k, grid.z.Δᵃᵃᶜ)
@inline Δrᵃᵃᶠ(i, j, k, grid) = getspacing(k, grid.z.Δᵃᵃᶠ)

@inline Δzᵃᵃᶜ(i, j, k, grid) = getspacing(k, grid.z.Δᵃᵃᶜ)
@inline Δzᵃᵃᶜ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid) = getspacing_3d(i, j, k, grid.z.Δᵃᵃᶜ)

@inline Δzᵃᵃᶠ(i, j, k, grid) = getspacing(k, grid.z.Δᵃᵃᶠ)
@inline Δzᵃᵃᶠ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid) = getspacing_3d(i, j, k, grid.z.Δᵃᵃᶠ)

#####
#####
##### One - Dimensional Horizontal Spacings
#####
#####

#####
##### Rectilinear Grids (Flat grids already have Δ = 1)
#####

@inline Δxᶠᵃᵃ(i, j, k, grid::RG) = @inbounds grid.Δxᶠᵃᵃ[i]
@inline Δxᶜᵃᵃ(i, j, k, grid::RG) = @inbounds grid.Δxᶜᵃᵃ[i]

for sym in (:Δxᶠᵃᵃ, :Δxᶜᵃᵃ)
    @eval @inline function $sym(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::RG)
        x = $sym(i, 1, 1, grid)
        bc = Base.Broadcast.Broadcasted(
            Base.Broadcast.BroadcastStyle(typeof(x)),
            Base.identity,
            (Base.Broadcast.Extruded(x, (true, false, false), (1, 1, 1)),),
            (Base.axes(i, 1), Base.axes(j, 2), Base.axes(k, 3)))
        return Base.Broadcast.materialize(bc)
    end
end

@inline Δyᵃᶠᵃ(i, j, k, grid::RG) = @inbounds grid.Δyᵃᶠᵃ[j]
@inline Δyᵃᶜᵃ(i, j, k, grid::RG) = @inbounds grid.Δyᵃᶜᵃ[j]

for sym in (:Δyᵃᶠᵃ, :Δyᵃᶜᵃ)
    @eval @inline function $sym(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::RG)
        y = $sym(1, j, 1, grid)
        bc = Base.Broadcast.Broadcasted(
            Base.Broadcast.BroadcastStyle(typeof(y)),
            Base.identity,
            (Base.Broadcast.Extruded(y, (false, true, false), (1, 1, 1)),),
            (Base.axes(i, 1), Base.axes(j, 2), Base.axes(k, 3)))
        return Base.Broadcast.materialize(bc)
    end
end

### XRegularRG

@inline Δxᶠᵃᵃ(i, j, k, grid::RGX) = grid.Δxᶠᵃᵃ
@inline Δxᶜᵃᵃ(i, j, k, grid::RGX) = grid.Δxᶜᵃᵃ

for sym in (:Δxᶠᵃᵃ, :Δxᶜᵃᵃ)
    @eval @inline function $sym(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::RGX)
        val = $sym(1, 1, 1, grid)
        Base.fill(val, (length(Base.axes(i, 1)),
                        length(Base.axes(j, 2)),
                        length(Base.axes(k, 3))))
    end
end

### YRegularRG

@inline Δyᵃᶠᵃ(i, j, k, grid::RGY) = grid.Δyᵃᶠᵃ
@inline Δyᵃᶜᵃ(i, j, k, grid::RGY) = grid.Δyᵃᶜᵃ

for sym in (:Δyᵃᶠᵃ, :Δyᵃᶜᵃ)
    @eval @inline function $sym(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::RGY)
        val = $sym(1, 1, 1, grid)
        Base.fill(val, (length(Base.axes(i, 1)),
                        length(Base.axes(j, 2)),
                        length(Base.axes(k, 3))))
    end
end

#####
##### LatitudeLongitude Grids (define both precomputed and non-precomputed metrics)
#####

### Curvilinear spacings

@inline Δλᶜᵃᵃ(i, j, k, grid::LLG) = @inbounds grid.Δλᶜᵃᵃ[i]
@inline Δλᶠᵃᵃ(i, j, k, grid::LLG) = @inbounds grid.Δλᶠᵃᵃ[i]

for sym in (:Δλᶜᵃᵃ, :Δλᶠᵃᵃ)
    @eval @inline function $sym(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLG)
        x = $sym(i, 1, 1, grid)
        bc = Base.Broadcast.Broadcasted(Base.Broadcast.BroadcastStyle(typeof(x)),
                                        Base.identity,
                                        (Base.Broadcast.Extruded(x, (true, false, false), (1, 1, 1)),),
                                        (Base.axes(i, 1), Base.axes(j, 2), Base.axes(k, 3))) # i, j, k are 3D arrays
        return Base.Broadcast.materialize(bc)
    end
end
@inline Δλᶜᵃᵃ(i, j, k, grid::LLGX) = @inbounds grid.Δλᶜᵃᵃ
@inline Δλᶠᵃᵃ(i, j, k, grid::LLGX) = @inbounds grid.Δλᶠᵃᵃ

for sym in (:Δλᶜᵃᵃ, :Δλᶠᵃᵃ)
    @eval @inline function $sym(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLGX)
        val = $sym(1, 1, 1, grid)
        Base.fill(val, (length(Base.axes(i, 1)),
                        length(Base.axes(j, 2)),
                        length(Base.axes(k, 3))))
    end
end

@inline Δφᵃᶜᵃ(i, j, k, grid::LLG) = @inbounds grid.Δφᵃᶜᵃ[j]
@inline Δφᵃᶠᵃ(i, j, k, grid::LLG) = @inbounds grid.Δφᵃᶠᵃ[j]

for sym in (:Δφᵃᶜᵃ, :Δφᵃᶠᵃ)
    @eval @inline function $sym(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLG)
        y = $sym(1, j, 1, grid)
        bc = Base.Broadcast.Broadcasted(
            Base.Broadcast.BroadcastStyle(typeof(y)),
            Base.identity,
            (Base.Broadcast.Extruded(y, (false, true, false), (1, 1, 1)),),
            (Base.axes(i, 1), Base.axes(j, 2), Base.axes(k, 3)))
        return Base.Broadcast.materialize(bc)
    end
end
@inline Δφᵃᶜᵃ(i, j, k, grid::LLGY) = @inbounds grid.Δφᵃᶜᵃ
@inline Δφᵃᶠᵃ(i, j, k, grid::LLGY) = @inbounds grid.Δφᵃᶠᵃ

for sym in (:Δφᵃᶜᵃ, :Δφᵃᶠᵃ)
    @eval @inline function $sym(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLGY)
        val = $sym(1, 1, 1, grid)
        Base.fill(val, (length(Base.axes(i, 1)),
                        length(Base.axes(j, 2)),
                        length(Base.axes(k, 3))))
    end
end

### Linear spacings

### Precomputed metrics

@inline Δyᵃᶜᵃ(i, j, k, grid::LLGY) = grid.Δyᶠᶜᵃ
@inline Δyᵃᶠᵃ(i, j, k, grid::LLGY) = grid.Δyᶜᶠᵃ

for sym in (:Δyᵃᶜᵃ, :Δyᵃᶠᵃ)
    @eval @inline function $sym(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLGY)
        val = $sym(1, 1, 1, grid)
        return Base.fill(val, (length(Base.axes(i, 1)),
                               length(Base.axes(j, 2)),
                               length(Base.axes(k, 3)))) # return a 3D array
    end
end

@inline Δyᵃᶜᵃ(i, j, k, grid::LLG) = @inbounds grid.Δyᶠᶜᵃ[j]
@inline Δyᵃᶠᵃ(i, j, k, grid::LLG) = @inbounds grid.Δyᶜᶠᵃ[j]

for sym in (:Δyᵃᶜᵃ, :Δyᵃᶠᵃ)
    @eval @inline function $sym(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLG)
        y = $sym(1, j, 1, grid)
        bc = Base.Broadcast.Broadcasted(Base.Broadcast.BroadcastStyle(typeof(y)),
                                        Base.identity,
                                        (Base.Broadcast.Extruded(y, (false, true, false), (1, 1, 1)),), # y is a 1D array
                                        (Base.axes(i, 1), Base.axes(j, 2), Base.axes(k, 3))) # i, j, k are 3D arrays
        return Base.Broadcast.materialize(bc) # return a 3D array
    end
end

### On-the-fly metrics

@inline Δyᵃᶠᵃ(i, j, k, grid::LLGFY) = grid.radius * deg2rad(grid.Δφᵃᶠᵃ)
@inline Δyᵃᶜᵃ(i, j, k, grid::LLGFY) = grid.radius * deg2rad(grid.Δφᵃᶜᵃ)
@inline Δyᵃᶠᵃ(i, j, k, grid::LLGF)  = @inbounds grid.radius * deg2rad(grid.Δφᵃᶠᵃ[j])
@inline Δyᵃᶜᵃ(i, j, k, grid::LLGF)  = @inbounds grid.radius * deg2rad(grid.Δφᵃᶜᵃ[j])

#####
#####
##### Two-dimensional horizontal spacings
#####
#####

#####
##### LatitudeLongitudeGrid (only the Δx are required, Δy, Δλ, and Δφ are 1D)
#####

### Pre computed metrics

@inline Δxᶜᶠᵃ(i, j, k, grid::LLG) = @inbounds grid.Δxᶜᶠᵃ[i, j]
@inline Δxᶠᶜᵃ(i, j, k, grid::LLG) = @inbounds grid.Δxᶠᶜᵃ[i, j]
@inline Δxᶠᶠᵃ(i, j, k, grid::LLG) = @inbounds grid.Δxᶠᶠᵃ[i, j]
@inline Δxᶜᶜᵃ(i, j, k, grid::LLG) = @inbounds grid.Δxᶜᶜᵃ[i, j]

for sym in (:Δxᶜᶠᵃ, :Δxᶠᶜᵃ, :Δxᶠᶠᵃ, :Δxᶜᶜᵃ)
    @eval @inline function $sym(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLG)
        xy = $sym(i, j, 1, grid)
        bc = Base.Broadcast.Broadcasted(
            Base.Broadcast.BroadcastStyle(typeof(xy)),
            Base.identity,
            (Base.Broadcast.Extruded(xy, (true, true, false), (1, 1, 1)),),
            (Base.axes(i, 1), Base.axes(j, 2), Base.axes(k, 3)))
        return Base.Broadcast.materialize(bc)
    end
end


### XRegularLLG with pre computed metrics

@inline Δxᶠᶜᵃ(i, j, k, grid::LLGX) = @inbounds grid.Δxᶠᶜᵃ[j]
@inline Δxᶜᶠᵃ(i, j, k, grid::LLGX) = @inbounds grid.Δxᶜᶠᵃ[j]
@inline Δxᶠᶠᵃ(i, j, k, grid::LLGX) = @inbounds grid.Δxᶠᶠᵃ[j]
@inline Δxᶜᶜᵃ(i, j, k, grid::LLGX) = @inbounds grid.Δxᶜᶜᵃ[j]

for sym in (:Δxᶠᶜᵃ, :Δxᶜᶠᵃ, :Δxᶠᶠᵃ, :Δxᶜᶜᵃ)
    @eval @inline function $sym(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::LLGX)
        x = $sym(1, j, 1, grid)
        bc = Base.Broadcast.Broadcasted(
            Base.Broadcast.BroadcastStyle(typeof(x)),
            Base.identity,
            (Base.Broadcast.Extruded(x, (false, true, false), (1, 1, 1)),),
            (Base.axes(i, 1), Base.axes(j, 2), Base.axes(k, 3)))
        return Base.Broadcast.materialize(bc)
    end
end


### On-the-fly metrics

@inline Δxᶠᶜᵃ(i, j, k, grid::LLGF) = @inbounds grid.radius * deg2rad(grid.Δλᶠᵃᵃ[i]) * hack_cosd(grid.φᵃᶜᵃ[j])
@inline Δxᶜᶠᵃ(i, j, k, grid::LLGF) = @inbounds grid.radius * deg2rad(grid.Δλᶜᵃᵃ[i]) * hack_cosd(grid.φᵃᶠᵃ[j])
@inline Δxᶠᶠᵃ(i, j, k, grid::LLGF) = @inbounds grid.radius * deg2rad(grid.Δλᶠᵃᵃ[i]) * hack_cosd(grid.φᵃᶠᵃ[j])
@inline Δxᶜᶜᵃ(i, j, k, grid::LLGF) = @inbounds grid.radius * deg2rad(grid.Δλᶜᵃᵃ[i]) * hack_cosd(grid.φᵃᶜᵃ[j])

### XRegularLLG with on-the-fly metrics

@inline Δxᶠᶜᵃ(i, j, k, grid::LLGFX) = @inbounds grid.radius * deg2rad(grid.Δλᶠᵃᵃ) * hack_cosd(grid.φᵃᶜᵃ[j])
@inline Δxᶜᶠᵃ(i, j, k, grid::LLGFX) = @inbounds grid.radius * deg2rad(grid.Δλᶜᵃᵃ) * hack_cosd(grid.φᵃᶠᵃ[j])
@inline Δxᶠᶠᵃ(i, j, k, grid::LLGFX) = @inbounds grid.radius * deg2rad(grid.Δλᶠᵃᵃ) * hack_cosd(grid.φᵃᶠᵃ[j])
@inline Δxᶜᶜᵃ(i, j, k, grid::LLGFX) = @inbounds grid.radius * deg2rad(grid.Δλᶜᵃᵃ) * hack_cosd(grid.φᵃᶜᵃ[j])

#####
#####  OrthogonalSphericalShellGrid (does not have one-dimensional spacings)
#####

### Curvilinear spacings

@inline Δλᶜᶜᵃ(i, j, k, grid::OSSG) = δxᶜᵃᵃ(i, j, k, grid, λnode, Face(),   Center(), nothing)
@inline Δλᶠᶜᵃ(i, j, k, grid::OSSG) = δxᶠᵃᵃ(i, j, k, grid, λnode, Center(), Center(), nothing)
@inline Δλᶜᶠᵃ(i, j, k, grid::OSSG) = δxᶜᵃᵃ(i, j, k, grid, λnode, Face(),   Face(),   nothing)
@inline Δλᶠᶠᵃ(i, j, k, grid::OSSG) = δxᶠᵃᵃ(i, j, k, grid, λnode, Center(), Face(),   nothing)

@inline Δφᶜᶜᵃ(i, j, k, grid::OSSG) = δyᵃᶜᵃ(i, j, k, grid, λnode, Center(), Face(),   nothing)
@inline Δφᶠᶜᵃ(i, j, k, grid::OSSG) = δyᵃᶜᵃ(i, j, k, grid, λnode, Face(),   Face(),   nothing)
@inline Δφᶜᶠᵃ(i, j, k, grid::OSSG) = δyᵃᶠᵃ(i, j, k, grid, λnode, Center(), Center(), nothing)
@inline Δφᶠᶠᵃ(i, j, k, grid::OSSG) = δyᵃᶠᵃ(i, j, k, grid, λnode, Face(),   Center(), nothing)

### Linear spacings

@inline Δxᶜᶜᵃ(i, j, k, grid::OSSG) = @inbounds grid.Δxᶜᶜᵃ[i, j]

@inline Δxᶜᶜᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::OSSG) = Base.stack(collect(Δxᶜᶜᵃ(i, j, 1, grid) for _ in k))

@inline Δxᶠᶜᵃ(i, j, k, grid::OSSG) = @inbounds grid.Δxᶠᶜᵃ[i, j]
@inline Δxᶠᶜᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::OSSG) = Base.stack(collect(Δxᶠᶜᵃ(i, j, 1, grid) for _ in k))

@inline Δxᶜᶠᵃ(i, j, k, grid::OSSG) = @inbounds grid.Δxᶜᶠᵃ[i, j]
@inline Δxᶜᶠᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::OSSG) = Base.stack(collect(Δxᶜᶠᵃ(i, j, 1, grid) for _ in k))

@inline Δxᶠᶠᵃ(i, j, k, grid::OSSG) = @inbounds grid.Δxᶠᶠᵃ[i, j]
@inline Δxᶠᶠᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::OSSG) = Base.stack(collect(Δxᶠᶠᵃ(i, j, 1, grid) for _ in k))

@inline Δyᶜᶜᵃ(i, j, k, grid::OSSG) = @inbounds grid.Δyᶜᶜᵃ[i, j]
@inline Δyᶜᶜᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::OSSG) = Base.stack(collect(Δyᶜᶜᵃ(i, j, 1, grid) for _ in k))

@inline Δyᶠᶜᵃ(i, j, k, grid::OSSG) = @inbounds grid.Δyᶠᶜᵃ[i, j]
@inline Δyᶠᶜᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::OSSG) = Base.stack(collect(Δyᶠᶜᵃ(i, j, 1, grid) for _ in k))

@inline Δyᶜᶠᵃ(i, j, k, grid::OSSG) = @inbounds grid.Δyᶜᶠᵃ[i, j]
@inline Δyᶜᶠᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::OSSG) = Base.stack(collect(Δyᶜᶠᵃ(i, j, 1, grid) for _ in k))

@inline Δyᶠᶠᵃ(i, j, k, grid::OSSG) = @inbounds grid.Δyᶠᶠᵃ[i, j]
@inline Δyᶠᶠᵃ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid::OSSG) = Base.stack(collect(Δyᶠᶠᵃ(i, j, 1, grid) for _ in k))

#####
#####
##### Areas!!
#####
#####

# We do the same thing as for the spacings: define general areas and then specialize for each grid.
# Areas need to be at least 2D so we use the respective 2D spacings to define them.
for L1 in (:ᶜ, :ᶠ), L2 in (:ᶜ, :ᶠ)

    Δxˡˡᵃ = Symbol(:Δx, L1, L2, :ᵃ)
    Δxˡᵃˡ = Symbol(:Δx, L1, :ᵃ, L2)
    Δyˡˡᵃ = Symbol(:Δy, L1, L2, :ᵃ)
    Δyᵃˡˡ = Symbol(:Δy, :ᵃ, L1, L2)
    Δzˡᵃˡ = Symbol(:Δz, L1, :ᵃ, L2)
    Δzᵃˡˡ = Symbol(:Δz, :ᵃ, L1, L2)

    # 2D areas
    Axᵃˡˡ = Symbol(:Ax, :ᵃ, L1, L2)
    Ayˡᵃˡ = Symbol(:Ay, L1, :ᵃ, L2)
    Azˡˡᵃ = Symbol(:Az, L1, L2, :ᵃ)

    @eval begin
        @inline $Axᵃˡˡ(i, j, k, grid) = $Δyᵃˡˡ(i, j, k, grid) * $Δzᵃˡˡ(i, j, k, grid)
        @inline $Ayˡᵃˡ(i, j, k, grid) = $Δxˡᵃˡ(i, j, k, grid) * $Δzˡᵃˡ(i, j, k, grid)
        @inline $Azˡˡᵃ(i, j, k, grid) = $Δxˡˡᵃ(i, j, k, grid) * $Δyˡˡᵃ(i, j, k, grid)
    end

    for L3 in (:ᶜ, :ᶠ)
        # 3D spacings
        Δxˡˡˡ = Symbol(:Δx, L1, L2, L3)
        Δyˡˡˡ = Symbol(:Δy, L1, L2, L3)
        Δzˡˡˡ = Symbol(:Δz, L1, L2, L3)

        # 3D areas
        Axˡˡˡ = Symbol(:Ax, L1, L2, L3)
        Ayˡˡˡ = Symbol(:Ay, L1, L2, L3)
        Azˡˡˡ = Symbol(:Az, L1, L2, L3)

        @eval begin
            @inline $Axˡˡˡ(i, j, k, grid) = $Δyˡˡˡ(i, j, k, grid) * $Δzˡˡˡ(i, j, k, grid)
            @inline $Axˡˡˡ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid) = $Δyˡˡˡ(i, j, k, grid) .* $Δzˡˡˡ(i, j, k, grid)
            @inline $Ayˡˡˡ(i, j, k, grid) = $Δxˡˡˡ(i, j, k, grid) * $Δzˡˡˡ(i, j, k, grid)
            @inline $Ayˡˡˡ(i::AbstractArray, j::AbstractArray, k::AbstractArray, grid) = $Δxˡˡˡ(i, j, k, grid) .* $Δzˡˡˡ(i, j, k, grid)

            # For the moment the horizontal area is independent of `z`. This might change if
            # we want to implement deep atmospheres where Az is a function of z
            @inline $Azˡˡˡ(i, j, k, grid) = $Azˡˡᵃ(i, j, k, grid)
        end
    end
end

####
#### Special 2D z Areas for LatitudeLongitudeGrid and OrthogonalSphericalShellGrid
####

@inline Azᶠᶜᵃ(i, j, k, grid::LLGF)  = @inbounds grid.radius^2 * deg2rad(grid.Δλᶠᵃᵃ[i]) * (hack_sind(grid.φᵃᶠᵃ[j+1]) - hack_sind(grid.φᵃᶠᵃ[j]))
@inline Azᶜᶠᵃ(i, j, k, grid::LLGF)  = @inbounds grid.radius^2 * deg2rad(grid.Δλᶜᵃᵃ[i]) * (hack_sind(grid.φᵃᶜᵃ[j])   - hack_sind(grid.φᵃᶜᵃ[j-1]))
@inline Azᶠᶠᵃ(i, j, k, grid::LLGF)  = @inbounds grid.radius^2 * deg2rad(grid.Δλᶠᵃᵃ[i]) * (hack_sind(grid.φᵃᶜᵃ[j])   - hack_sind(grid.φᵃᶜᵃ[j-1]))
@inline Azᶜᶜᵃ(i, j, k, grid::LLGF)  = @inbounds grid.radius^2 * deg2rad(grid.Δλᶜᵃᵃ[i]) * (hack_sind(grid.φᵃᶠᵃ[j+1]) - hack_sind(grid.φᵃᶠᵃ[j]))

@inline Azᶠᶜᵃ(i, j, k, grid::LLGFX) = @inbounds grid.radius^2 * deg2rad(grid.Δλᶠᵃᵃ) * (hack_sind(grid.φᵃᶠᵃ[j+1]) - hack_sind(grid.φᵃᶠᵃ[j]))
@inline Azᶜᶠᵃ(i, j, k, grid::LLGFX) = @inbounds grid.radius^2 * deg2rad(grid.Δλᶜᵃᵃ) * (hack_sind(grid.φᵃᶜᵃ[j])   - hack_sind(grid.φᵃᶜᵃ[j-1]))
@inline Azᶠᶠᵃ(i, j, k, grid::LLGFX) = @inbounds grid.radius^2 * deg2rad(grid.Δλᶠᵃᵃ) * (hack_sind(grid.φᵃᶜᵃ[j])   - hack_sind(grid.φᵃᶜᵃ[j-1]))
@inline Azᶜᶜᵃ(i, j, k, grid::LLGFX) = @inbounds grid.radius^2 * deg2rad(grid.Δλᶜᵃᵃ) * (hack_sind(grid.φᵃᶠᵃ[j+1]) - hack_sind(grid.φᵃᶠᵃ[j]))

for LX in (:ᶠ, :ᶜ), LY in (:ᶠ, :ᶜ)

    z_area_2D = Symbol(:Az, LX, LY, :ᵃ)

    @eval begin
        @inline $z_area_2D(i, j, k, grid::OSSG) = @inbounds grid.$z_area_2D[i, j]
        @inline $z_area_2D(i, j, k, grid::LLG)  = @inbounds grid.$z_area_2D[i, j]
        @inline $z_area_2D(i, j, k, grid::LLGX) = @inbounds grid.$z_area_2D[j]
    end
end

#####
#####
##### Volumes!! (always 3D)
#####
#####

for LX in (:ᶠ, :ᶜ), LY in (:ᶠ, :ᶜ), LZ in (:ᶠ, :ᶜ)

    volume    = Symbol(:V,  LX, LY, LZ)
    z_area    = Symbol(:Az, LX, LY, LZ)
    z_spacing = Symbol(:Δz, LX, LY, LZ)

    @eval begin
        @inline $volume(i, j, k, grid) = $z_area(i, j, k, grid) * $z_spacing(i, j, k, grid)
    end
end

#####
##### Generic functions for specified locations
#####
##### For example, Δx(i, j, k, Face, Center, Face) is equivalent to = Δxᶠᶜᶠ(i, j, k, grid).
#####
##### We also use the function "volume" rather than `V`.
#####

function location_from_superscript(val::Symbol)
    if val == :ᶜ
        return :Center
    elseif val == :ᶠ
        return :Face
    else
        return :Nothing
    end
end

export Δx, Δy, Δz, Δλ, Δφ, Δr, Ax, Ay, Az

for ℓ1 in (:ᶜ, :ᶠ), ℓ2 in (:ᶜ, :ᶠ, :ᵃ), ℓ3 in (:ᶜ, :ᶠ, :ᵃ)

    L1 = location_from_superscript(ℓ1)
    L2 = location_from_superscript(ℓ2)
    L3 = location_from_superscript(ℓ3)

    spacing_x = Symbol(:Δx, ℓ1, ℓ2, ℓ3)
    spacing_λ = Symbol(:Δλ, ℓ1, ℓ2, ℓ3)
    spacing_y = Symbol(:Δy, ℓ2, ℓ1, ℓ3)
    spacing_φ = Symbol(:Δφ, ℓ2, ℓ1, ℓ3)
    spacing_z = Symbol(:Δz, ℓ2, ℓ3, ℓ1)
    spacing_r = Symbol(:Δr, ℓ2, ℓ3, ℓ1)

    @eval begin
        Δx(i, j, k, grid, ::$L1, ::$L2, ::$L3) = $spacing_x(i, j, k, grid)
        Δλ(i, j, k, grid, ::$L1, ::$L2, ::$L3) = $spacing_λ(i, j, k, grid)
        Δy(i, j, k, grid, ::$L2, ::$L1, ::$L3) = $spacing_y(i, j, k, grid)
        Δφ(i, j, k, grid, ::$L2, ::$L1, ::$L3) = $spacing_φ(i, j, k, grid)
        Δz(i, j, k, grid, ::$L2, ::$L3, ::$L1) = $spacing_z(i, j, k, grid)
        Δr(i, j, k, grid, ::$L2, ::$L3, ::$L1) = $spacing_r(i, j, k, grid)

        export $spacing_x, $spacing_λ, $spacing_y, $spacing_φ, $spacing_z, $spacing_r
    end
end

for ℓ1 in (:ᶜ, :ᶠ), ℓ2 in (:ᶜ, :ᶠ), ℓ3 in (:ᶜ, :ᶠ, :ᵃ)

    L1 = location_from_superscript(ℓ1)
    L2 = location_from_superscript(ℓ2)
    L3 = location_from_superscript(ℓ3)

    area_x = Symbol(:Ax, ℓ3, ℓ1, ℓ2)
    area_y = Symbol(:Ay, ℓ1, ℓ3, ℓ2)
    area_z = Symbol(:Az, ℓ1, ℓ2, ℓ3)

    @eval begin
        @eval Ax(i, j, k, grid, ::$L3, ::$L1, ::$L2) = $area_x(i, j, k, grid)
        @eval Ay(i, j, k, grid, ::$L1, ::$L3, ::$L2) = $area_y(i, j, k, grid)
        @eval Az(i, j, k, grid, ::$L1, ::$L2, ::$L3) = $area_z(i, j, k, grid)

        export $area_x, $area_y, $area_z
    end
end

for ℓ1 in (:ᶜ, :ᶠ), ℓ2 in (:ᶜ, :ᶠ), ℓ3 in (:ᶜ, :ᶠ)

    L1 = location_from_superscript(ℓ1)
    L2 = location_from_superscript(ℓ2)
    L3 = location_from_superscript(ℓ3)
    V  =  Symbol(:V, ℓ1, ℓ2, ℓ3)

    @eval begin
        @inline volume(i, j, k, grid, ::$L1, ::$L2, ::$L3) = $V(i, j, k, grid)
        export $V
    end
end

# One-dimensional convenience spacings (for grids that support them)

Δx(i, grid, ℓx) = Δx(i, 1, 1, grid, ℓx, nothing, nothing)
Δy(j, grid, ℓy) = Δy(1, j, 1, grid, nothing, ℓy, nothing)
Δz(k, grid, ℓz) = Δz(1, 1, k, grid, nothing, nothing, ℓz)
Δλ(i, grid, ℓx) = Δλ(i, 1, 1, grid, ℓx, nothing, nothing)
Δφ(j, grid, ℓy) = Δφ(1, j, 1, grid, nothing, ℓy, nothing)
Δr(k, grid, ℓz) = Δr(1, 1, k, grid, nothing, nothing, ℓz)

# Two-dimensional horizontal convenience spacings (for grids that support them)

Δx(i, j, grid, ℓx, ℓy) = Δx(i, j, 1, grid, ℓx, ℓy, nothing)
Δy(i, j, grid, ℓx, ℓy) = Δy(i, j, 1, grid, ℓx, ℓy, nothing)
Δλ(i, j, grid, ℓx, ℓy) = Δλ(i, j, 1, grid, ℓx, ℓy, nothing)
Δφ(i, j, grid, ℓx, ℓy) = Δφ(i, j, 1, grid, ℓx, ℓy, nothing)
