image.jl

export evaluate_bilinear, halve, warp_backwards

# ImageFunction

# TODO: inherit from some abstract function type
struct ImageFunction{Img}
    mesh::Mesh
    img::Img
    cell::Base.RefValue{Int}
end

ImageFunction(mesh, img) =
    ImageFunction(mesh, img, Ref(1))

bind!(f::ImageFunction, cell) = f.cell[] = cell
# transform coordinates to image/matrix indexing space
img_coord(img, x) = (size(img, 1) - x[2] + 1, x[1])
evaluate(f::ImageFunction, xloc) =
    evaluate_bilinear(f.img,
        img_coord(f.img, elmap(f.mesh, f.cell[])(xloc)))

# FIXME: unused?
from_img(img) = permutedims(reverse(arr; dims = 1))

eval_neumann(img, x) = img[clamp.(Tuple(x), axes(img))...]

# uses native array indexing, i.e. 1:n
function evaluate_bilinear(img, x)
    x0 = floor.(Int, x)
    x1 = x0 .+ 1

    val = zero(eltype(img))
    for idx in CartesianIndices(ntuple(_->0:1, ndims(img)))
	cornerbool = Bool.(Tuple(idx))
	λ = ifelse.(cornerbool, x .- x0, x1 .- x)
	corner = ifelse.(cornerbool, x1, x0)
	val += prod(λ) * eval_neumann(img, corner)
    end
    return val
end

function warp_backwards(img, u)
    d = ndims(img)
    axes(img) == axes(u)[2:end] && 1:d == axes(u, 1) ||
	throw(ArgumentError("invalid dimensions"))

    res = similar(img)
    for I in CartesianIndices(img)
	x = Tuple(I) .+ ntuple(i -> u[i, I], d)
	res[I] = evaluate_bilinear(img, x)
    end
    return res
end

function halve(img)
    all(iseven.(size(img))) || throw(ArgumentError("uneven size"))

    res = similar(img, size(img) .÷ 2)
    fill!(res, zero(eltype(res)))

    for i in CartesianIndices(img)
        j = CartesianIndex((Tuple(i) .+ 1) .÷ 2)
        res[j] += img[i]
    end
    res ./= 2 ^ ndims(img)

    return res
end

# Sampling

# evaluate the function on a rectangular grid of coordinates with integer
# spacing bounded by rounded minimum and maximum vertex coordinates.
# i.e. pixel center coordinates of an image. output indexing as for images:
# first index is y downwards, second index is x rightwards
# TODO: get rid of this hack as soon as we use size=() for scalar functions
function sample(f::FeFunction)
    out = _sample(f)
    return f.space.size == (1,) ?
        reshape(out, size(out)[2:end]) : out
end

function _sample(f::FeFunction)
    # assumes dual grid
    mesh = f.space.mesh

    m0 = NTuple{2}(minimum(mesh.vertices, dims = 2))
    m1 = NTuple{2}(maximum(mesh.vertices, dims = 2))

    fcolon = ntuple(_ -> :, length(f.space.size))

    outsize = (f.space.size..., round.(Int, m1 .- m0 .+ 1)...)
    out = similar(f.data, outsize)
    for cell in cells(mesh)
        bind!(f, cell)
	A = SArray{Tuple{ndims_space(mesh), nvertices_cell(mesh)}}(
	    view(mesh.vertices, :, view(mesh.cells, :, cell)))

        # TODO: cleanup when minimum/maximum with dims on StaticArrays becomes
        # type-stable
	I0 = round.(Int, SVector(minimum(A[1, :]), minimum(A[2, :])) .- m0) .+ 1
	I1 = round.(Int, SVector(maximum(A[1, :]), maximum(A[2, :])) .- m0) .+ 1

	for I in CartesianIndex(I0[1], I0[2]):CartesianIndex(I1[1], I1[2])
            x = SVector(Tuple(I) .- 1 .+ m0)
            # TODO: move to global-to-local function or inside-triangle test
            xloc = (A[:, SUnitRange(2, 3)] .- A[:, 1])::SArray \
                (x .- A[:, 1])
            ε = eps()
            if all(xloc .>= -ε) && sum(xloc) <= 1 + ε
                out[fcolon..., I] .= evaluate(f, xloc)
            end
	end
    end

    # convert to image indexing
    d = length(f.space.size)
    reverse!(out, dims = d + 2) # reverse y
    out2 = permutedims(out, (ntuple(identity, d)..., d + 2, d + 1)) # flip xy

    return out2
end

# Interpolation

"""
interprets `img` as a bilinearly interpolated continuous function and applies
the default interpolation operator for the discrete function u.
"""
function interpolate!(u::FeFunction, img::AbstractArray)
    f = ImageFunction(u.space.mesh, img)
    interpolate!(u, @inline (x; f) -> f; f)
end

# TODO: refine interface: projection vs interpolation, unify different
# algorithms

project_l2_lagrange(space::FeSpace, img) =
    (u = FeFunction(space); project_l2_lagrange!(u, img))

function project_l2_lagrange!(u::FeFunction, img::AbstractArray)
    d = 2 # domain dimension
    space = u.space
    mesh = space.mesh
    f = ImageFunction(mesh, img)
    opparams = (; f)

    nrdims = prod(space.size)
    # number of element dofs (i.e. local dofs not counting range dimensions)
    nldofs = ndofs(space.element)

    a(xloc, u, du, v, dv; f) = dot(u, v)
    l(xloc, v, dv; f) = dot(f, v)

    I = Float64[]
    J = Float64[]
    V = Float64[]
    b = zeros(ndofs(space))

    # mesh cells
    for cell in cells(mesh)
	foreach(f -> bind!(f, cell), opparams)
	# cell map is assumed to be constant per cell
	delmap = jacobian(elmap(mesh, cell), SA[0., 0.])
	delmapinv = inv(delmap)
	intel = abs(det(delmap))

	p = ceil(Int, diam(mesh, cell))
	qw, qx = quadrature_composite_lagrange_midpoint(p)
	nqpts = length(qw) # number of quadrature points

	qphi = zeros(nrdims, nrdims, nldofs, nqpts)
	dqphi = zeros(nrdims, d, nrdims, nldofs, nqpts)
	for k in axes(qx, 2)
	    for r in 1:nrdims
		qphi[r, r, :, k] .= evaluate_basis(space.element, SVector{d}(view(qx, :, k)))
		dqphi[r, :, r, :, k] .= transpose(jacobian(x -> evaluate_basis(space.element, x), SVector{d}(view(qx, :, k))))
	    end
	end

	# quadrature points
	for k in axes(qx, 2)
	    xhat = SVector{d}(view(qx, :, k))
	    x = elmap(mesh, cell)(xhat)
	    opvalues = map(f -> evaluate(f, xhat), opparams)

	    # local test-function dofs
	    for jdim in 1:nrdims, ldofj in 1:nldofs
		gdofj = space.dofmap[jdim, ldofj, cell]

		phij = SArray{Tuple{space.size...}}(view(qphi, :, jdim, ldofj, k))
		dphij = SArray{Tuple{space.size..., d}}(
		    SArray{Tuple{nrdims, d}}(view(dqphi, :, :, jdim, ldofj, k)) * delmapinv)

		lv = qw[k] * l(x, phij, dphij; opvalues...) * intel
		b[gdofj] += lv

		# local trial-function dofs
		for idim in 1:nrdims, ldofi in 1:nldofs
		    gdofi = space.dofmap[idim, ldofi, cell]

		    phii = SArray{Tuple{space.size...}}(view(qphi, :, idim, ldofi, k))
		    dphii = SArray{Tuple{space.size..., d}}(
			SArray{Tuple{nrdims, d}}(view(dqphi, :, :, idim, ldofi, k)) * delmapinv)

		    av = qw[k] * a(x, phii, dphii, phij, dphij; opvalues...) * intel
		    push!(I, gdofi)
		    push!(J, gdofj)
		    push!(V, av)
		end
	    end
	end
    end

    ngdofs = ndofs(space)
    A = sparse(I, J, V, ngdofs, ngdofs)

    u.data .= A \ b
    return u
end

project_qi_lagrange(space::FeSpace, img) =
    (u = FeFunction(space); project_qi_lagrange!(u, img))

# L1-stable quasi-interpolation operator
# (https://arxiv.org/pdf/1505.06931.pdf)
#
# FIXME: currently only approximate quadrature by sampling on lagrange lattice
# based on element size. Exact evaluation is tricky to implement (lots of
# intersection handling).
function project_qi_lagrange!(u::FeFunction, img)
    u.space.element == P1() ||
        throw(ArgumentError("element unsupported"))
    d = 2 # domain dimension
    space = u.space
    mesh = space.mesh
    f = ImageFunction(mesh, img)
    # count contributions to respective dof for subsequent averaging
    gdofcount = zeros(Int, size(u.data))
    area_refel = 0.5

    u.data .= zero(eltype(u.data))
    for cell in cells(mesh)
        bind!(f, cell)

        # size: nrdims × nldofs
        # TODO: should be statically sized
        gdofs = u.space.dofmap[:, :, cell]

        p = ceil(Int, diam(mesh, cell))
        qw, qx = quadrature_composite_lagrange_midpoint(p)
	for k in axes(qx, 2)
	    xhat = SVector{d}(view(qx, :, k))
	    x = elmap(mesh, cell)(xhat)

            # size: nrdims
            value = evaluate(f, xhat)
            # size: nldofs
            phix = evaluate_basis(space.element, xhat)

            # no integration element used
            u.data[gdofs] .+= qw[k] ./ area_refel .* 3 .* value * phix'
            #u.data[gdofs] .+= qw[k] ./ area_refel * value * (12 .* phix' .- 3)
	end
        gdofcount[gdofs] .+= 1
    end
    u.data ./= gdofcount
    return u
end

struct PixelIterator{A, G}
    "cell vertices"
    A::A
    "CartesianIndices covering the cell"
    grid::G
end

Base.IteratorSize(::PixelIterator) = Base.SizeUnknown()
Base.eltype(it::PixelIterator) = eltype(it.grid)

"""
    PixelIterator(mesh, cell)
    PixelIterator(mesh, A)

Return an iterator over all integer cartesian indices (corresponding to pixel
centers) intersecting a given `cell` or geometry `A`.
"""
function PixelIterator(mesh::Mesh, cell::Int)
    A = SArray{Tuple{ndims_space(mesh), nvertices_cell(mesh)}}(
        view(mesh.vertices, :, view(mesh.cells, :, cell)))
    return PixelIterator(mesh, A)
end

function PixelIterator(mesh::Mesh, A)
    # rounding is probably correct since we don't miss any coordinates even
    # when we round inwards to the element
    I0 = round.(Int, SVector(minimum(A[1, :]), minimum(A[2, :])))
    I1 = round.(Int, SVector(maximum(A[1, :]), maximum(A[2, :])))
    return PixelIterator(A, CartesianIndex(I0...):CartesianIndex(I1...))
end

function _intersects(it::PixelIterator, I)
    x = SVector(Tuple(I))
    # TODO: move to global-to-local function or inside-triangle test
    xloc = (it.A[:, SUnitRange(2, 3)] .- it.A[:, 1])::SArray \
        (x .- it.A[:, 1])
    ε = eps()
    return all(xloc .>= -ε) && sum(xloc) <= 1 + ε
end

function Base.iterate(it::PixelIterator, state...)
    y = iterate(it.grid, state...)
    isnothing(y) && return nothing
    while !_intersects(it, first(y))
        y = iterate(it.grid, last(y))
        isnothing(y) && return nothing
    end
    return y
end

function project_l2_pixel!(u::FeFunction, img)
    d = 2 # domain dimension
    space = u.space
    mesh = space.mesh
    f = ImageFunction(mesh, img)
    opparams = (; f)

    nrdims = prod(space.size)
    # number of element dofs (i.e. local dofs not counting range dimensions)
    nldofs = ndofs(space.element)

    # first loop to count number of triangles intersecting pixel evaluation
    # points
    ncells = zeros(axes(img))
    for cell in cells(mesh)
        pixels = PixelIterator(mesh, cell)
        for I in pixels
            ncells[I] += 1
        end
    end
    all(!iszero, ncells) ||
        throw(ErrorException("not all pixels are covered by a cell"))
    pixelweights = inv.(ncells)

    # typical L2 projection of f
    a(xloc, u, du, v, dv; f) = dot(u, v)
    l(xloc, v, dv; f) = dot(f, v)

    I = Float64[]
    J = Float64[]
    V = Float64[]
    b = zeros(ndofs(space))

    # buffer for basis value/derivative at a single point
    qphi = zero(MArray{Tuple{nrdims, nrdims, nldofs}})
    dqphi = zero(MArray{Tuple{nrdims, d, nrdims, nldofs}})

    for cell in cells(mesh)
	foreach(f -> bind!(f, cell), opparams)
	# cell map is assumed to be constant per cell
	delmap = jacobian(elmap(mesh, cell), SA[0., 0.])
	delmapinv = inv(delmap)
	intel = abs(det(delmap))

	#p = ceil(Int, diam(mesh, cell))
	#qw, qx = quadrature_composite_lagrange_midpoint(p)
	#nqpts = length(qw) # number of quadrature points

        # pixels as quadrature points
        pixels = PixelIterator(mesh, cell)
        for P in pixels
            x = SVector(Tuple(P))
            # TODO: move to global-to-local function or inside-triangle test
            xhat = (pixels.A[:, SUnitRange(2, 3)] .- pixels.A[:, 1])::SArray \
                (x .- pixels.A[:, 1])
	    opvalues = map(f -> evaluate(f, xhat), opparams)

            # fill basis value/derivative buffer
	    for r in 1:nrdims
		qphi[r, r, :] .=
                    evaluate_basis(space.element, SVector{d}(xhat))
                dqphi[r, :, r, :]::MArray{Tuple{d, nldofs}, Float64, 2, d * nldofs} .=
                    transpose(jacobian(
                        x -> evaluate_basis(space.element, x),
                        SVector{d}(xhat)))
	    end

	    # local test-function dofs
	    for jdim in 1:nrdims, ldofj in 1:nldofs
		gdofj = space.dofmap[jdim, ldofj, cell]

		phij = SArray{Tuple{space.size...}}(view(qphi, :, jdim, ldofj))
		dphij = SArray{Tuple{space.size..., d}}(
		    SArray{Tuple{nrdims, d}}(view(dqphi, :, :, jdim, ldofj)) * delmapinv)

                lv = pixelweights[P] * l(x, phij, dphij; opvalues...) * intel
		b[gdofj] += lv

		# local trial-function dofs
		for idim in 1:nrdims, ldofi in 1:nldofs
		    gdofi = space.dofmap[idim, ldofi, cell]

		    phii = SArray{Tuple{space.size...}}(view(qphi, :, idim, ldofi))
		    dphii = SArray{Tuple{space.size..., d}}(
			SArray{Tuple{nrdims, d}}(view(dqphi, :, :, idim, ldofi)) * delmapinv)

		    av = pixelweights[P] * a(x, phii, dphii, phij, dphij; opvalues...) * intel
		    push!(I, gdofi)
		    push!(J, gdofj)
		    push!(V, av)
		end
	    end
	end
    end

    ngdofs = ndofs(space)
    A = sparse(I, J, V, ngdofs, ngdofs)

    u.data .= A \ b
    return u

end