mesh.jl

export init_grid, init_hgrid, save, refine!, cells, vertices

using WriteVTK
using StaticArrays: SVector

struct HMesh
    vertices::Vector{SVector{2, Float64}}
    cells::Vector{NTuple{3, Int}}
end

Base.show(io::IO, ::MIME"text/plain", x::HMesh) =
    print("$(nameof(typeof(x))), $(ncells(x)) cells")

ndims_domain(::HMesh) = 2
ndims_space(::HMesh) = 2
nvertices_cell(::HMesh) = 3
nvertices(x::HMesh) = length(x.vertices)
ncells(x::HMesh) = length(x.cells)

cells(mesh::HMesh) = axes(mesh.cells, 1)
vertices(mesh::HMesh, cell) = mesh.cells[cell]

function init_hgrid(m::Int, n::Int = m, v0 = (0., 0.), v1 = (1., 1.))
    r1 = LinRange(v0[1], v1[1], m + 1)
    r2 = LinRange(v0[2], v1[2], n + 1)
    coords = collect(Iterators.product(r1, r2))

    vertices = SVector{2, Float64}.(vec(coords))
    cells = Vector{NTuple{3, Int}}()

    vidx = LinearIndices((m + 1, n + 1))
    e1 = CartesianIndex(1, 0)
    e2 = CartesianIndex(0, 1)
    k = 0
    for I in CartesianIndices((m, n))
	push!(cells, (vidx[I], vidx[I + e1], vidx[I + e1 + e2]))
	push!(cells, (vidx[I], vidx[I + e1 + e2], vidx[I + e2]))
    end

    return HMesh(vertices, cells)
end

init_hgrid(img::Array{<:Any, 2}; type=:vertex) =
    type == :vertex ?
	init_grid(size(img, 1) - 1, size(img, 2) - 1, (1.0, 1.0), size(img)) :
	init_grid(size(img, 1), size(img, 2), (0.5, 0.5), size(img) .- (0.5, 0.5))


# 2d, simplex grid
struct Mesh
    vertices::Array{Float64, 2}
    cells::Array{Int, 2}
end

Base.show(io::IO, ::MIME"text/plain", x::Mesh) =
    print("$(nameof(typeof(x))), $(ncells(x)) cells")

ndims_domain(::Mesh) = 2
ndims_space(::Mesh) = 2
nvertices_cell(::Mesh) = 3
nvertices(x::Mesh) = size(x.vertices, 2)
ncells(x::Mesh) = size(x.cells, 2)

cells(mesh::Mesh) = axes(mesh.cells, 2)
vertices(mesh::Mesh, cell) = ntuple(i -> mesh.cells[i, cell], nvertices_cell(mesh))

#function Base.getproperty(obj::Mesh, sym::Symbol)
#    if sym === :vertices
#	return reinterpret(SVector{ndims_space(obj), Float64}, vec(obj.vertices))
#    else
#	return getfield(obj, sym)
#    end
#end

function init_grid_old(m::Int, n::Int = m, v0 = (0., 0.), v1 = (1., 1.))
    r1 = LinRange(v0[1], v1[1], m + 1)
    r2 = LinRange(v0[2], v1[2], n + 1)
    coords = collect(Iterators.product(r1, r2))

    vertices = [x[i] for i in 1:2, x in vec(coords)]
    cells = Array{Int, 2}(undef, 3, 2 * m * n)

    vidx = LinearIndices((m + 1, n + 1))
    e1 = CartesianIndex(1, 0)
    e2 = CartesianIndex(0, 1)
    k = 0
    for I in CartesianIndices((m, n))
	cells[:, k += 1] .= (vidx[I], vidx[I + e1], vidx[I + e1 + e2])
	cells[:, k += 1] .= (vidx[I], vidx[I + e1 + e2], vidx[I + e2])
    end

    return Mesh(vertices, cells)
end

function init_grid(m::Int, n::Int = m, v0 = (0., 0.), v1 = (1., 1.))
    r1 = LinRange(v0[1], v1[1], m + 1)
    r2 = LinRange(v0[2], v1[2], n + 1)
    coords = collect(Iterators.product(r1, r2))

    vertices = [x[i] for i in 1:2, x in vec(coords)]
    cells = Array{Int, 2}(undef, 3, 2 * m * n)

    vidx = zeros(Int, m + 1, n + 1)
    k = 0
    for j = 1 : n + 1
	for i = 1 : m + 1
	    if iseven(i + j)
		vidx[i, j] = k += 1
	    end
	end
    end
    for j = 1 : n + 1
	for i = 1 : m + 1
	    if isodd(i + j)
		vidx[i, j] = k += 1
	    end
	end
    end
    vidxinv = reshape(invperm(vec(vidx)), m + 1, n + 1)
    vertices = reshape(vertices[:, vidxinv], 2, :)

    e1 = CartesianIndex(1, 0)
    e2 = CartesianIndex(0, 1)
    k = 0
    for I in CartesianIndices((m, n))
	if iseven(I[1] + I[2])
	    cells[:, k += 1] .= (vidx[I], vidx[I + e1], vidx[I + e1 + e2])
	    cells[:, k += 1] .= (vidx[I], vidx[I + e1 + e2], vidx[I + e2])
	else
	    cells[:, k += 1] .= (vidx[I], vidx[I + e1], vidx[I + e2])
	    cells[:, k += 1] .= (vidx[I + e1], vidx[I + e1 + e2], vidx[I + e2])
	end
    end

    return Mesh(vertices, cells)
end

init_grid(img::Array{<:Any, 2}; type=:vertex) =
    type == :vertex ?
	init_grid(size(img, 1) - 1, size(img, 2) - 1, (1.0, 1.0), size(img)) :
	init_grid(size(img, 1), size(img, 2), (0.5, 0.5), size(img) .- (0.5, 0.5))

function refine!(mesh::HMesh, marked_cells::Set)
    refined_cells = Set{Int}()
    # assemble edge -> cells map
    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)
    end

    function refine_cell(c1)
	c2 = -1
	# c1 -> c11 + c12
	# c2 -> c21 + c22
	c1_vs = sort(SVector(vertices(mesh, c1)))
	c2_arr = setdiff(edgemap[(c1_vs[1], c1_vs[2])], c1)

	if !isempty(c2_arr)
	    @assert(length(c2_arr) == 1)
	    c2 = c2_arr[begin]
	    c2_vs = sort(SVector(vertices(mesh, c2)))

	    if c1_vs[end] != c2_vs[end]
		# cannot refine `cellop` compatibly, recurse
		refine_cell(c2)
		# refetch c2 because topology has changed
		c2_arr = setdiff(edgemap[(c1_vs[1], c1_vs[2])], c1)
		@assert(!isempty(c2_arr))
		c2 = c2_arr[begin]
		c2_vs = sort(SVector(vertices(mesh, c2)))
	    end
	    # c2 exists and compatibility is guaranteed
	    @assert(c1_vs[1:2] == c2_vs[1:2])
	end

	# create bisection vertex
	xbisect = (mesh.vertices[c1_vs[1]] + mesh.vertices[c1_vs[2]]) / 2
	push!(mesh.vertices, xbisect)
	vbisect = lastindex(mesh.vertices)

	push!(mesh.cells, Tuple(sort(replace(c1_vs, c1_vs[1] => vbisect))))
	c3 = lastindex(mesh.cells)
	replace!(edgemap[NTuple{2}(setdiff(c1_vs, c1_vs[1]))], c1 => c3)

	push!(mesh.cells, Tuple(sort(replace(c1_vs, c1_vs[2] => vbisect))))
	c4 = lastindex(mesh.cells)
	replace!(edgemap[NTuple{2}(setdiff(c1_vs, c1_vs[2]))], c1 => c4)

        # flip is correct
	edgemap[(c1_vs[1], vbisect)] = [c4]
	edgemap[(c1_vs[2], vbisect)] = [c3]
	edgemap[(c1_vs[3], vbisect)] = [c3, c4]

	delete!(marked_cells, c1)
	push!(refined_cells, c1)

	if c2 > 0
	    push!(mesh.cells, Tuple(sort(replace(c2_vs, c1_vs[1] => vbisect)))) # c1_vs is correct
	    c5 = lastindex(mesh.cells)
	    replace!(edgemap[NTuple{2}(setdiff(c2_vs, c1_vs[1]))], c1 => c5)

	    push!(mesh.cells, Tuple(sort(replace(c2_vs, c1_vs[2] => vbisect)))) # c1_vs is correct
	    c6 = lastindex(mesh.cells)
	    replace!(edgemap[NTuple{2}(setdiff(c2_vs, c1_vs[2]))], c1 => c6)

	    # flip is correct
	    push!(edgemap[(c1_vs[1], vbisect)], c6)
	    push!(edgemap[(c1_vs[2], vbisect)], c5)
	    edgemap[(c1_vs[3], vbisect)] = [c5, c6]

	    delete!(marked_cells, c2)
	    push!(refined_cells, c1)
	end
    end

    while !isempty(marked_cells)
	refine_cell(first(marked_cells))
    end

    deleteat!(mesh.cells, sort(collect(refined_cells)))
    return mesh
end

#function cell_contains(mesh, cell, v)
#    geo = mesh.vertices[:, mesh.cells[:, cell]]
#    J = jacobian(x -> geo * [1 - x[1] - x[2], x[1], x[2]], [0., 0.])
#    λ = J \ (v - geo[:, 1])
#    return all(λ .>= 0) && sum(λ) <= 1
#end

function vtk_mesh(filename, mesh::Mesh)
    cells = [MeshCell(VTKCellTypes.VTK_TRIANGLE, 3*(i-1)+1:3*(i-1)+3)
	for i in axes(mesh.cells, 2)]
    vertices = reshape(mesh.vertices[:, mesh.cells], 2, :)

    return vtk_grid(filename, vertices, cells)
end

"convenience function for saving to vtk"
function save(filename::String, mesh::Mesh, fs...)
    vtk = vtk_mesh(filename, mesh)
    for f in fs
	f.space.mesh == mesh ||
	    throw(ArgumentError("meshes do not match"))
        vtk_append!(vtk, f)
    end
    vtk_save(vtk)
end