warp.jl

using Interpolations, ImageTransformations, StaticArrays, OffsetArrays, DiskArrays
using ProgressMeter, Morton
import Interpolations: degree, value_weights
import OffsetArrays: IdOffsetRange

"""
    Lanczos4OpenCVFaithful()

Similar to Lanczos4OpenCV in Interpolations.jl, but does *exactly* what OpenCV
does, and explicitly does not try to make it better.  This rendition also works
with Float32.
"""
struct Lanczos4OpenCVFaithful{T} <: Interpolations.AbstractLanczos end
Lanczos4OpenCVFaithful() = Lanczos4OpenCVFaithful{Float64}()

degree(::Lanczos4OpenCVFaithful) = 4

value_weights(::Lanczos4OpenCVFaithful{T}, δx) where {T} =
        _lanczos4_opencv(float(T), float(T).(l4_2d_cs), δx)

# s45 = sqrt(2)/2
const s45 = 0.70710678118654752440084436210485

# l4_2d_cs is a lookup table that could be generated by
# x = (0:7)*45*5
# l4_2d_cs = [cosd.(x) sind.(x)]'
const l4_2d_cs = SA[1 0; -s45 -s45; 0 1; s45 -s45; -1 0; s45 s45; 0 -1; -s45 s45]

function _lanczos4_opencv(::Type{F}, l4_2d_cs, δx) where {F}
    p_4 = π * F(0.25)
    y0 = -(δx + F(3)) * p_4
    s0, c0 = sincos(y0)
    cs = ntuple(8) do i
        y = δx + 4 - i
        if abs(y) >= F(1e-6)
            y *= p_4
            (F(l4_2d_cs[i, 1]) * s0 + F(l4_2d_cs[i, 2]) * c0) / y^2
        else
            F(1e30)
        end
    end
    sum_cs = sum(cs)
    normed_cs = ntuple(i -> cs[i] / sum_cs, Val(8))
    return normed_cs
end

parentindex(r::IdOffsetRange, i::Integer) = i - r.offset
parentindex(r::IdOffsetRange, i::AbstractRange) = (first(i) - r.offset):(last(i) - r.offset)
parentindex(r::IdOffsetRange, ::Colon) = Colon()

@inline function Base.getindex(A::OffsetArray, I::Union{Colon, Int, AbstractRange}...)
    @boundscheck checkbounds(A, I...)
    J = map(parentindex, axes(A), I)
    @inbounds parent(A)[J...]
end

@inline function Base.setindex!(A::OffsetArray, val, I::Union{Colon, Int, AbstractRange}...)
    @boundscheck checkbounds(A, I...)
    J = map(parentindex, axes(A), I)
    @inbounds parent(A)[J...] = val
    A
end

function warp_chunk(chunk, warped, invmap_xextrema, invmap_yextrema, tform, ccurr, zs)
    out = Array{eltype(warped)}(undef, length(chunk[1]), length(chunk[2]), length(zs))
    for (iz,z) in enumerate(zs)
        ax1 = (max(floor(Int,chunk[1][1]  +invmap_xextrema[iz][1]), 1) :
               min( ceil(Int,chunk[1][end]+invmap_xextrema[iz][2]), size(ccurr,1)))
        ax2 = (max(floor(Int,chunk[2][1]  +invmap_yextrema[iz][1]), 1) :
               min( ceil(Int,chunk[2][end]+invmap_yextrema[iz][2]), size(ccurr,2)))
        oacurr = OffsetArray(ccurr[ax1, ax2, z], ax1, ax2)
        _out = warp(oacurr, tform[iz], (chunk[1],chunk[2]);  method=Lanczos4OpenCVFaithful())
        out[:,:,iz] = round.(clamp.(OffsetArrays.no_offset_view(_out),
                                          typemin(eltype(warped)), typemax(eltype(warped))))
    end
    warped[chunk[1],chunk[2],zs] = out
end

function warp_slab(warped, ccurr, zs)
    invmap_xextrema = [extrema(filter(!isnan, invmap[:,:,z,1])) for z in zs]
    invmap_yextrema = [extrema(filter(!isnan, invmap[:,:,z,2])) for z in zs]
    tform = map(zs) do z
        itpx = extrapolate(scale(interpolate(invmap[:,:,z,1], BSpline(Linear())), sx, sy), NaN32)
        itpy = extrapolate(scale(interpolate(invmap[:,:,z,2], BSpline(Linear())), sx, sy), NaN32)
        function tform(p::SVector{2})
            x::Float32 = itpx(p[1], p[2])
            y::Float32 = itpy(p[1], p[2])
            SVector{2}(p[1]+x, p[2]+y)
        end
    end

    iz = findfirst(x->all(in.(zs,Ref(x[3]))), DiskArrays.eachchunk(warped)[1,1,:])
    chunks = DiskArrays.eachchunk(warped)[:,:,iz]
    p = Progress(prod(size(chunks)))
    Threads.@threads :greedy for i in 1:cartesian2morton([size(chunks)...])
        ix, iy = morton2cartesian(i)
        (ix > size(chunks,1) || iy > size(chunks,2)) && continue
        chunk = chunks[ix,iy]
        warp_chunk(chunk, warped, invmap_xextrema, invmap_yextrema, tform, ccurr, zs)
        next!(p)
    end
    finish!(p)
end