image.jl

export evaluate_bilinear, halve, warp_backwards

# General Utility Methods on Images

function mean_squared_error(img, img_ref)
    axes(img) == axes(img_ref) ||
        throw(ArgumentError("images need matching axes"))
    n = length(img)
    return sum((img .- img_ref) .^ 2 ./ n)
end

function peak_signal_to_noise_ratio(img, img_ref)
    mse = mean_squared_error(img, img_ref)
    imax = 1.0 # maximum intensity
    return 10 * log10(imax / mse)
end

# 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

# ArrayFunction mesh function wrapper

# TODO: inherit from some abstract mesh function type
# TODO: make mesh type a parameter for performance
struct ArrayFunction{Img}
    mesh::Mesh
    img::Img
    cell::Base.RefValue{Int}
end

ArrayFunction(mesh, img) = ArrayFunction(mesh, img, Ref(0))

bind!(f::ArrayFunction, cell) = f.cell[] = cell
evaluate(f::ArrayFunction, xloc) = evaluate_bilinear(f.img,
    elmap(f.mesh, f.cell[])(xloc))

# 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 + ε
                # use view here to handle scalar data assigment as well
                view(out, fcolon..., I) .= evaluate(f, xloc)
            end
        end
    end

    return out
end


# Mesh Function Projection Methods

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

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

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

"""
standard l2 projection but using an adaptive quadrature rule based on cell size
"""
# TODO: this could be a general l2 projection, accepting a mesh function and
# some custom quadrature rule?
function project_l2_lagrange!(u::FeFunction, img::AbstractArray)
    d = 2 # domain dimension
    space = u.space
    mesh = space.mesh
    f = ArrayFunction(mesh, img)
    opparams = (; f)

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

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

    I = Int[]
    J = Int[]
    V = Float64[]
    n = length(img) * nrdims * nldofs * nrdims * nldofs
    sizehint!(I, n)
    sizehint!(J, n)
    sizehint!(V, n)
    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 = ArrayFunction(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
            # actually: transformation to barycentric coordinates!
            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)

            #rhox = SA[-12 -12; 12 0; 0 12] * xhat + SA[9, -3, -3]
            #u.data[gdofs] .+= qw[k] ./ area_refel * value * rhox'
        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`.
The returned indices represent coordinates in the given mesh.
"""
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
    # TODO: do something more sophisticated than walking the whole cartesian
    # bounding box
    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 = ArrayFunction(mesh, img)
    opparams = (; f)

    nrdims = prod(space.size)
    # number of element dofs (i.e. local dofs not counting codomain 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))

        # 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