## Scatter API
#   - Scatter:
#     - scatter(op, src, idx)
#     - scatter!(op, dst, src, idx)
#   - Scatter destination backpropagation
#     - ∇scatter!_dst
#   - Scatter source backpropagation
#     - ∇scatter_src
#     - ∇scatter!_src
#

typelength(::Type{<:Number}) = 1
typelength(::Type{<:NTuple{M}}) where M = M
typelength(::Type{CartesianIndex{M}}) where M = M

"""
Performs dimensional consistency checks and return the
dimensionality of the scattered objects.
"""
function scatter_dims(
    X::AbstractArray{Tx,Nx}, Y::AbstractArray{Ty,Ny},
    idx::AbstractArray{Tidx,Nidx},
) where {Tx,Ty,Tidx,Nx,Ny,Nidx}
    dims = scatter_dims(Nx, Ny, typelength(Tidx), Nidx)
    size(X)[1:dims] == size(Y)[1:dims] || throw(ArgumentError("Incompatible input shapes."))
    size(Y)[dims+1:end] == size(idx) || throw(ArgumentError("Incompatible input shapes."))
    return dims
end

function scatter_dims(Nx, Ny, M, Nidx)
    @assert Nx - M == Ny - Nidx "Incompatible input shapes of (dst, src, idx) = ($Nx, $Ny, $Nidx)."
    dims = Nx - M
    dims < 0 && throw(ArgumentError("dims must be non-negative but got dims=$dims."))
    return dims
end

_view(X, colons, k) = view(X, colons..., k...)
_view(X, colons, k::Union{Integer, CartesianIndex}) = view(X, colons..., k)

# Move `idx` onto the same backend as the GPU array `ref`, so that `gather`/`scatter`
# work when the source/destination is on the GPU but the index was left on the CPU.
# Returns `idx` unchanged when it already lives on a GPU (issue #415).
_to_backend(ref::AnyGPUArray, idx::AnyGPUArray) = idx
function _to_backend(ref::AnyGPUArray, idx::AbstractArray)
    idx_dev = similar(ref, eltype(idx), size(idx))
    return copyto!(idx_dev, idx)
end

# Check that every entry of `idx` is within the `sz` trailing dimensions it indexes
# into. The CPU `gather!`/`scatter!` get this for free from `view`, but the GPU
# kernels use `@inbounds`, so an out-of-range index would otherwise silently corrupt
# memory instead of erroring (issue #416). Runs on the CPU or the GPU (one reduction).
function checkbounds_idx(idx::AbstractArray, sz::Dims)
    isempty(idx) || all(k -> _idx_inbounds(k, sz), idx) || throw(ArgumentError(
        "out-of-bounds index in `idx`: every index must address the trailing size $sz"))
    return nothing
end

@inline _idx_inbounds(k::Integer, sz::Dims{1}) = 1 ≤ k ≤ sz[1]
@inline _idx_inbounds(k::CartesianIndex, sz::Dims) = _idx_inbounds(Tuple(k), sz)
@inline function _idx_inbounds(k::NTuple{M,<:Integer}, sz::Dims{M}) where M
    inbounds = true
    for d in 1:M
        inbounds &= 1 ≤ k[d] ≤ sz[d]
    end
    return inbounds
end

"""
    NNlib.scatter!(op, dst, src, idx)

Scatter operation, which writes data in `src` into `dst` at locations `idx`.
A binary reduction operator `op` is applied during the scatter.
For each index `k` in `idx`, accumulates values in `dst` according to

    dst[:, ..., idx[k]...] = (op).(dst[:, ..., idx[k]...], src[:, ..., k...])

See also [`scatter`](@ref), [`gather`](@ref).

# Arguments

- `op`: Operations to be applied on `dst` and `src`, e.g. `+`, `-`, `*`, `/`, `max`, `min` and `mean`.
- `dst`: The destination for `src` to aggregate to. This argument will be mutated.
- `src`: The source data for aggregating.
- `idx`: The mapping for aggregation from source (index) to destination (value).
         The `idx` array can contain either integers or tuples.

# Examples
```jldoctest
julia> NNlib.scatter!(+, ones(3), [10,100], [1,3])
3-element Vector{Float64}:
  11.0
   1.0
 101.0

julia> NNlib.scatter!(*, fill(0.5, 2, 4), [1 10; 100 1000], [3,2])
2×4 Matrix{Float64}:
 0.5    5.0   0.5  0.5
 0.5  500.0  50.0  0.5
```
"""
function scatter!(op::OP, dst::AbstractArray, src::AbstractArray, idx::AbstractArray) where OP
    dims = scatter_dims(dst, src, idx)
    colons = Base.ntuple(_->Colon(), dims)
    for k in CartesianIndices(idx)
        dst_v = _view(dst, colons, idx[k])
        src_v = _view(src, colons, k)
        dst_v .= (op).(dst_v, src_v)
    end
    dst
end

for AT in (AbstractArray, AnyGPUArray)
    @eval function scatter!(op::typeof(mean), dst::$AT, src::$AT, idx::$AT)
        Ns = scatter!(+, zero(dst), one.(src), idx)
        dst_ = scatter!(+, zero(dst), src, idx)
        dst .+= safe_div.(dst_, Ns)
        return dst
    end
end

function scatter!(op::OP, dst::AnyGPUArray, src::AnyGPUArray, idx::AnyGPUArray) where OP
    n_dims = scatter_dims(dst, src, idx)
    checkbounds_idx(idx, size(dst)[n_dims+1:end])
    args = if n_dims == 0
        ndrange = length(idx)
        ()
    else
        dims = size(dst)[1:n_dims]
        max_dims_idx = prod(dims)
        ndrange = max_dims_idx * length(idx)
        (CartesianIndices(dims), max_dims_idx)
    end
    _scatter!(KernelAbstractions.get_backend(dst))(
        op, dst, src, idx, args...; ndrange)
    dst
end

# Destination on the GPU, index on the CPU: move the index to the GPU and dispatch to
# the kernel above, rather than falling back to slow scalar indexing (issue #415). The
# `mean` method resolves dispatch ambiguity with the `op::typeof(mean)` methods above.
scatter!(op::OP, dst::AnyGPUArray, src::AnyGPUArray, idx::AbstractArray) where OP =
    scatter!(op, dst, src, _to_backend(dst, idx))
scatter!(op::typeof(mean), dst::AnyGPUArray, src::AnyGPUArray, idx::AbstractArray) =
    scatter!(op, dst, src, _to_backend(dst, idx))

@kernel function _scatter!(op::OP, dst, src, idxs) where OP
    i = @index(Global)
    @inbounds idx = Tuple(_convert_i64(idxs[i]))
    @inbounds _atomic_scatter!(dst, idx, op, src[i])
end

@kernel function _scatter!(
    op::OP, dst, src, idxs, dim_ids::CartesianIndices, max_dims_idx::Int,
) where OP
    i = @index(Global)
    j, k = divrem(i - 1, max_dims_idx)
    @inbounds idx = (Tuple(dim_ids[k + 1])..., Tuple(_convert_i64(idxs[j + 1]))...)
    @inbounds _atomic_scatter!(dst, idx, op, src[i])
end

# Atomically perform `dst[idx...] = op(dst[idx...], val)` inside a scatter kernel.
# The default routes through Atomix, which covers the full op set on CUDA and
# AMDGPU. Metal is special-cased in NNlibMetalExt because Atomix's Metal atomics
# only support a subset of ops (e.g. they error on float `max`/`min` and silently
# no-op on `*`/`/`); see https://github.com/FluxML/NNlib.jl/issues/534.
@inline _atomic_scatter!(dst, idx, op::OP, val) where OP =
    Atomix.modify!(Atomix.IndexableRef(dst, idx), op, val)

# Allow non-Int64 indices by converting them to Int64 when index eltype <: Integer.
# All other index types (tuples, cartesian indices) must be in Int64 already.
@inline _convert_i64(x::Int) = x
@inline _convert_i64(x::Integer) = Int(x)
@inline _convert_i64(x) = x

"""
    NNlib.scatter(op, src, idx; [init, dstsize])

Scatter operation allocating a destination array `dst` and
calling `scatter!(op, dst, src, idx)` on it.

* If keyword `init` is provided, it is used to initialize the content of `dst`.
  Otherwise, the init values is inferred from the reduction operator `op`
  for some common operators (e.g. `init = 0` for `op = +`).

* If `dstsize` is provided, it will be used to define the size of
  destination array, otherwise it will be inferred by `src` and `idx`.

See [`scatter!`](@ref) for full details on how `idx` works.

# Examples
```jldoctest
julia> NNlib.scatter(+, [10,100,1000], [3,1,2])
3-element Vector{Int64}:
  100
 1000
   10

julia> NNlib.scatter(+, [1 2 3 4; 5 6 7 8], [2,1,1,5])
2×5 Matrix{Int64}:
  5  1  0  0  4
 13  5  0  0  8

julia> NNlib.scatter(*, [10,200,3000], [1,4,2]; init = 10, dstsize = 6)
6-element Vector{Int64}:
   100
 30000
    10
  2000
    10
    10
```
"""
function scatter(
    op::OP, src::AbstractArray{Tsrc,Nsrc}, idx::AbstractArray{Tidx,Nidx};
    init = nothing, dstsize = nothing,
) where {Tsrc,Tidx,Nsrc,Nidx,OP}
    dims = Nsrc - Nidx
    dstsz = isnothing(dstsize) ? (size(src)[1:dims]..., maximum_dims(idx)...) : dstsize
    dst = similar(src, Tsrc, dstsz)
    xinit = isnothing(init) ? scatter_empty(op, Tsrc) : init
    fill!(dst, xinit)
    scatter!(op, dst, src, idx)
end

scatter_empty(op, T) = Base.reduce_empty(op, T)
scatter_empty(op::typeof(-), T) = zero(T)
scatter_empty(op::typeof(/), T) = one(T)
scatter_empty(op::typeof(min), T) = typemax(T)
scatter_empty(op::typeof(max), T) = typemin(T)
scatter_empty(op::typeof(mean), T) = zero(T)

## Gradients

∇scatter!_src(op, Δ, dst, src, idx) = ∇scatter_src(op, Δ, dst, src, idx)
∇scatter!_src(op::Union{typeof(*),typeof(/)}, Δ, dst, src, idx) =
    gather(dst, idx) .* ∇scatter_src(op, Δ, dst, src, idx)
∇scatter!_dst(op, Δ, dst, y) = Δ
∇scatter!_dst(op::Union{typeof(max),typeof(min)}, Δ, dst_old, dst) =
    (dst_old .== op.(dst_old, dst)) .* Δ

modify_src(::typeof(+), X) = X
modify_src(::typeof(-), X) = -X
modify_src(::typeof(*), X, Y) = X
modify_src(::typeof(/), X, Y) = .-X ./ Y.^2

∇scatter_src(op::Union{typeof(+),typeof(-)}, Δ, dst, src, idx) =
    modify_src(op, gather(Δ, idx))
∇scatter_src(::Union{typeof(max),typeof(min)}, Δ, dst, src, idx) =
    (src .== gather(dst, idx)) .* gather(Δ, idx)

function ∇scatter_src(
    op::Union{typeof(*),typeof(/)}, Δ, dst,
    src::AbstractArray{Tsrc,Nsrc}, idx::AbstractArray{Tidx,Nidx},
) where {Tsrc,Tidx,Nsrc,Nidx}
    dims = Nsrc - Nidx
    Δsrc = modify_src(op, gather(Δ, idx), src)
    rev_idx = reverse_indices(idx)
    ax = CartesianIndices(axes(src)[1:dims])
    for k in CartesianIndices(idx)
        inds = filter(x -> x != k, rev_idx[idx[k]])
        for i in ax
            # `init` handles values mapped by a single source position, where
            # `inds` is empty and the product of siblings is the identity.
            Δsrc[i, k] = op(Δsrc[i, k], prod(j -> src[i, j], inds; init=one(Tsrc)))
        end
    end
    Δsrc
end

# The product-rule gradient needs, for each source position, the product of its
# "siblings" (the other source values scattering to the same destination). That
# equals the product over the whole group divided by the element itself, and the
# group product is just a forward `*`-scatter. Expressing it with
# `gather`/`scatter`/broadcast keeps everything on the device and avoids the
# ragged reverse-index buffer, which cannot be compiled on backends such as Metal.
function ∇scatter_src(
    op::Union{typeof(*), typeof(/)}, Δ, dst,
    src::AnyGPUArray{Tsrc, Nsrc}, idx::AnyGPUArray{Tidx, Nidx},
) where {Tsrc, Nsrc, Tidx, Nidx}
    Δsrc = NNlib.modify_src(op, NNlib.gather(Δ, idx), src)
    prod_siblings = NNlib.gather(NNlib.scatter(*, src, idx), idx) ./ src
    return op.(Δsrc, prod_siblings)
end

function ∇scatter_src(
    ::typeof(mean), Δ, dst,
    src::AbstractArray{Tsrc,Nsrc}, idx::AbstractArray{Tidx,Nidx},
) where {Tsrc,Tidx,Nsrc,Nidx}
    M = typelength(Tidx)
    num = gather(Δ, idx)
    # `Int32` (rather than `Int`) keeps the count buffer compatible with backends
    # that lack 64-bit atomics, e.g. Metal. Counts are small, so this never overflows.
    counts = fill!(similar(Δ, Int32, size(Δ)[end-M+1:end]), 0)
    scatter!(+, counts, fill!(similar(idx, Int32), 1), idx)
    den = gather(counts, idx)
    # make num and den broadcast compatible
    for i in 1:ndims(num)-ndims(den)
        den = unsqueeze(den)
    end
    return safe_div.(num, den)
end

∇scatter_src(op, Δ, dst, src, idx) = ∇scatter_src(op, unthunk(Δ), dst, src, idx)

function rrule(::typeof(scatter!), op, dst::AbstractArray, src::AbstractArray, idx::AbstractArray)
    dst_old = copy(dst)
    scatter!(op, dst, src, idx)
    scatter!_pullback(Δ) = (NoTangent(), NoTangent(), ∇scatter!_dst(op, unthunk(Δ), dst_old, dst), ∇scatter!_src(op, unthunk(Δ), dst, src, idx), NoTangent())
    dst, scatter!_pullback
end

function rrule(::typeof(scatter), op, src::AbstractArray, idx::AbstractArray; kws...)
    y = scatter(op, src, idx; kws...)
    scatter_pullback(Δ) = (NoTangent(), NoTangent(), ∇scatter_src(op, unthunk(Δ), y, src, idx), NoTangent())
    y, scatter_pullback
end
