const multiary_operators = Set()

struct MultiaryOperation{LX, LY, LZ, N, O, A, IN, G, T} <: AbstractOperation{LX, LY, LZ, G, T}
    op :: O
    args :: A
    ▶ :: IN
    grid :: G

    function MultiaryOperation{LX, LY, LZ}(op::O, args::A, ▶::IN, grid::G,
                                            ::Type{T}=Base.promote_op(op, map(eltype, args)...)) where {LX, LY, LZ, O, A, IN, G, T}
        N = length(args)
        return new{LX, LY, LZ, N, O, A, IN, G, T}(op, args, ▶, grid)
    end
end

# NOTE!!! The GPU compiler has trouble inferring this operation which might lead to dynamic evaluations.
@inline Base.getindex(Π::MultiaryOperation{LX, LY, LZ, N}, i, j, k)  where {LX, LY, LZ, N} =
    Π.op(ntuple(γ -> Π.▶[γ](i, j, k, Π.grid, Π.args[γ]), Val(N))...)

# Try to improve inferrability by hardcoding methods for small N
@inline Base.getindex(Π::MultiaryOperation{LX, LY, LZ, 1}, i, j, k)  where {LX, LY, LZ} =
    Π.op(Π.▶[1](i, j, k, Π.grid, Π.args[1]))

@inline Base.getindex(Π::MultiaryOperation{LX, LY, LZ, 2}, i, j, k)  where {LX, LY, LZ} =
    Π.op(Π.▶[1](i, j, k, Π.grid, Π.args[1]),
         Π.▶[2](i, j, k, Π.grid, Π.args[2]))

@inline Base.getindex(Π::MultiaryOperation{LX, LY, LZ, 3}, i, j, k)  where {LX, LY, LZ} =
    Π.op(Π.▶[1](i, j, k, Π.grid, Π.args[1]),
         Π.▶[2](i, j, k, Π.grid, Π.args[2]),
         Π.▶[3](i, j, k, Π.grid, Π.args[3]))

@inline Base.getindex(Π::MultiaryOperation{LX, LY, LZ, 4}, i, j, k)  where {LX, LY, LZ} =
    Π.op(Π.▶[1](i, j, k, Π.grid, Π.args[1]),
         Π.▶[2](i, j, k, Π.grid, Π.args[2]),
         Π.▶[3](i, j, k, Π.grid, Π.args[3]),
         Π.▶[4](i, j, k, Π.grid, Π.args[4]))

@inline Base.getindex(Π::MultiaryOperation{LX, LY, LZ, 5}, i, j, k)  where {LX, LY, LZ} =
    Π.op(Π.▶[1](i, j, k, Π.grid, Π.args[1]),
         Π.▶[2](i, j, k, Π.grid, Π.args[2]),
         Π.▶[3](i, j, k, Π.grid, Π.args[3]),
         Π.▶[4](i, j, k, Π.grid, Π.args[4]),
         Π.▶[5](i, j, k, Π.grid, Π.args[5]))

@inline Base.getindex(Π::MultiaryOperation{LX, LY, LZ, 6}, i, j, k)  where {LX, LY, LZ} =
    Π.op(Π.▶[1](i, j, k, Π.grid, Π.args[1]),
         Π.▶[2](i, j, k, Π.grid, Π.args[2]),
         Π.▶[3](i, j, k, Π.grid, Π.args[3]),
         Π.▶[4](i, j, k, Π.grid, Π.args[4]),
         Π.▶[5](i, j, k, Π.grid, Π.args[5]),
         Π.▶[6](i, j, k, Π.grid, Π.args[6]))

#####
##### MultiaryOperation construction
#####

indices(Π::MultiaryOperation) = construct_regionally(intersect_indices, location(Π), Π.args...)

function _multiary_operation(L::Tuple{LX, LY, LZ}, op, args, Largs, grid) where {LX, LY, LZ}
    args = map(validate_operand, args)
    ▶ = Tuple(interpolation_operator(La, L) for La in Largs)
    return MultiaryOperation{LX, LY, LZ}(op, Tuple(a for a in args), ▶, grid)
end

# Recompute location of multiary operation
@inline at(loc, Π::MultiaryOperation) = Π.op(loc, Tuple(at(loc, a) for a in Π.args)...)

"""Return an expression that defines an abstract `MultiaryOperator` named `op` for `AbstractField`."""
function define_multiary_operator(op)
    return quote
        function $op(Lop::Tuple{<:$Location, <:$Location, <:$Location},
                     a::Union{Function, Number, Oceananigans.Fields.AbstractField},
                     b::Union{Function, Number, Oceananigans.Fields.AbstractField},
                     c::Union{Function, Number, Oceananigans.Fields.AbstractField},
                     d::Union{Function, Number, Oceananigans.Fields.AbstractField}...)

            args = tuple(a, b, c, d...)
            grid = $(validate_grid)(args...)

            # Convert any functions to FunctionFields
            args = Tuple(Oceananigans.Fields.fieldify_function(Lop, a, grid) for a in args)
            Largs = Tuple(Oceananigans.Fields.location(a) for a in args)

            return $(_multiary_operation)(Lop, $op, args, Largs, grid)
        end

        # Instantiate location if types are passed
        $op(Lop::Tuple,
            a::Union{Function, Number, Oceananigans.Fields.AbstractField},
            b::Union{Function, Number, Oceananigans.Fields.AbstractField},
            c::Union{Function, Number, Oceananigans.Fields.AbstractField},
            d::Union{Function, Number, Oceananigans.Fields.AbstractField}...) = $op((Lop[1](), Lop[2](), Lop[3]()), a, b, c, d...)

        $op(a::Oceananigans.Fields.AbstractField,
            b::Union{Function, Number, Oceananigans.Fields.AbstractField},
            c::Union{Function, Number, Oceananigans.Fields.AbstractField},
            d::Union{Function, Number, Oceananigans.Fields.AbstractField}...) = $op(Oceananigans.Fields.instantiated_location(a), a, b, c, d...)
    end
end

"""
    @multiary op1 op2 op3...

Turn each multiary operator in the list `(op1, op2, op3...)`
into a multiary operator on `Oceananigans.Fields` for use in `AbstractOperations`.

Note that a multiary operator:
  * is a function with two or more arguments: for example, `+(x, y, z)` is a multiary function;
  * must be imported to be extended if part of `Base`: use `import Base: op; @multiary op`;
  * can only be called on `Oceananigans.Field`s if the "location" is noted explicitly; see example.

Example
=======

```jldoctest
julia> using Oceananigans, Oceananigans.AbstractOperations

julia> harmonic_plus(a, b, c) = 1/3 * (1/a + 1/b + 1/c)
harmonic_plus (generic function with 1 method)

julia> c, d, e = Tuple(CenterField(RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1))) for i = 1:3);

julia> harmonic_plus(c, d, e) # before magic @multiary transformation
BinaryOperation at (Center, Center, Center)
├── grid: 1×1×1 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×1 halo
└── tree:
    * at (Center, Center, Center)
    ├── 0.3333333333333333
    └── + at (Center, Center, Center)
        ├── / at (Center, Center, Center)
        │   ├── 1
        │   └── 1×1×1 Field{Center, Center, Center} on RectilinearGrid on CPU
        ├── / at (Center, Center, Center)
        │   ├── 1
        │   └── 1×1×1 Field{Center, Center, Center} on RectilinearGrid on CPU
        └── / at (Center, Center, Center)
            ├── 1
            └── 1×1×1 Field{Center, Center, Center} on RectilinearGrid on CPU

julia> @multiary harmonic_plus
Set{Any} with 3 elements:
  :+
  :harmonic_plus
  :*

julia> harmonic_plus(c, d, e)
MultiaryOperation at (Center, Center, Center)
├── grid: 1×1×1 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×1 halo
└── tree:
    harmonic_plus at (Center, Center, Center)
    ├── 1×1×1 Field{Center, Center, Center} on RectilinearGrid on CPU
    ├── 1×1×1 Field{Center, Center, Center} on RectilinearGrid on CPU
    └── 1×1×1 Field{Center, Center, Center} on RectilinearGrid on CPU
```
"""
macro multiary(ops...)
    expr = Expr(:block)

    for op in ops
        defexpr = define_multiary_operator(op)
        push!(expr.args, :($(esc(defexpr))))

        add_to_operator_lists = quote
            push!($(operators), Symbol($op))
            push!($(multiary_operators), Symbol($op))
        end

        push!(expr.args, :($(esc(add_to_operator_lists))))
    end

    return expr
end

#####
##### Nested computations
#####

function compute_at!(Π::MultiaryOperation, time)
    for a in Π.args
        compute_at!(a, time)
    end
    return Π
end

#####
##### GPU capabilities
#####

"Adapt `MultiaryOperation` to work on the GPU via KernelAbstractions."
Adapt.adapt_structure(to, multiary::MultiaryOperation{LX, LY, LZ}) where {LX, LY, LZ} =
    MultiaryOperation{LX, LY, LZ}(Adapt.adapt(to, multiary.op),
                                  Adapt.adapt(to, multiary.args),
                                  Adapt.adapt(to, multiary.▶),
                                  Adapt.adapt(to, multiary.grid),
                                  eltype(multiary))

Architectures.on_architecture(to, multiary::MultiaryOperation{LX, LY, LZ}) where {LX, LY, LZ} =
    MultiaryOperation{LX, LY, LZ}(on_architecture(to, multiary.op),
                                  on_architecture(to, multiary.args),
                                  on_architecture(to, multiary.▶),
                                  on_architecture(to, multiary.grid),
                                  eltype(multiary))
