plots.jl

theme = Dict(:LarvaPlot => Dict(
                                :path_color => :lightgrey,
                                :path_linewidth => 1,
                                :shape_linewidth => 2,
                                :shape_alpha => 0.1f0,
                                :fallback_color => :lightgrey,
                               ))

@recipe(LarvaPlot) do scene
    Theme(theme[:LarvaPlot]...)
end

#

const CompatLarvaID = Int32
const ActiveLarva = Union{Nothing, CompatLarvaID}

function withalpha(outline_color, alpha_value)
    color = to_color(outline_color)
    return RGBA(color.r, color.g, color.b, alpha_value)
end

withalpha(outline_color::Observable, alpha_value::Observable) = lift(withalpha, outline_color, alpha_value)

const MinimalLarva = Tuple{PathOrOutline, PathOrOutline, Color, Bool}

const wgl = startswith(string(Makie.current_backend[]), "WGL")

function Makie.plot!(plot::LarvaPlot{<:MinimalLarva})
    path = plot[1]
    shape_outline = plot[2]
    shape_color = plot[3]
    visibility = plot[4]
    lines!(plot, path;
           color=plot.path_color,
           linewidth=plot.path_linewidth,
           visible=visibility)
    lines!(plot, shape_outline;
           color=shape_color,
           linewidth=plot.shape_linewidth,
           visible=visibility)
    translate!(plot.plots[end], Makie.Vec3f(0, 0, 1))
    if !wgl
        poly!(plot, shape_outline;
              color=withalpha(shape_color, plot.shape_alpha),
              transparency=true,
              visible=visibility)
        translate!(plot.plots[end], Makie.Vec3f(0, 0, 0.5))
    end
end

"""
Larva controller.
"""
struct StatefulLarva
    # model side
    model::LarvaModel
    tag_lut::AbstractObservable{<:TagLUT}
    timestep::AbstractObservable{Int}
    visible::AbstractObservable{Bool}
    # view side
    shape_outline::AbstractObservable{PathOrOutline}
    shape_color::AbstractObservable{<:Color}
end

function StatefulLarva(larva::LarvaModel,
        tag_lut::AbstractObservable{<:TagLUT},
        timestep::AbstractObservable{Int},
        parent_visible::AbstractObservable{Bool};
        fallback_color::OptionalColor=nothing,
    )
    timeseries = larva.fullstates
    alignedsteps = larva.alignedsteps
    first_timestep = alignedsteps[1]
    last_timestep = alignedsteps[end]

    fallback_color = @something fallback_color theme[:LarvaPlot][:fallback_color]

    # initial state
    _, state = timeseries[1]
    _, color = gettag(tag_lut, larva, larva.alignedsteps[1], fallback_color)

    # key observables
    shape_outline = Observable(outline_or_spine(Makie.Point2f, state))
    shape_color = Observable(html_color(color))
    visibility = Observable(false)

    onany(timestep, tag_lut, larva.usertags) do timestep, tag_lut, _
        parent_visible[] || return
        if first_timestep<=timestep && timestep<=last_timestep
            step = timestep - first_timestep + 1
            if step ∈ larva.missingsteps
                visibility[] = false
                return
            end
            step -= count(larva.missingsteps .< step)
            _, state = timeseries[step]
            shape_outline.val = outline_or_spine(Makie.Point2f, state)
            #
            _, color = gettag(tag_lut, larva, timestep, fallback_color)
            shape_color.val = html_color(color)
            #
            visibility[] = true # notify all subplots
        elseif visibility[]
            visibility[] = false
        end
    end

    StatefulLarva(larva, tag_lut, timestep, visibility, shape_outline, shape_color)
end

function Makie.plot!(plot::LarvaPlot{Tuple{StatefulLarva}})
    larva = plot[1][]
    path = larva.model.path
    shape_outline = larva.shape_outline
    shape_color = larva.shape_color
    visibility = larva.visible
    larvaplot!(plot, path, shape_outline, shape_color, visibility)
end

struct SingleLarvaView
    controller
    model::Observable
    usertags::Observable
    visible::Observable{Bool}
    path::Observable{PathOrOutline}
    pathtree::Observable{<:NNTree}
    outline::Observable{PathOrOutline}
    outline_color::Observable{String}
    #
    segment::Union{Nothing, Observable{PathOrOutline}}
    segment_color::Observable{String}
    segment_visible::Observable{Bool}
end

function SingleLarvaView(larvae::Vector{LarvaModel}, controller; editabletags::Bool=true)
    fallback_color = theme[:LarvaPlot][:fallback_color]

    # initial values
    larva = larvae[1]
    _, state = larva.fullstates[1]
    _, color = gettag(gettags(controller), larva, larva.alignedsteps[1], fallback_color)

    # observables
    model = Observable(larva)
    usertags = Observable(larva.usertags)
    visible = Observable(false)
    path = Observable(larva.path)
    pathtree = map(KDTree, path)
    shape_outline = Observable(outline_or_spine(Makie.Point2f, state))
    shape_color = Observable(html_color(color))

    # TODO: move the callbacks to `plot!`

    # callbacks
    on(getactivelarva(controller)) do id
        if isnothing(id)
            visible[] = false
        else
            for larva in larvae
                if id == larva.id
                    model[] = larva
                    break
                end
            end
        end
    end

    tags = gettags(controller)
    timestep = gettimestep(controller)
    player = getplayer(controller)

    f = on(usertags[]) do _
        _, color = gettag(tags[], model[], timestep[], fallback_color)
        shape_color[] = html_color(color)
    end

    on(model) do larva
        path[] = larva.path
        usertags[] = larva.usertags
        #
        Observables.off(f)
        f = on(larva.usertags) do _
            _, color = gettag(tags[], larva, timestep[], fallback_color)
            shape_color[] = html_color(color)
        end
    end

    onany(model,
          timestep,
          tags,
          usertags,
         ) do larva, timestep, tag_lut, _
        first_timestep = larva.alignedsteps[1]
        last_timestep = larva.alignedsteps[end]
        step = timestep - first_timestep + 1
        if (first_timestep <= timestep <= last_timestep &&
            getactivelarva(controller)[] == model[].id &&
            step ∉ larva.missingsteps
           )
            step -= count(larva.missingsteps .< step)
            _, state = larva.fullstates[step]
            shape_outline.val = outline_or_spine(Makie.Point2f, state)
            #
            _, color = gettag(tag_lut, larva, timestep, fallback_color)
            shape_color.val = html_color(color)
            #
            visible[] = true # notify all subplots
            timestep == last_timestep && pause!(player)
        elseif visible[]
            visible[] = false
        end
    end

    segment = editabletags ? Observable(PathOrOutline()) : nothing
    segment_color = Observable(shape_color[])
    segment_visible = Observable(false)

    SingleLarvaView(controller, model, usertags, visible,
                    path, pathtree, shape_outline, shape_color,
                    segment, segment_color, segment_visible)
end

iseditable(larvaview::SingleLarvaView) = !isnothing(larvaview.segment)

norm2(p, q) = @. norm2(p - q)
norm2(p) = p[1]^2 + p[2]^2

function pick_timestep(scene, larva)
    pos = mouseposition(scene)
    best, dist = nn(larva.pathtree[], pos)
    timestep = larva.model[].alignedsteps[best]
    return timestep, best, dist
end

function setinitialstep!(initialstep, larva)
    step = initialstep[]
    if step == 0
        timestep = gettimestep(larva.controller)[]
        firststep = larva.model[].alignedsteps[1]
        if timestep < firststep
            @warn "Current time step is earlier than current track"
            timestep = firststep
        end
        # faster but possibly wrong:
        step = timestep - firststep + 1 #findfirst(larva.model[].alignedsteps == timestep)
        initialstep[] = step
    end
    return step
end

function pick_timesegment(scene, larva, initialstep)
    timestep, best, dist = pick_timestep(scene, larva)
    path = larva.path[]
    if initialstep isa Ref
        initialstep = initialstep[]
    end
    segment = initialstep <= best ? path[initialstep:best] : path[best:initialstep]
    return timestep, segment, dist
end

function assign_tag_to_segment!(larvaview, firststep)
    iseditable(larvaview) || return
    segment_color = larvaview.segment_color[]
    controller = larvaview.controller
    larva = larvaview.model[]
    firststep′ = larva.alignedsteps[firststep]
    laststep′ = gettimestep(controller)[]
    laststep = laststep′ - larva.alignedsteps[1] + 1
    firststep″, laststep″ = min(firststep, laststep), max(firststep, laststep)
    # history
    previous_usertags = freezetags(larva.usertags[],
                                   larva.alignedsteps[firststep″],
                                   larva.alignedsteps[laststep″])
    transaction = nothing
    #
    lut = gettags(controller)
    usertags = getusertags!(larva, firststep′; lut=lut)
    tag = gettag(usertags, lut)
    if isnothing(tag)
        @info "Removing all tags on segment" t0=larva.fullstates[firststep][1] t1=larva.fullstates[laststep][1]
        # history
        transaction = RemoveAllTags(larva.id,
                                    (firststep′, laststep′),
                                    previous_usertags)
        #
        for step in firststep″:laststep″
            usertags = getusertags!(larva, larva.alignedsteps[step]; lut=lut)
            deletetags!(usertags)
        end
    else
        @assert segment_color == tag.color[]
        @info "Assigning tag to segment" tag=tag.name[] t0=larva.fullstates[firststep][1] t1=larva.fullstates[laststep][1]
        # history
        transaction = OverrideTags(larva.id,
                                   (firststep′, laststep′),
                                   tag.name[],
                                   previous_usertags)
        #
        for step in firststep″:laststep″
            usertags = getusertags!(larva, larva.alignedsteps[step]; lut=lut)
            settag!(usertags, tag; single=true)
        end
    end
    push!(transactions(controller), transaction)
    notify(larva.usertags)
    notify(larvaview.usertags)
    notify(gettimestep(controller)) # full refresh
    flag_active_larva_as_edited(controller)
end

function setmouseevents!(scene, larva::SingleLarvaView;
        blocking=true, priority=10, kwargs...)
    # WGLMakie does not support addmouseevents! with plots as additional input args
    mouseevents = Makie.addmouseevents!(scene; priority=priority)

    Makie.onmouseleftclick(mouseevents) do _
        timestep, _, _ = pick_timestep(scene, larva)
        settimestep!(larva.controller, timestep)
        return Consume(false)
    end

    reference_larva_size = getmedianlarvasize(larva.controller)

    # dead code
    modifier_active() = ispressed(scene, Keyboard.left_shift)

    function larva_hovered()
        # checkout the current model at once to compensate for lacking atomicity
        currentstep = gettimestep(larva.controller)[]
        path = larva.path[]
        model = larva.model[]
        #
        firststep = model.alignedsteps[1]
        step = currentstep - firststep + 1
        # no check for missing steps here, as we do not need single-step precision
        step -= count(model.missingsteps .< step)
        # note: we are in single-larvae view; `path` has not been processed by `downsample`
        if length(path) < step
            @error "BoundsError: attempt to access $(length(path))-point path at index $(step)" id=model.id firststep=firststep laststep=model.alignedsteps[end] currentstep=currentstep
            return false
        end
        centroid = path[step]
        pointer = mouseposition(scene)
        return norm2(centroid - pointer) <= reference_larva_size^2
    end

    dragging = Ref(false)
    initialstep = Ref(0)
    dragstart = Ref(0.0)

    function start()
        dragstart[] = time()
        dragging[] = true
        setinitialstep!(initialstep, larva)
        if iseditable(larva)
            larva.segment_visible[] = false
            larva.segment[] = PathOrOutline()
            larva.segment_color[] = larva.outline_color[]
        end
    end

    function reset()
        dragging[] = false
        initialstep[] = 0
        if iseditable(larva)
            larva.segment_visible[] = false
            larva.segment[] = PathOrOutline()
        end
    end

    stop = reset

    Makie.onmouseleftdragstart(mouseevents) do _
        if larva_hovered()
            @debug "Start to drag"
            start()
            return Consume(blocking)
        end
        return Consume(false)
    end

    dodrag = newobservable(Cooldown(0.1), true)

    Makie.onmouseleftdrag(mouseevents) do _
        if dragging[]
            dodrag[] = true
            return Consume(blocking)
        end
        return Consume(false)
    end

    Makie.onmouseleftdragstop(mouseevents) do _
        if dragging[]
            dodrag[] = false
        end
        return Consume(false)
    end

    on(dodrag) do b
        if b
            @debug "Dragging"
            if iseditable(larva)
                timestep, segment, dist = pick_timesegment(scene, larva, initialstep)
                settimestep!(larva.controller, timestep)
                if dist > 2reference_larva_size
                    assignmentfailed!(larva.controller, "mouse pointer away")
                    stop()
                else
                    if !(isempty(segment) || larva.segment_visible[])
                        larva.segment_visible[] = true
                    end
                    larva.segment[] = segment
                end
            else
                timestep, _, _ = pick_timestep(scene, larva)
                settimestep!(larva.controller, timestep)
            end
        else
            @debug "Stop dragging"
            if iseditable(larva)
                if time() - dragstart[] > 1
                    assign_tag_to_segment!(larva, initialstep[])
                else
                    assignmentfailed!(larva.controller, "drag duration too short")
                end
            end
            stop()
        end
    end

    on(taggingevents(larva.controller).assignmentfailed) do reason
        @info "Tag assignment cancelled ($reason)"
    end

    on(events(scene).mousebutton) do event
        if event.action == Mouse.press && event.button != Mouse.left
            assignmentfailed!(larva.controller, "$(event.button) mouse button pressed")
            stop()
        end
    end

end

function Makie.plot!(plot::LarvaPlot{Tuple{SingleLarvaView}})
    larva = plot[1][]
    scene = plot.parent
    setmouseevents!(scene, larva)
    setkeyboardevents!(scene, larva)
    larvaplot!(plot, larva.path, larva.outline, larva.outline_color, larva.visible)
    if iseditable(larva)
        lines!(plot, larva.segment; color=larva.segment_color, linewidth=6, visible=larva.segment_visible)
        translate!(plot.plots[end], Makie.Vec3f(0, 0, 0.1))
    end
end

setkeyboardevents!(scene::Scene, larva::SingleLarvaView) = setkeyboardevents!(scene, larva.controller)
function setkeyboardevents!(scene::Scene, controller)
    player = getplayer(controller)
    on(events(scene).keyboardbutton) do event
        if event.action == Keyboard.press
            if event.key in (Keyboard.left, Keyboard.down)
                stepbackward!(player)
            elseif event.key in (Keyboard.right, Keyboard.up)
                stepforward!(player)
            end
        end
    end
end

##

struct DecoratedLarva
    larva::StatefulLarva
    activearea
    label
    decorated::AbstractObservable{Bool}
end

function DecoratedLarva(larva::StatefulLarva)
    decorated = Observable(false)
    on(larva.visible) do visible
        # on larva disappearance
        if !visible && decorated[]
            decorated[] = false
        end
    end
    l, r = extrema([p[1] for p in larva.model.path])
    b, t = extrema([p[2] for p in larva.model.path])
    margin = theme[:DecoratedLarva][:decoration_margin]
    if 0 < margin
        l -= margin
        b -= margin
        r += margin
        t += margin
    end
    bl, tr = Meshes.Point2f(l, b), Meshes.Point2f(r, t)
    br, tl = Meshes.Point2f(r, b), Meshes.Point2f(l, t)
    activearea = Meshes.Quadrangle(bl, br, tr, tl)
    label = string(larva.model.id)
    DecoratedLarva(larva, activearea, label, decorated)
end

function DecoratedLarva(larva::LarvaModel, args...; kwargs...)
    DecoratedLarva(StatefulLarva(larva, args...; kwargs...))
end

theme[:DecoratedLarva] = Dict(:hover_color => :red,
                              :hover_linewidth => 4,
                              :decoration_margin => 2.0f0,
                             )

function Makie.plot!(plot::LarvaPlot{Tuple{DecoratedLarva}})
    decoratedlarva = plot[1][]
    decorated = decoratedlarva.decorated
    larva = decoratedlarva.larva
    # on(decorated) do b
    #     if b
    #         @info "larva #$(larva.model.id) decorated"
    #     end
    # end

    # plot
    larvaplot!(plot, larva)
    #decoratedlarva.primitive_child = plot.plots[1]

    # decoration
    p = Makie.Point2f ∘ coordinates
    outline = p.(Vector(vertices(decoratedlarva.activearea)))
    # TODO: design a more systematic way to prevent path closing of spine, i.e. when outline
    # data is not available and `outline` is actually a spine
    close!(path) = if 7 < length(path)
        push!(path, path[1])
    else
        path
    end
    lines!(plot, close!(outline);
           color=theme[:DecoratedLarva][:hover_color],
           linewidth=theme[:DecoratedLarva][:hover_linewidth],
           visible=decorated,
          )
    text!(plot, mean(outline); text=decoratedlarva.label, visible=decorated)
end

struct DecoratedLarvae
    larvae::Vector{DecoratedLarva}
    centers::Array{Float32, 2}
    norm2::Array{Float32, 2}
    hovered_larva::AbstractObservable{Int}
    hovering_active::AbstractObservable{Bool}
end

function DecoratedLarvae(larvae::Vector{LarvaModel}, args...; kwargs...)
    larvae′ = DecoratedLarva[]
    for larva in larvae
        try
            push!(larvae′, DecoratedLarva(larva, args...; kwargs...))
        catch e
            @warn "Skipping track on failure" id=larva.id cause=e
        end
    end
    DecoratedLarvae(larvae′)
end

function DecoratedLarvae(larvae::Vector{DecoratedLarva})
    # not-a-larva is encoded as index 0
    current_larva = Observable(0)
    hovered_larva = Observable(0)
    hovering_active = Observable(true)
    on(hovering_active) do active
        if !(active || isnothing(hovered_larva[]))
            hovered_larva[] = 0
        end
    end
    # note: if an active larva turns invisible, decorations will disappear,
    #       but the above observables won't be cleared
    on(hovered_larva) do j
        @assert hovering_active[]
        i = current_larva[]
        if i == j
            @warn "moving too fast? - please file an issue at https://gitlab.pasteur.fr/nyx/larvatagger.jl/-/issues"
        end
        if 0 < i
            i_decorated = larvae[i].decorated
            if i_decorated[]
                i_decorated[] = false
            else
                @warn begin
                    i_id = larvae[i].larva.model.id
                    "Larva #$(i_id)'s decorations already removed"
                end
            end
        end
        current_larva.val = j
        if 0 < j
            if larvae[j].larva.visible[]
                larvae[j].decorated[] = true
            else
                @error begin
                    j_id = larvae[j].larva.model.id
                    "cannot decorate invisible larva #$(j_id) - please file an issue at https://gitlab.pasteur.fr/nyx/larvatagger.jl/-/issues"
                end
                hovered_larva.val = 0
            end
        end
    end
    center = Meshes.coordinates ∘ Meshes.centroid
    centers = cat((center(larva.activearea) for larva in larvae)...; dims=2)
    norm2 = sum(centers .* centers, dims=1)
    DecoratedLarvae(larvae, centers, norm2, hovered_larva, hovering_active)
end

function find(larvae::DecoratedLarvae, position)
    pos = Vector(position)'
    p2 = pos * pos'
    visible = [i for (i, larva) in enumerate(larvae.larvae) if larva.larva.visible[]]
    isempty(visible) && return 0
    dist2 = vec(larvae.norm2[:,visible] .+ p2 .- 2 * pos * larvae.centers[:,visible])
    bestlarva = visible[argmin(dist2)]
    if Meshes.Point2f(pos...) ∉ larvae.larvae[bestlarva].activearea
        # not-a-larva is encoded as index 0
        bestlarva = 0
    end
    return bestlarva
end

function Makie.plot!(plot::LarvaPlot{Tuple{DecoratedLarvae}})
    content = plot[1][]
    for larva in content.larvae
        larvaplot!(plot, larva)
    end
    scene = plot.parent
    on(events(scene).mouseposition) do _
        if content.hovering_active[]
            #pos = Makie.to_world(scene, Makie.screen_relative(scene, mp))
            pos = mouseposition(scene)
            # both `find` and `hovered_larva` return larva indices (not IDs)
            larva = find(content, pos)
            if larva != content.hovered_larva[]
                content.hovered_larva[] = larva
            end
        end
        return Consume(false)
    end
end

Meshes.boundingbox(larva::StatefulLarva) = Meshes.boundingbox(larva.model)
Meshes.boundingbox(larva::DecoratedLarva) = Meshes.boundingbox(larva.larva)
Meshes.boundingbox(larvae::DecoratedLarvae) = Meshes.boundingbox(map(Meshes.boundingbox, larvae.larvae))

function setmouseevents!(scene, plot::DecoratedLarvae, ctrl; consume=false)
    on(events(scene).mousebutton) do mb
        consumed = consume
        if mb.button == Mouse.left && mb.action == Mouse.release
            larva_ix = plot.hovered_larva[]
            if 0 < larva_ix
                larva_id = plot.larvae[larva_ix].larva.model.id
                activatelarva!(ctrl, larva_id)
                consumed = true
            end
        end
        return Consume(consumed)
    end
end

const FIGSIZE = (1200, 800)

function autosize!(axis, size)
    target_ratio = Float32(size[1] / size[2])
    on(axis.targetlimits) do limits
        width, height = limits.widths
        ratio = width / height
        if 1e-3 <= abs(ratio - target_ratio)
            # enlarge the shortest dimension at constant aspect ratio to approach the target
            # width/height ratio
            x0, y0 = limits.origin
            if ratio < target_ratio
                width′= height * target_ratio
                origin = [x0 - (width′- width) / 2, y0]
                widths = [width′, height]
            else#if target_ratio < ratio
                height′= width / target_ratio
                origin = [x0, y0 - (height′- height) / 2]
                widths = [width, height′]
            end
            axis.targetlimits[] = typeof(limits)(origin, widths)
        end
    end
end