using Oceananigans.Grids: AbstractGrid
using Oceananigans.Fields: AbstractField, instantiated_location

const unary_operators = Set()

"""
    UnaryOperation{LX, LY, LZ}(op, arg, ▶, grid)

Returns an abstract `UnaryOperation` representing the action of `op` on `arg`,
and subsequent interpolation by `▶` on `grid`.
"""
struct UnaryOperation{LX, LY, LZ, O, A, IN, G, T} <: AbstractOperation{LX, LY, LZ, G, T}
    op :: O
    arg :: A
    ▶ :: IN
    grid :: G

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

@inline Base.getindex(υ::UnaryOperation, i, j, k) = υ.▶(i, j, k, υ.grid, υ.op, υ.arg)

#####
##### UnaryOperation construction
#####

indices(υ::UnaryOperation) = indices(υ.arg)

"""Create a unary operation for `operator` acting on `arg` which interpolates the
result from `Larg` to `L`."""
function _unary_operation(L::Tuple{LX, LY, LZ}, operator, arg, Larg, grid) where {LX, LY, LZ}
    arg = validate_operand(arg)
    ▶ = interpolation_operator(Larg, L)
    return UnaryOperation{LX, LY, LZ}(operator, arg, ▶, grid)
end

# Recompute location of unary operation
@inline at(loc, υ::UnaryOperation) = υ.op(loc, at(loc, υ.arg))

"""
    @unary op1 op2 op3...

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

Note: a unary function is a function with one argument: for example, `sin(x)` is a unary function.

Also note: a unary function in `Base` must be imported to be extended: use `import Base: op; @unary op`.

Example
=======

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

julia> square_it(x) = x^2
square_it (generic function with 1 method)

julia> @unary square_it
Set{Any} with 15 elements:
  :+
  :log10
  :interpolate_identity
  :cos
  :exp
  :tanh
  :abs
  :log
  :cosh
  :square_it
  :-
  :sqrt
  :tan
  :sinh
  :sin

julia> c = CenterField(RectilinearGrid(size=(1, 1, 1), extent=(1, 1, 1)));

julia> square_it(c)
UnaryOperation at (Center, Center, Center)
├── grid: 1×1×1 RectilinearGrid{Float64, Periodic, Periodic, Bounded} on CPU with 1×1×1 halo
└── tree:
    square_it at (Center, Center, Center) via identity
    └── 1×1×1 Field{Center, Center, Center} on RectilinearGrid on CPU
```
"""
macro unary(ops...)
    expr = Expr(:block)

    for op in ops
        define_unary_operator = quote
            @inline $op(i, j, k, grid::$(AbstractGrid), a) = @inbounds $op(a[i, j, k])
            @inline $op(i, j, k, grid::$(AbstractGrid), a::Number) = $op(a)

            """
                $($op)(Lop::Tuple, a::AbstractField)

            Returns an abstract representation of the operator `$($op)` acting on the Oceananigans `Field`
            `a`, and subsequently interpolated to the location indicated by `Lop`.
            """
            function $op(Lop::Tuple{<:$Location, <:$Location, <:$Location}, a::$(AbstractField))
                L = $(instantiated_location)(a)
                return $(_unary_operation)(Lop, $op, a, L, a.grid)
            end

            # instantiate location if types are passed
            $op(Lc::Tuple, a::$(AbstractField)) = $op((Lc[1](), Lc[2](), Lc[3]()), a)

            $op(a::$(AbstractField)) = $op($(instantiated_location)(a), a)

            push!($(operators), Symbol($op))
            push!($(unary_operators), Symbol($op))
        end

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

    return expr
end

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

compute_at!(υ::UnaryOperation, time) = compute_at!(υ.arg, time)

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

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

Architectures.on_architecture(to, unary::UnaryOperation{LX, LY, LZ}) where {LX, LY, LZ} =
    UnaryOperation{LX, LY, LZ}(on_architecture(to, unary.op),
                               on_architecture(to, unary.arg),
                               on_architecture(to, unary.▶),
                               on_architecture(to, unary.grid),
                               eltype(unary))
