filtering.jl

module Filtering

export imfilter, imfilter!, trim_kernel

@inline function _shift_range(range, shift::Integer)
  Base.UnitRange(first(range)+shift : last(range)+shift)
end

## TODO: allow border to be tuple
#function _innershift_indices(a::Array{T, d}, shift::NTuple{d, Integer}, border = abs(maximum(shift))) where {T, d}
#  abs(maximum(shift)) <= border || throw("cannot shift by $shift with border $border")
#  range(i) = (1 + shift[i] + border):(lastindex(a, i) + shift[i] - border)
#  ntuple(range, Val(ndims(a)))
#end
#
#function innershift_view(a::Array{T, d}, shift::NTuple{d, Integer} = ntuple(i -> 0, d), border = abs(maximum(shift))) where {T, d}
#  view(a, _innershift_indices(a, shift, border)...)
#end

@inline function _bidx_to_range(range::AbstractUnitRange, bidx::Integer)
  if bidx > 0
    return last(range):last(range)
  elseif bidx < 0
    return first(range):first(range)
  else
    return first(range)+1:last(range)-1
  end
end


# we return Base.Slice since that is what `OffsetArray` uses for `axes()`
function _trim_kernel_axes(kaxes, bidx)
  ntuple(Val(length(kaxes))) do i
    if bidx[i] < 0
      return Base.Slice(0:last(kaxes[i]))
    elseif bidx[i] > 0
      return Base.Slice(first(kaxes[i]):0)
    else
      return Base.Slice(kaxes[i])
    end
  end
end

"""
    trim_kernel(kernel, bidx)

Trim `kernel` so that the resulting kernel is usable on the boundary
indicated by `bidx`
"""
function trim_kernel(kernel, bidx)
  dst_axes = _trim_kernel_axes(axes(kernel), bidx)
  out = similar(kernel, dst_axes)
  range = CartesianIndices(out)
  copyto!(out, range, kernel, range)
end



# TODO: weirdly Base.Slice instead of Base.UnitRange leads to axes that are not
# compatible when broadcasting, figure out why!
function imfilter!(dst, img, kern, slice=_bidx_to_range.(axes(img), ntuple(i->0, ndims(img))))
  ndims(dst) == ndims(img) == ndims(kern) == length(slice) ||
      throw(ArgumentError("dst, image, kernel or slice dimensions disagree ($(ndims(dst)) vs $(ndims(img)) vs $(ndims(kern)) vs $(length(slice)))"))
  axes(dst) == axes(img) || throw(ArgumentError("output and input image axes disagree"))
  all(size(kern) .<= size(img)) || throw(ArgumentError("kernel bigger than image"))
  last.(slice) .+ last.(axes(kern)) <= last.(axes(img)) || first.(slice) .+ first.(axes(kern)) >= first.(axes(img)) ||
      throw(ArgumentError("kern and/or slice out of bounds for image"))

  dstslice = view(dst, UnitRange.(slice)...)
  for i in CartesianIndices(kern)
    iszero(kern[i]) && continue
    dstslice .+= kern[i] .* view(img, (_shift_range.(slice, Tuple(i)))...)
  end
  dst
end

"""
    filter(kernelf, img)

Filter `img` using a kernel function `kernelf(bidx) = kernel` specifying a
kernel for each type of node.
"""
function imfilter(kernelf::Function, img)
  out = zeros(axes(img))
  for bidx in Iterators.product(ntuple(i -> -1:1, ndims(img))...)
    range = _bidx_to_range.(axes(img), bidx)
    imfilter!(out, img, kernelf(bidx), range)
  end
  out
end


"""
    imfilter(img, kern)

Filter `img` by `kern` using correlation.
Boundary is zeroed out.
"""
imfilter(img, kern) = imfilter(bidx -> trim_kernel(kern, bidx), img)
# TODO: use keyword arguments
imfilter_inner(img, kern) = imfilter(bidx -> all(bidx .== 0) ? kern : zeros(ntuple(i -> 0, ndims(img))), img)

end