function.jl

using Statistics: mean
using StaticArrays: SA, SArray, MVector, MMatrix, SUnitRange
using DataFrames, CSV
export FeSpace, Mapper, FeFunction, ImageFunction, P1, DP0, DP1
export interpolate!, sample, bind!, evaluate, integrate, nabla, save_csv

# Finite Elements
struct P1 end
struct DP0 end
struct DP1 end

ndofs(::P1) = 3
ndofs(::DP0) = 1
ndofs(::DP1) = 3

# evaluate all (1-dim) local basis functions against x
evaluate_basis(::P1, x) = SA[1 - x[1] - x[2], x[1], x[2]]
evaluate_basis(::DP0, x) = SA[1]
evaluate_basis(::DP1, x) = evaluate_basis(P1(), x)

# non-mixed
# scalar elements

struct FeSpace{M, Fe, S}
    mesh::M
    element::Fe
    dofmap::Array{Int, 3} # (rdim, eldof, cell) -> gdof
    ndofs::Int # = maximum(dofmap)
end

# TODO: size should probably be a type argument to make the constructor type
# safe
function FeSpace(mesh, el::P1, size_=(1,))
    # special case for P1: dofs correspond to vertices
    rdims = prod(size_)
    dofmap = Array{Int, 3}(undef, rdims, 3, ncells(mesh))
    for i in CartesianIndices((rdims, 3, ncells(mesh)))
	dofmap[i] = rdims * (mesh.cells[i[2], i[3]] - 1) + i[1]
    end
    return FeSpace{typeof(mesh), typeof(el), size_}(
	mesh, el, dofmap, rdims * nvertices(mesh))
end

function FeSpace(mesh, el::DP0, size_=(1,))
    # special case for DP1: dofs correspond to cells
    rdims = prod(size_)
    dofmap = Array{Int, 3}(undef, rdims, 1, ncells(mesh))
    for i in CartesianIndices((rdims, 1, ncells(mesh)))
	dofmap[i] = rdims * (i[3] - 1) + i[1]
    end
    return FeSpace{typeof(mesh), typeof(el), size_}(
	mesh, el, dofmap, rdims * ncells(mesh))
end

function FeSpace(mesh, el::DP1, size_=(1,))
    # special case for DP1: dofs correspond to vertices
    rdims = prod(size_)
    dofmap = Array{Int, 3}(undef, rdims, 3, ncells(mesh))
    for i in LinearIndices((rdims, 3, ncells(mesh)))
	dofmap[i] = i
    end
    return FeSpace{typeof(mesh), typeof(el), size_}(
	mesh, el, dofmap, rdims * 3 * ncells(mesh))
end

Base.show(io::IO, ::MIME"text/plain", x::FeSpace) =
    print("$(nameof(typeof(x))), $(nameof(typeof(x.element))) elements, size $(x.size), $(ndofs(x)) dofs")

function Base.getproperty(obj::FeSpace{<:Any, <:Any, S}, sym::Symbol) where S
    if sym === :size
	return S
    else
	return getfield(obj, sym)
    end
end

ndofs(space::FeSpace) = space.ndofs

# evaluate at local point
@inline function evaluate(space::FeSpace, ldofs, xloc)
    bv = evaluate_basis(space.element, xloc)
    ldofs_ = SArray{Tuple{prod(space.size), ndofs(space.element)}}(ldofs)
    v = ldofs_ * bv
    return SArray{Tuple{space.size...}}(v)
end



# dof ordering for vector valued functions:
#   (rsize..., eldofs)

# Array-valued function
struct FeFunction{Sp, Ld}
    space::Sp
    data::Vector{Float64} # gdof -> data
    name::String
    ldata::Ld # ldof -> data
end

function FeFunction(space; name=string(gensym("f")))
    data = Vector{Float64}(undef, ndofs(space))
    ldata = zero(MVector{prod(space.size) * ndofs(space.element)})
    return FeFunction(space, data, name, ldata)
end

Base.show(io::IO, ::MIME"text/plain", f::FeFunction) =
    print("$(nameof(typeof(f))), size $(f.space.size) with $(length(f.data)) dofs")

interpolate!(dst::FeFunction, expr::Function; params...) =
    interpolate!(dst, dst.space.element, expr; params...)


myvec(x) = vec(x)
myvec(x::Number) = x

function interpolate!(dst::FeFunction, ::P1, expr::Function; params...)
    params = NamedTuple(params)
    space = dst.space
    mesh = space.mesh
    for cell in cells(mesh)
	for f in params
	    bind!(f, cell)
	end

        F = elmap(mesh, cell)
        vloc = SA[0. 1. 0.; 0. 0. 1.]
	for eldof in 1:nvertices_cell(mesh)
            xloc = vloc[:, eldof]
            x = F(xloc)::SArray

	    opvalues = map(f -> evaluate(f, xloc), params)

	    gdofs = space.dofmap[:, eldof, cell]
	    dst.data[gdofs] .= myvec(expr(x; opvalues...))
	end
    end
end

function interpolate!(dst::FeFunction, ::DP0, expr::Function; params...)
    params = NamedTuple(params)
    space = dst.space
    mesh = space.mesh
    for cell in cells(mesh)
	for f in params
	    bind!(f, cell)
	end

        F = elmap(mesh, cell)
        lcentroid = SA[1/3, 1/3]
        centroid = F(lcentroid)

	opvalues = map(f -> evaluate(f, lcentroid), params)

	gdofs = space.dofmap[:, 1, cell]
	dst.data[gdofs] .= myvec(expr(centroid; opvalues...))
    end
end

interpolate!(dst::FeFunction, ::DP1, expr::Function; params...) =
    interpolate!(dst, P1(), expr; params...)


project!(dst::FeFunction, expr::Function; params...) =
    project!(dst, dst.space.element, expr; params...)

function project!(dst::FeFunction, ::DP0, expr; params...)
    # same as interpolate!, but scaling by cell size
    params = NamedTuple(params)
    space = dst.space
    mesh = space.mesh
    for cell in cells(mesh)
	delmap = jacobian(elmap(mesh, cell), SA[0., 0.])
	intel = abs(det(delmap))

	for f in params
	    bind!(f, cell)
	end

	vertices = mesh.vertices[:, mesh.cells[:, cell]]
	centroid = SArray{Tuple{ndims_domain(mesh)}}(mean(vertices, dims = 2))
	lcentroid = SA[1/3, 1/3]

	opvalues = map(f -> evaluate(f, lcentroid), params)

	gdofs = space.dofmap[:, 1, cell]
	dst.data[gdofs] .= myvec(expr(centroid; opvalues...)) .* intel
    end
end

function integrate(mesh::Mesh, expr; params...)
    opparams = NamedTuple(params)

    qw, qx = quadrature()

    # only for type inference
    opvalues = map(f -> evaluate(f, SA[0., 0.]), opparams)
    val = zero(expr(SA[0., 0.]; opvalues...))
    for cell in cells(mesh)
	foreach(f -> bind!(f, cell), opparams)

	delmap = jacobian(elmap(mesh, cell), SA[0., 0.])
	intel = abs(det(delmap))

	# quadrature points
	for k in axes(qx, 2)
	    xhat = qx[:, k]
	    x = elmap(mesh, cell)(xhat)
	    opvalues = map(f -> evaluate(f, xhat), opparams)

	    val += qw[k] * expr(x; opvalues...) * intel
	end
    end
    return val
end

function bind!(f::FeFunction, cell)
    gdofs_view = view(f.space.dofmap, :, :, cell)
    gdofs = SArray{Tuple{prod(f.space.size) * ndofs(f.space.element)}}(gdofs_view)
    f.ldata .= f.data[gdofs]
    return f
end

# evaluate at local point (needs bind! call before)
evaluate(f::FeFunction, x) = evaluate(f.space, f.ldata, x)

# allow any non-function to act as a constant function
bind!(c, cell) = c
evaluate(c, xloc) = c


# TODO: inherit from some abstract function type
struct Derivative{F, M}
    f::F
    delmapinv::M
end

Base.show(io::IO, ::MIME"text/plain", x::Derivative) =
    print("$(nameof(typeof(x))) of $(typeof(x.f))")

# TODO: should be called "derivative"
function nabla(f)
    d = ndims_domain(f.space.mesh)
    Derivative(f, zero(MMatrix{d, d, Float64}))
end

function bind!(df::Derivative, cell)
    bind!(df.f, cell)
    delmap = jacobian(elmap(df.f.space.mesh, cell), SA[0., 0.])
    df.delmapinv .= inv(delmap)
end

function evaluate(df::Derivative, x)
    # size: (:, d)
    jac = jacobian(x -> evaluate(df.f.space, df.f.ldata, x), x)
    jac *= df.delmapinv
    return SArray{Tuple{df.f.space.size..., length(x)}}(jac)
end


# 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
# TODO: precompute the offset from minimum mesh vertex
evaluate(f::ImageFunction, xloc) =
    interpolate_bilinear(f.img, elmap(f.mesh, f.cell[])(xloc) .+ (0.5, 0.5))


struct FacetDivergence{F}
    f::F
    cell::Base.RefValue{Int}
    edgemap::Dict{NTuple{2, Int}, Vector{Int}}
end

# TODO: incomplete!
function FacetDivergence(f::FeFunction)
    f.space.element == DP0() || throw(ArgumentError("unimplemented"))
    mesh = f.space.mesh
    edgedivergence = Dict{NTuple{2, Int}, Float64}()
    edgemap = Dict{NTuple{2, Int}, Vector{Int}}()

    for cell in cells(mesh)
	vs = sort(SVector(vertices(mesh, cell)))

	e1 = (vs[1], vs[2])
	e2 = (vs[1], vs[3])
	e3 = (vs[2], vs[3])

	edgemap[e1] = push!(get!(edgemap, e1, []), cell)
	edgemap[e2] = push!(get!(edgemap, e2, []), cell)
	edgemap[e3] = push!(get!(edgemap, e3, []), cell)

	bind!(f, cell)
	p = evaluate(f, SA[0., 0.])

	A = SArray{Tuple{ndims_space(mesh), nvertices_cell(mesh)}}(
	    view(mesh.vertices, :, view(mesh.cells, :, cell)))

	v1 = (A[:,2] - A[:,1])
	v2 = (A[:,3] - A[:,2])
	v3 = (A[:,1] - A[:,3])

	p1 = norm(v1) * (p[1] * v1[2] - p[2] * v1[1])
	p2 = norm(v2) * (p[1] * v2[2] - p[2] * v2[1])
	p3 = norm(v3) * (p[1] * v3[2] - p[2] * v3[1])

	edgedivergence[e1] = get!(edgedivergence, e1, 0.) + p1
	edgedivergence[e3] = get!(edgedivergence, e3, 0.) + p2
	edgedivergence[e2] = get!(edgedivergence, e2, 0.) + p3
    end

    return FacetDivergence(f, Ref(1), edgemap)
end

bind!(f::FacetDivergence, cell) = f.cell[] = cell


# 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
    return out
end


prolong!(new_f, old_cell, new_cells) =
    prolong!(new_f, old_cell, new_cells, new_f.space.element)

prolong!(new_f, old_cell, new_cells, ::DP1) =
    prolong!(new_f, old_cell, new_cells, P1())

function prolong!(new_f, old_cell, new_cells, ::P1)
    old_f = new_f
    old_cell_vs = collect(vertices(old_f.space.mesh, old_cell))
    new_cell1_vs = collect(vertices(new_f.space.mesh, new_cells[1]))
    new_cell2_vs = collect(vertices(new_f.space.mesh, new_cells[2]))

    # copy over data for common vertices
    common_vs = intersect(old_cell_vs, new_cell1_vs)
    old_ldofs = indexin(common_vs, old_cell_vs)
    old_gdofs = old_f.space.dofmap[:, old_ldofs, old_cell]
    new_ldofs = indexin(common_vs, new_cell1_vs)
    new_gdofs = new_f.space.dofmap[:, new_ldofs, new_cells[1]]
    new_f.data[new_gdofs] .= old_f.data[old_gdofs]

    common_vs = intersect(old_cell_vs, new_cell2_vs)
    old_ldofs = indexin(common_vs, old_cell_vs)
    old_gdofs = old_f.space.dofmap[:, old_ldofs, old_cell]
    new_ldofs = indexin(common_vs, new_cell2_vs)
    new_gdofs = new_f.space.dofmap[:, new_ldofs, new_cells[2]]
    new_f.data[new_gdofs] .= old_f.data[old_gdofs]

    # vertices of bisection edge
    avg_vs = symdiff(new_cell1_vs, new_cell2_vs)
    old_ldofs = indexin(avg_vs, old_cell_vs)
    old_gdofs = old_f.space.dofmap[:, old_ldofs, old_cell]

    avg_data = (old_f.data[old_gdofs[:, 1]] .+ old_f.data[old_gdofs[:, 2]]) ./ 2

    _, newldof = findmax(new_cell1_vs)
    new_gdofs = new_f.space.dofmap[:, newldof, new_cells[1]]
    new_f.data[new_gdofs] .= avg_data

    _, newldof = findmax(new_cell2_vs)
    new_gdofs = new_f.space.dofmap[:, newldof, new_cells[2]]
    new_f.data[new_gdofs] .= avg_data
end

function prolong!(new_f, old_cell, new_cells, ::DP0)
    old_f = new_f
    # simply copy over the data
    old_gdofs = old_f.space.dofmap[:, 1, old_cell]

    new_gdofs = new_f.space.dofmap[:, 1, new_cells[1]]
    new_f.data[new_gdofs] .= old_f.data[old_gdofs]
    new_gdofs = new_f.space.dofmap[:, 1, new_cells[2]]
    new_f.data[new_gdofs] .= old_f.data[old_gdofs]
end


vtk_append!(vtkfile, f::FeFunction, fs::FeFunction...) =
    (vtk_append!(vtkfile, f); vtk_append!(vtkfile, fs...))

vtk_append!(vtkfile, f::FeFunction) =
    vtk_append!(vtkfile, f, f.space.element)

function vtk_append!(vtkfile, f::FeFunction, ::P1)
    # separate out data per cell since different cells have distinct vertices
    # in the output vtk
    fd = vec(f.data[f.space.dofmap])
    vtk_point_data(vtkfile, fd, f.name)
end

function vtk_append!(vtkfile, f::FeFunction, ::DP0)
    vtk_cell_data(vtkfile, f.data, f.name)
end

vtk_append!(vtkfile, f::FeFunction, ::DP1) =
    vtk_append!(vtkfile, f, P1())

function save_csv(path, f::FeFunction)
    f.space.element::P1
    f.space.size == (1,) ||
	throw(ArgumentError("non-scalar functions unsupported"))
    mesh = f.space.mesh
    df = DataFrame()
    X = SA[0 1 0; 0 0 1] # corners in local coordinates
    for cell in cells(mesh)
	bind!(f, cell)
	A = SArray{Tuple{ndims_space(mesh), nvertices_cell(mesh)}}(
	    view(mesh.vertices, :, view(mesh.cells, :, cell)))

	for k in axes(A, 2)
	    push!(df, (x = A[1, k], y = A[2, k],
                f_xy = evaluate(f, X[:, k])[1]))
	end
    end
    CSV.write(path, df)
end