using MetagraphOptimization
using DataFrames
using CSV
using BenchmarkTools
using StatsBase

results_filename = "qed_gen_results_$(Threads.nthreads()).csv"

df = DataFrame(
    process_name = String[],
    cpu_threads = Int[],
    graph_gen_samples = Int[],
    graph_gen_mean = Float64[],
    graph_gen_std = Float64[],
    graph_gen_median = Float64[],
    graph_nodes = Int[],
    graph_data_nodes = Int[],
    graph_u_nodes = Int[],
    graph_v_nodes = Int[],
    graph_s1_nodes = Int[],
    graph_s2_nodes = Int[],
    graph_edges = Int[],
    graph_nodes_reduced = Int[],
    graph_data_nodes_reduced = Int[],
    graph_u_nodes_reduced = Int[],
    graph_v_nodes_reduced = Int[],
    graph_s1_nodes_reduced = Int[],
    graph_s2_nodes_reduced = Int[],
    graph_edges_reduced = Int[],
    graph_mem = Float64[],
    graph_mem_reduced = Float64[],
    graph_elapsed_reduce = Float64[],
)

function bench_process(process::AbstractString; warmup = false, optimize = true)
    println("Benchmarking $process...")
    model = QEDModel()

    proc = parse_process(process, model)

    gen_bench = @benchmark gen_graph($proc) gcsample = true seconds = 5

    graph = gen_graph(proc)

    props = GraphProperties(graph)
    node_dict = countmap(typeof.(graph.nodes))
    graph_size = Base.summarysize(graph)

    reduce_elapsed = -1.0
    node_dict_reduced = Dict()
    graph_size_reduced = -1.0
    props_reduced = GraphProperties()
    if optimize
        reduce_elapsed = @elapsed optimize_to_fixpoint!(ReductionOptimizer(), graph)

        props_reduced = GraphProperties(graph)
        node_dict_reduced = countmap(typeof.(graph.nodes))
        graph_size_reduced = Base.summarysize(graph)
    end

    if warmup
        return nothing
    end

    push!(
        df,
        Dict(
            :process_name => process,
            :cpu_threads => Threads.nthreads(),
            :graph_gen_samples => length(gen_bench.times),
            :graph_gen_mean => mean(gen_bench.times),
            :graph_gen_std => std(gen_bench.times),
            :graph_gen_median => median(gen_bench.times),
            :graph_nodes => props.noNodes,
            :graph_data_nodes => get(node_dict, DataTaskNode{DataTask}, 0),
            :graph_u_nodes => get(node_dict, ComputeTaskNode{ComputeTaskQED_U}, 0),
            :graph_v_nodes => get(node_dict, ComputeTaskNode{ComputeTaskQED_V}, 0),
            :graph_s1_nodes => get(node_dict, ComputeTaskNode{ComputeTaskQED_S1}, 0),
            :graph_s2_nodes => get(node_dict, ComputeTaskNode{ComputeTaskQED_S2}, 0),
            :graph_edges => props.noEdges,
            :graph_nodes_reduced => props_reduced.noNodes,
            :graph_data_nodes_reduced => get(node_dict_reduced, DataTaskNode{DataTask}, 0),
            :graph_u_nodes_reduced => get(node_dict_reduced, ComputeTaskNode{ComputeTaskQED_U}, 0),
            :graph_v_nodes_reduced => get(node_dict_reduced, ComputeTaskNode{ComputeTaskQED_V}, 0),
            :graph_s1_nodes_reduced => get(node_dict_reduced, ComputeTaskNode{ComputeTaskQED_S1}, 0),
            :graph_s2_nodes_reduced => get(node_dict_reduced, ComputeTaskNode{ComputeTaskQED_S2}, 0),
            :graph_edges_reduced => props_reduced.noEdges,
            :graph_mem => graph_size,
            :graph_mem_reduced => graph_size_reduced,
            :graph_elapsed_reduce => reduce_elapsed,
        ),
    )
    return nothing
end

processes = [
    ("ke->ke", true),
    ("ke->kke", true),
    ("ke->kkke", true),
    ("ke->kkkke", true),
    ("ke->kkkkke", true),
    ("ke->kkkkkke", true),
    ("ke->kkkkkkke", true),
    #("ke->kkkkkkkke", false),
    #("ke->kkkkkkkkke", false),
]

df = DataFrame(
    process_name = String[],
    cpu_threads = Int[],
    graph_gen_samples = Int[],
    graph_gen_mean = Float64[],
    graph_gen_std = Float64[],
    graph_gen_median = Float64[],
    graph_nodes = Int[],
    graph_data_nodes = Int[],
    graph_u_nodes = Int[],
    graph_v_nodes = Int[],
    graph_s1_nodes = Int[],
    graph_s2_nodes = Int[],
    graph_edges = Int[],
    graph_nodes_reduced = Int[],
    graph_data_nodes_reduced = Int[],
    graph_u_nodes_reduced = Int[],
    graph_v_nodes_reduced = Int[],
    graph_s1_nodes_reduced = Int[],
    graph_s2_nodes_reduced = Int[],
    graph_edges_reduced = Int[],
    graph_mem = Float64[],
    graph_mem_reduced = Float64[],
    graph_elapsed_reduce = Float64[],
)

# if they exist, read existing results and append new ones
if isfile(results_filename)
    df = CSV.read(results_filename, DataFrame)
end

bench_process("ke->kke", warmup = true)

for (process, opt) in processes
    bench_process(process, optimize = opt)
    CSV.write(results_filename, df)
end
