module QuantumSavoryMakie

using QuantumSavory
using Graphs
using NetworkLayout
using ConcurrentSim: ConcurrentSim
using LinearAlgebra: eigen
using Makie: Makie, Theme, Figure, Axis, Axis3, Aspect, Label, PolyElement, Legend, get_scene, @L_str,
    @recipe, lift, @lift, Observable,
    Point2, Point2f, Rect2f, Rect3f,
    scatter!, heatmap!, barplot!, poly!, linesegments!, lines!, hlines!, vlines!, mesh!, text!,
    xlims!, ylims!, zlims!,
    xticks!, yticks!,
    hidedecorations!, hidespines!,
    colsize!, resize_to_layout!,
    deregister_interaction!, interactions,
    DataInspector, Slider, Colorbar, axislegend, cgrad

import QuantumSavory: registernetplot, registernetplot!, registernetplot_axis, resourceplot_axis, showonplot, showmetadata
using QuantumSavory: compactstr
using QuantumSavory.ProtocolZoo: ProtocolZoo, EntanglerProt, EntanglementConsumer
using QuantumSavory: _mode_mean

using Gabs: Gabs
using QuantumClifford: QuantumClifford
using QuantumOpticsBase: QuantumOpticsBase, dm


##

@recipe(RegisterNetPlot, regnet) do scene
    Theme(
        colormap = :Spectral,
        colorrange = (-1., 1.),
        register_color = :gray90,
        slotcolor = :gray60,
        slotmarker = :rect,
        scale = 1.0,
        slotsize = 0.8,
        observables_marker = :circle,
        observables_markersize = 0.55,
        observables_linewidth = 5,
        state_markersize = 0.4,
        state_marker = :diamond,
        state_markercolor = :black,
        state_linecolor = :gray90,
        lock_marker = '⚿',
        registercoords = nothing,
        observables = nothing,
    )
end

function Makie.plot!(rn::RegisterNetPlot{<:Tuple{RegisterNet}})
    networkobs = rn[1]
    registers = networkobs[].registers

    register_rectangles = Observable(Rect2f[])     # Makie rectangles that will be plotted for each register
    register_slots_coords = Observable(Point2f[])  # Makie marker locations that will be plotted for each register slot
    state_coords = Observable(fill(Point2f(NaN, NaN), sum(nsubsystems.(registers))))  # NaN-padded fixed size
    lock_coords = Observable(fill(Point2f(NaN, NaN), sum(nsubsystems.(registers))))  # NaN-padded fixed size
    state_links = Observable(fill(Point2f(NaN, NaN), 2 * sum(nsubsystems.(registers))))  # NaN-padded fixed size; GLMakie GPU buffers cannot resize after plot creation
    observables_coords = Observable(Point2f[])     # Makie marker locations that will be plotted for each subsystem on which an observable is evaluated
    observables_links = Observable(Point2f[])      # The links between observed subsystems
    observables_vals = Observable(Float64[])       # Values of the observables (stored per marker)
    observables_linkvals = Observable(Float64[])   # Values of the observables (stored per link)
    register_backref = Observable(Any[])           # A backreference to the register object for register square
    register_slots_coords_backref = Observable(Tuple{Any,Int,Int}[]) # A backreference to the register object and reference indices for each register slot marker
    state_coords_backref = Observable(Tuple{Any,Any,Int,Int,Int}[])  # A backreference to the state object and register object and reference indices for each state marker
    observables_backref = Observable(Tuple{Any,Float64}[])           # A backreference to the observable (and its value) for each colored dot visualizing an observable
    observables_links_backref = Observable(Tuple{Any,Float64}[])     # same as above but for the links
    _extras = Dict{Symbol, Any}(
        :register_backref => register_backref,
        :register_slots_coords_backref => register_slots_coords_backref,
        :state_coords_backref => state_coords_backref,
        :observables_backref => observables_backref,
        :observables_links_backref => observables_links_backref,
    )

    # Optional arguments
    ## registercoords -- updates handled explicitly by an `onany` call below
    if !isnothing(rn[:registercoords][])
        registercoords = rn[:registercoords][]
        registercoords isa Vector{<:Point2} || throw(ArgumentError("While plotting a network layout an incorrect argument was given: `registercoords` has to be of type `Vector{<:Point2}`. You can leave it empty to autogenerate a register layout. You can generate it manually or with packages like `NetworkLayout`."))
    else
        adj_matrix = adjacency_matrix(networkobs[].graph)
        registercoords = spring(adj_matrix, iterations=400, C=2*maximum(nsubsystems.(registers)))
        rn.registercoords[] = registercoords
    end
    ## slotcolor -- updates handled implicitly (used only in a single `scatter` call)
    slotcolor_resolved = if rn[:slotcolor][] isa Vector{<:Vector} # A vector of vector of colors (i.e. a vector of colors per register)
        lift(x->reduce(vcat, x), rn[:slotcolor]) # Turn it into a vector of colors
    else
        rn[:slotcolor]
    end
    ## observables -- updates handled explicitly by an `onany` call below
    if !isnothing(rn[:observables][])
        observablesval = rn[:observables][]
        if observablesval isa Vector{<:Tuple{Any,Tuple{Vararg{Tuple{Int,Int}}}}}
            observables = Tuple{Any, Tuple{Vararg{Tuple{Int,Int}}}, Vector{Tuple{Int,Int}}}[]
            for (O, rsidx) in observablesval
                links = Tuple{Int,Int}[]
                if length(rsidx)>1
                    for (i, (iʳᵉᵍ, iˢˡᵒᵗ)) in enumerate(rsidx)
                        push!(links, (iʳᵉᵍ, iˢˡᵒᵗ))
                        i == 1 || i == length(rsidx) || push!(links, (iʳᵉᵍ, iˢˡᵒᵗ))
                    end
                end
                push!(observables, (O, rsidx, links))
            end
            rn.observables[] = observables
        elseif observablesval isa Vector{<:Tuple{Any,Tuple{Vararg{Tuple{Int,Int}}},Vector{Tuple{Int,Int}}}}
            # the expected most general format
        else
            throw(ArgumentError("While plotting a network layout an incorrect argument was given: `observables` has to be of type `Vector{<:Tuple{Any, Tuple{...}}}`, i.e. it has to be a vector in which each element is similar to `(X⊗X, ((1,1), (1,2)))`, giving the observable operator and the `(register, slot)` indices for each observed subsystem. There is also a support for adding a third tuple element, a vector specifying the exact links to be drawn."))
        end
    end

    # this handles the majority of conversions from input data to graphics coordinates/metadata
    function update_plot(network)
        registers = network.registers
        registercoords = rn[:registercoords][]
        all_nodes = [ # TODO it is rather wasteful to replot everything... do it smarter
            register_rectangles, register_slots_coords,
            register_slots_coords_backref, state_coords_backref,
            observables_coords, observables_links, observables_vals, observables_linkvals
        ]
        for a in all_nodes # using a naive `lift` would allocate, so instead we just empty each array and refill it; can still be done more elegantly with lift and preallocation
            empty!(a[])
        end
        state_coord_idx = 0
        lock_coord_idx = 0
        state_link_idx = 0

        # the location of the registers and the slots inside of the registers
        for (iʳᵉᵍ,reg) in enumerate(registers)
            xʳᵉᵍ  = registercoords[iʳᵉᵍ][1]-0.3*rn[:scale][]
            yʳᵉᵍ  = registercoords[iʳᵉᵍ][2]-0.3*rn[:scale][]
            Δxʳᵉᵍ = 0.6*rn[:scale][]
            Δyʳᵉᵍ = (nsubsystems(reg)-0.4)*rn[:scale][]
            push!(register_rectangles[], Rect2f(xʳᵉᵍ, yʳᵉᵍ, Δxʳᵉᵍ, Δyʳᵉᵍ))
            for iˢˡᵒᵗ in 1:nsubsystems(reg)
                xˢˡᵒᵗ = registercoords[iʳᵉᵍ][1]
                yˢˡᵒᵗ = registercoords[iʳᵉᵍ][2]+(iˢˡᵒᵗ-1)*rn[:scale][]
                push!(register_slots_coords[], Point2f(xˢˡᵒᵗ,yˢˡᵒᵗ))
                push!(register_slots_coords_backref[], (reg,iʳᵉᵍ,iˢˡᵒᵗ))
                if reg.locks[iˢˡᵒᵗ].level >= 1
                    lock_coord_idx += 1
                    lock_coords[][lock_coord_idx] = Point2f(xˢˡᵒᵗ, yˢˡᵒᵗ)
                end
            end
        end

        # the locations of the state subsystem markers and the lines connecting them (denoting belonging to the same system)
        states = unique(Iterators.flatten(((s for s in r.staterefs if !isnothing(s)) for r in registers)))
        for s in states
            for (iˢ,(reg,iˢˡᵒᵗ)) in enumerate(zip(s.registers,s.registerindices))
                isnothing(reg) && continue # TODO -- the state does not belong to a register... this should be a warning
                whichreg = findfirst(o->===(reg,o),registers) # TODO -- some form of caching or a backref would be valuable to significantly optimize this (skip the need for a O(n) search)
                isnothing(whichreg) && continue # TODO -- the state does not belong to a register in the network... maybe it is in a temporary message buffer register?
                xˢ = registercoords[whichreg][1]
                yˢ = registercoords[whichreg][2]+(iˢˡᵒᵗ-1)*rn[:scale][]
                pˢ = Point2f(xˢ, yˢ)
                state_coord_idx += 1
                state_coords[][state_coord_idx] = pˢ
                push!(state_coords_backref[], (s, network[whichreg], whichreg, iˢˡᵒᵗ, iˢ))
                if nsubsystems(s) != 1
                    state_link_idx += 1
                    state_links[][state_link_idx] = pˢ
                end
                if !(iˢ == 1 || iˢ == nsubsystems(s))
                    state_link_idx += 1
                    state_links[][state_link_idx] = pˢ
                end
            end
        end

        ## the colors and locations for various observables
        if !isnothing(rn[:observables][])
        for (O, rsidx, links) in rn[:observables][]
            val = real(observable(tuple((network[rs...] for rs in rsidx)...), O; something=NaN))
            # TODO issue a warning if val has (percentage-wise) significant imaginary component (here, for plotting, when we implicitly are taking the real part)
            for (iʳᵉᵍ, iˢˡᵒᵗ) in rsidx
                xˢ = registercoords[iʳᵉᵍ][1]
                yˢ = registercoords[iʳᵉᵍ][2]+(iˢˡᵒᵗ-1)*rn[:scale][]
                pˢ = Point2f(xˢ, yˢ)
                push!(observables_coords[], pˢ)
                push!(observables_vals[], val)
                push!(observables_backref[], (O, val))
            end
            for (iʳᵉᵍ, iˢˡᵒᵗ) in links
                xˢ = registercoords[iʳᵉᵍ][1]
                yˢ = registercoords[iʳᵉᵍ][2]+(iˢˡᵒᵗ-1)*rn[:scale][]
                pˢ = Point2f(xˢ, yˢ)
                push!(observables_links[], pˢ)
                push!(observables_linkvals[], val)
                push!(observables_links_backref[], (O, val))
            end
        end
        end

        for i in (state_coord_idx + 1):length(state_coords[])
            state_coords[][i] = Point2f(NaN, NaN)
        end
        notify(state_coords)

        for i in (lock_coord_idx + 1):length(lock_coords[])
            lock_coords[][i] = Point2f(NaN, NaN)
        end
        notify(lock_coords)

        for i in (state_link_idx + 1):length(state_links[])
            state_links[][i] = Point2f(NaN, NaN)
        end
        notify(state_links)

        for a in all_nodes
            notify(a)
        end
    end

    # populate all graphical coordinates / metadata for the first time
    update_plot(networkobs[])

    # set up event modification notifications
    Makie.Observables.onany(update_plot, networkobs)
    Makie.Observables.onany([rn[:registercoords], rn[:observables]]) do _
        update_plot(networkobs)
    end

    # generate the actual graphics
    register_polyplot = poly!(rn, register_rectangles, color=rn[:register_color],
        inspector_label = (self, i, p) -> "a register")
    register_polyplot.inspectable[] = false # TODO this `Poly` plot does not seem to be properly inspectable
    register_slots_scatterplot = scatter!(
        rn, register_slots_coords,
        marker=rn[:slotmarker], markersize=rn[:slotsize][]*rn[:scale][], color=slotcolor_resolved,
        markerspace=:data,
        inspector_label = (self, i, p) -> get_slots_vis_string(register_slots_coords_backref[],i))
    observables_scatterplot = scatter!(
        rn, observables_coords,
        marker=rn[:observables_marker], markersize=rn[:observables_markersize][]*rn[:scale][], markerspace=:data,
        color=observables_vals, colormap=rn[:colormap], colorrange=rn[:colorrange],
        inspector_label = (self, i, p) -> get_observables_vis_string(observables_backref[],i))
    observables_linesegments = linesegments!(
        rn, observables_links,
        linewidth=rn[:observables_linewidth][]*rn[:scale][],
        color=observables_linkvals, colormap=rn[:colormap], colorrange=rn[:colorrange],
        inspector_label = (self, i, p) -> get_observables_vis_string(observables_links_backref[],i))
    state_scatterplot = scatter!(
        rn, state_coords,
        marker=rn[:state_marker], markersize=rn[:state_markersize][]*rn[:scale][], color=rn[:state_markercolor],
        markerspace=:data,
        inspector_label = (self, i, p) -> get_state_vis_string(state_coords_backref[],i))
    state_linesegmentsplot = linesegments!(rn, state_links, color=rn[:state_linecolor])
    state_linesegmentsplot.inspectable[] = false
    lock_scatterplot = scatter!(
        rn, lock_coords,
        marker=rn[:lock_marker], markersize=rn[:slotsize][]*rn[:scale][],
        markerspace=:data)
    lock_scatterplot.inspectable[] = false

    # TODO all of these should be wrapped into their own types in order to simplify DataInspector and process_interaction
    _extras[:register_polyplot] = register_polyplot
    _extras[:register_slots_scatterplot] = register_slots_scatterplot
    _extras[:observables_scatterplot] = observables_scatterplot
    _extras[:observables_linesegmentsplot] = observables_linesegments
    _extras[:state_scatterplot] = state_scatterplot
    _extras[:state_linesegmentsplot] = state_linesegmentsplot
    _extras[:lock_scatterplot] = lock_scatterplot
    rn[:_extras] = _extras
    rn
end

function get_observables_vis_string(backrefs, i)
    o, val = backrefs[i]
    return "Observable ⟨$(o)⟩ = $(val)"
end

function get_slots_vis_string(backrefs, i)
    register, registeridx, slot = backrefs[i]
    tags = QuantumSavory.peektags(register[slot])
    tags_str = if isempty(tags)
        "not tagged"
    else
        "tagged with:\n"*join((" • $(t)" for t in tags), "\n")
    end
    return "Register $(registeridx) | Slot $(slot)\n $(tags_str)"
end

function get_state_vis_string(backrefs, i)
    state, register, registeridx, slot, subsystem = backrefs[i]
    tags = [ti.tag for ti in values(register.tag_info) if ti.slot==slot]
    tags_str = if isempty(tags)
        "not tagged"
    else
        "tagged with:\n"*join((" • $(t)" for t in tags), "\n")
    end
    return "Subsystem $(subsystem) of a state of $(nsubsystems(state)) subsystems, stored in\nRegister $(registeridx) | Slot $(slot)\n $(tags_str)"
end

abstract type RegisterNetGraphHandler end
struct RNHandler <: RegisterNetGraphHandler
    rn
end

function Makie.process_interaction(handler::RNHandler, event::Makie.MouseEvent, axis)
    plot, index = Makie.pick(axis.scene)
    rn = handler.rn
    extras = rn[:_extras][]
    #if plot===extras[:register_polyplot]             # TODO this does not work because poly seems to be much too basic for `pick` to provide a useful reference
    #    register = extras[:register_backref][][index]
    #    run(`clear`)
    #    println("$(register)")
    #else
    if plot===extras[:register_slots_scatterplot]
        register, registeridx, slot = extras[:register_slots_coords_backref][][index]
        try run(`clear`) catch end
        println("Register $registeridx | Slot $(slot)\n Details: $(register)")
    elseif plot===extras[:state_scatterplot]
        state, reg, registeridx, slot, subsystem = extras[:state_coords_backref][][index]
        try run(`clear`) catch end
        println("Subsystem stored in Register $(registeridx) | Slot $(slot)\n Subsystem $(subsystem) of $(state)")
    elseif plot===extras[:observables_scatterplot]
        o, val = extras[:observables_backref][][index]
        try run(`clear`) catch end
        println("Observable $(o) has value $(val)")
    end
    false
end

##

"""Draw the given registers on a given Makie axis or a subfigure.

It returns a tuple of (subfigure, axis, plot, observable).
The observable can be used to issue a `notify` call that updates
the plot with the current state of the network."""
function registernetplot_axis end

function registernetplot_axis(ax::Makie.AbstractAxis, registersobservable; infocli=true, datainspector=true, autolimits=false, kwargs...)
    registersobservable = registersobservable isa Observable ? registersobservable : Observable(registersobservable)
    p = registernetplot!(ax, registersobservable; kwargs...)
    ax.aspect = Makie.DataAspect()
    if hasmethod(Makie.hidedecorations!, Tuple{typeof(ax)})
        Makie.hidedecorations!(ax)
    end
    if hasmethod(Makie.hidespines!, Tuple{typeof(ax)})
        Makie.hidespines!(ax)
    end
    Makie.deregister_interaction!(ax, :rectanglezoom)
    if infocli
        rnh = RNHandler(p)
        Makie.register_interaction!(ax, :registernet, rnh)
    end
    if datainspector
        DataInspector(ax.parent)
    end
    if autolimits
        if hasmethod(Makie.autolimits!, Tuple{typeof(ax)})
            Makie.autolimits!(ax)
        end
    end
    # translating the plot so that it is in front of the background map
    Makie.translate!(p, 0, 0, 10)
    ax.parent, ax, p, registersobservable
end

# subfig::Union{GridPosition, GridSubposition} but maybe other as well, so leave it unspecified
function registernetplot_axis(subfig, registersobservable; infocli=true, datainspector=true, kwargs...)
    registernetplot_axis(Makie.Axis(subfig[1,1]), registersobservable; infocli, datainspector, kwargs...)
end

function registernetplot_axis(registersobservable; infocli=true, datainspector=true, kwargs...)
    fig = Figure()
    ax = Axis(fig[1, 1])
    registernetplot_axis(ax, registersobservable; infocli, datainspector, kwargs...)
end

##

showonplot(r::ConcurrentSim.Resource) = islocked(r)
showonplot(b::Bool) = b

"""Draw the various resources and locks stored in the given meta-graph on a given Makie axis.

It returns a tuple of (subfigure, axis, plot, observable).
The observable can be used to issue a `notify` call that updates
the plot with the current state of the network."""
function resourceplot_axis(subfig, network, edgeresources, vertexresources; registercoords=nothing, title="")
    axis = Makie.Axis(subfig[1,1], title=title)
    networkobs = Makie.Observable(network)
    if isnothing(registercoords)
        # copied from registernetplot
        adj_matrix = adjacency_matrix(networkobs[].graph)
        registercoords = spring(adj_matrix, iterations=40, C=2*maximum(nsubsystems.(network.registers)))
    else
        registercoords = registercoords[]
    end
    baseplot = Makie.scatter!(axis, # just to set coordinates
            registercoords,
            color=:gray90,
            )
    for (i,vertexres) in enumerate(vertexresources)
        Makie.scatter!(axis,
            Makie.lift(x->Makie.Point2{Float64}[registercoords[j].+0.2*(i-1) for j in vertices(x) if showonplot(x[j,vertexres])],
                networkobs),
            color=Makie.Cycled(i),
            label="$(vertexres)"
            )
    end
    for (i,edgeres) in enumerate(edgeresources)
        Makie.linesegments!(axis,
            Makie.lift(x->Makie.Point2{Float64}[registercoords[n].+0.2*(i-1) for (;dst,src) in edges(x) for n in (dst,src) if showonplot(x[(dst,src),edgeres])],
                networkobs),
            color=Makie.Makie.Cycled(i),
            label="$(edgeres)"
            )
    end
    axis.aspect = Makie.DataAspect()
    Makie.Legend(subfig[2,1],axis,tellwidth = false, tellheight = true, orientation=:horizontal)
    Makie.hidedecorations!(axis)
    Makie.hidespines!(axis)
    Makie.autolimits!(axis)
    subfig, axis, baseplot, networkobs
end

##

function showmetadata(fig, ax, p, reg, slot)
    Makie.events(fig).mouseposition[] =
    tuple(Makie.shift_project(ax.scene, p.registercoords[][reg].+(0,slot-1))...);
end

##

include("state_explorer.jl")
include("show_state.jl")
include("show_protocol.jl")

end
