plot_program.jl

using GraphvizDotLang: digraph, node, edge, attr, save, subgraph
using GraphvizDotLang
using Serialization
# using ImageView
using PNGFiles

# load_folder = "ga_metrics/pong/bddcfc4c-4e09-483e-b46e-24edd159239b/checkpoint_2115.pickle"
# payloads = deserialize(load_folder);

# load_folder = "ga_metrics/pong/b42a276f-cbb1-4d29-be23-637dd4ac93c8/checkpoint_2270.pickle"
# payloads = deserialize(load_folder);

# freeway 
# load_folder = "ga_metrics/freeway/cf40dd45-6611-4e58-9526-43adb0f46824/checkpoint_820.pickle"
# payloads = deserialize(load_folder);

function plot_program(
    progs::UTCGP.IndividualPrograms,
    folder::String,
    names::Vector{String}
)
    W = "2"
    H = "2"
    isdir(folder) || mkpath(folder)
    g = digraph("G", rankdir="LR", bgcolor="#f5f5f5", ranksep="0.25", nodesep="0.25",
            # rankstep="2",
            layout="dot",
            ratio="0.4",
            center="2",
            compound="true",
            ordering="in",
            pad=".1",
        ) |>
        attr(:node; shape="rectangle", style="filled", fillcolor="#ffffff", fontname="Helvetica", fontsize="25", margin="0.1") |>
        attr(:edge, fontname="Helvetica", fontsize="23", color="#333333", constraint="false", margin="0.1")
    # g = digraph("G")
    clusters = []
    nodes = []
    first_nodes = []
    last_nodes = []
    n_outs = length(progs)
    cluster = subgraph(g, "cluster_inputs"; label="Inputs", color="#7393B3", style="dashed",
        rank="source", fontsize="30")
    # push!(clusters, cluster)
    all_nodes = []
    input_nodes = []
    for n_out in 1:n_outs
        # create a cluster for the program
        id = "cluster_prog$n_out"
        name = "Program $(names[n_out])"
        cluster_ = subgraph(g, id;
            label=name, color="#7393B3", style="dashed",
            rank="same", fontsize="30")
        push!(clusters, cluster_)
        push!(nodes, [])
        # Get first nodes
        first_calling_node = progs[n_out].program[begin].calling_node
        x = first_calling_node.x_position
        y = first_calling_node.y_position
        id = "$x $y"
        push!(first_nodes, id)

        # Get last nodes
        last_calling_node = progs[n_out].program[end].calling_node
        x = last_calling_node.x_position
        y = last_calling_node.y_position
        id = "$x $y"
        push!(last_nodes, id)
    end
    @show last_nodes

    for (ith_prog, pack) in enumerate(zip(progs.programs, clusters, nodes))
        program, current_g, n = pack
        # r = reverse(program.program)
        for op in program.program
            ### ADD THE CALLING NODE (result) ###
            calling_node = op.calling_node
            x = calling_node.x_position
            y = calling_node.y_position
            id = "$x $y"
            nd = nothing
            group = nothing
            form = "rectangle"
            if id in last_nodes
                group = "last_nodes"
                form = "hexagon"
            elseif id in first_nodes
                group = "first_nodes"
            else
                group = "$ith_prog"
            end
            if calling_node.value isa Number
                nd = GraphvizDotLang.node(id; label="$(round(calling_node.value, digits = 2))",
                    group=group,
                    # xlabel=id,
                    shape=form
                )
            elseif calling_node.value isa Tuple
                nd = GraphvizDotLang.node(id; label="($(calling_node.value[1]) $(calling_node.value[2]))")

            else
                tpath = "$folder/$id.png"
                if maximum(calling_node.value) < N0f8(0.5)
                    PNGFiles.save(tpath, (
                        calling_node.value ./ (maximum(calling_node.value) + eps(N0f8))
                    ))
                else
                    PNGFiles.save(tpath, calling_node.value)
                end
                nd = GraphvizDotLang.node(id; shape="box",
                    label="",
                    image=tpath,
                    width=W, height=H,
                    imagescale="true"
                )
            end
            if !(id in all_nodes)
                push!(all_nodes, id)
                push!(n, id)
                current_g |> nd
            end

            ### ADD THE INPUTS ###
            ins = []
            for op_in in op.inputs
                R_node = UTCGP._extract_input_node_from_operationInput(
                    program.program_inputs,
                    op_in,
                )
                node = @unwrap_or R_node throw(ErrorException("Could not extract the input from operation."))
                push!(ins, node)
            end

            for (ith_input, input) in enumerate(ins)
                x = input.x_position
                y = input.y_position
                edge_id = "$x $y"
                nd = nothing
                if input.value isa Number
                    nd = GraphvizDotLang.node(edge_id; label="$(round(input.value, digits = 2))")
                elseif input.value isa Tuple
                    nd = GraphvizDotLang.node(edge_id; label="($(input.value[1]) $(input.value[2]))")
                else
                    tpath = "$folder/$edge_id.png"
                    if maximum(input.value) < N0f8(0.5)
                        PNGFiles.save(tpath, (
                            input.value ./ (maximum(input.value) + eps(N0f8))
                        ))
                    else
                        PNGFiles.save(tpath, input.value)
                    end
                    # PNGFiles.save(tpath, input.value)
                    nd = GraphvizDotLang.node(edge_id; shape="box",
                        image=tpath, width=W, height=H, imagescale="true")
                end
                if !("$edge_id to $id" in all_nodes)
                    push!(all_nodes, "$edge_id to $id")
                    edge_name = ""
                    if ith_input == 1
                        edge_name *= replace(string(op.fn.name), "experimental_" => "", "_img2D" => "", "_image2D" => "", "_factory" => "")
                    end
                    current_g |> edge(edge_id, id; label=edge_name,
                        constraint="true",
                        weight="1", minlen="1",
                        #samehead="true" # ugly 
                    )
                end

                if input isa UTCGP.InputNode
                    if !(edge_id in input_nodes)
                        push!(input_nodes, edge_id)
                        cluster |> nd
                    end
                else
                    if !(edge_id in all_nodes)
                        push!(all_nodes, edge_id)
                        current_g |> nd
                    end
                end

            end
        end
    end

    # EDGES between input nodes so that they are aligned
    sort!(input_nodes)
    reverse!(input_nodes)
    @show length(input_nodes)
    prev_input = input_nodes[1]
    for next_input in input_nodes[2:end]
        @info "Inputs edge $prev_input, $next_input"
        g |> edge(prev_input, next_input; label="inputnode",
            constraint="true",
            weight="1",
            minlen="0.8",
            style="invis"
        )
        prev_input = next_input
    end

    # EDGES between CALLING NODES of the same program
    # sort!(n1)
    # reverse!(n1)
    for nodes_of_program in nodes
        prev_input = nodes_of_program[1]
        for next_input in nodes_of_program[2:end]
            # @info "Inputs edge $prev_input, $next_input"
            g |> edge(prev_input, next_input; label="PROG", constraint="true",
                weight="1",
                style="invis"
            )
            prev_input = next_input
        end
    end

    # EDGE CONNECTING THE FIRST TWO 
    for first_input in first_nodes[2:end]
        @info "Connecting first nodes"
        g |> edge(
            first_nodes[1],
            first_input;
            weight="100000000000000000000",
            minlen="0.01",
            contraint="true",
            label="FIRST",
            style="invis"
        )
    end

    # EDGES BETWEEN THE inputs and FIRST NODE
    # for input in input_nodes
    #     @info "Connecting to inputs"
    #     for first_input in first_nodes
    #         g |> edge(
    #             input,
    #             first_input;
    #             # weight="10",
    #             # l_tail="cluster_inputs",
    #             # l_head="cluster_prog1",
    #             # minlen="2", #e
    #             contraint="true",
    #             label="HOLDER",
    #             style="invis"
    #         )
    #     end
    # end

    # CONNECT THE LAST TWO
    # for last_node in last_nodes[2:end]
    #     @info "Connecting last nodes"
    #     g |> edge(
    #         last_nodes[1],
    #         last_node;
    #         weight="10000",
    #         minlen="0.0001",
    #         contraint="true",
    #         label="LAST",
    #         style="invis"
    #     )
    # end

    g, nodes
end

# function eval_plot(x, y, prog, ml, model_arch, shared_inputs, folder, file)
#     UTCGP.reset_programs!(prog)
#     input_nodes = [
#         InputNode(value, pos, pos, model_arch.inputs_types_idx[pos]) for
#         (pos, value) in enumerate(x)
#     ]
#     replace_shared_inputs!(prog, input_nodes)
#     outs = UTCGP.evaluate_individual_programs(
#         prog,
#         model_arch.chromosomes_types,
#         ml
#     )
#     @show outs
#     g, nodes = plot_program(prog, folder)
#     fn = "$(file)_$(y[1]).png"
#     save(g, fn)
# end

# progs = filter(i -> occursin("pickle", i), readdir("metrics_ga"))
# payloads = [
#     deserialize("metrics_ga/$i") for i in progs
# ];

# N images seen per run 
# n_images_per_run = @chain D begin
#     #@subset :filename .== best_on_val[1]
#     @subset :data .== "Val"
#     @groupby :filename
#     @subset :accuracy .== maximum(:accuracy)
#     @groupby :filename
#     @subset :iteration .== minimum(:iteration)
#     @select :iteration :filename :accuracy
#     #@combine :best = maximum(:accuracy)
# end
# [p["best_loss"] for p in payloads] |> argmax # 73
# best_per_run = @chain D begin
#     @subset :data .== "Val"
#     @groupby :filename
#     @combine :best = maximum(:accuracy)
# end
# best_on_val = deepcopy(best_per_run[73, :])
# n_images_per_run[n_images_per_run[:,"filename"] .== best_on_val[1],:] # gens of the best on val

# BEST 50% 
# med = median(best_per_run[:, "best"]) # mid 50% val point
# top50_ids = @chain best_per_run begin
#     @subset :best .>= med
#     @select :filename
# end
# top50_ids = top50_ids[:,1]
# function special_in(values, allowed_values)
#     ret = Bool[]
#     for v in values
#         push!(ret, v in allowed_values)
#     end
#     ret
# end

# test metrics -- 
# # best on val 
# [p["best_loss"] for p in payloads] |> argmax
# test_callback(best_programs[73])
# test_callback = acc_callback(testx, testy, "Test")
# test_acc = []
# for p in best_programs
#     f = test_callback(
#         p
#     )
#     push!(test_acc, f)
# end