using MetagraphOptimization
using CUDA
using UUIDs
using BenchmarkTools
using DataFrames
using CSV

results_filename = "bench_results_reduction_steps.csv"

df = DataFrame(
    threads = Int[],
    process = String[],
    operations = Int[],
    cumulative_optimization_time = Float64[],
    graph_nodes = Int[],
    graph_edges = Int[],
    graph_ce = Float64[],
    graph_dt = Float64[],
    graph_ci = Float64[],
    gen_func_t = Float64[],
    cpu_compile_t = Float64[],
    cpu_st_t = Float64[],
    cpu_mt_t = Float64[],
    gpu_compile_t = Float64[],
    gpu_t = Float64[],
)

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

function bench(func, inputs)
    compile_time = @elapsed func(inputs[1])

    single_thread = @benchmark $func.($inputs)
    multi_threaded = @benchmark Threads.@threads for i in eachindex($inputs)
        $func($inputs[i])
    end

    return (
        cpu_compile_time = compile_time,
        gpu_compile_time = 0.0,
        cpu_single_thread_time = mean(single_thread.times) / 1e9,
        cpu_multi_thread_time = mean(multi_threaded.times) / 1e9,
        gpu_time = 0.0,
    )
end

# preparation of machine
machine = Machine(
    [
        MetagraphOptimization.NumaNode(
            0,
            1,
            MetagraphOptimization.default_strategy(MetagraphOptimization.NumaNode),
            -1.0,
            UUIDs.uuid1(),
        ),
    ],
    [-1.0;;],
)

# bench and produce data
n_inputs = 50_000
optimizer = ReductionOptimizer()
processes = [("ke->kke", 5), ("ke->ke", 1), ("ke->kke", 1), ("ke->kkke", 1), ("ke->kkkke", 1), ("ke->kkkkke", 1)]

for (process_str, STEPSIZE) in processes
    n = 0
    opt_time_cum = 0

    process = parse_process(process_str, QEDModel())
    graph = gen_graph(process)
    inputs = [gen_process_input(process) for _ in 1:n_inputs]

    get_compute_function(graph, process, machine)

    while true
        func_gen_time = @elapsed func = get_compute_function(graph, process, machine)
        res = bench(func, inputs)

        graph_properties = get_properties(graph)
        push!(
            df,
            (
                Threads.nthreads(),
                process_str,
                n,
                opt_time_cum,
                graph_properties.noNodes,
                graph_properties.noEdges,
                graph_properties.computeEffort,
                graph_properties.data,
                graph_properties.computeIntensity,
                func_gen_time,
                res.cpu_compile_time,
                res.cpu_single_thread_time,
                res.cpu_multi_thread_time,
                res.gpu_compile_time,
                res.gpu_time,
            ),
        )
        CSV.write(results_filename, df)

        if fixpoint_reached(optimizer, graph)
            break
        end

        opt_time_cum += @elapsed optimize!(optimizer, graph, STEPSIZE)
        n += STEPSIZE
    end
end

CSV.write(results_filename, df)

for (process_str, STEPSIZE) in [("AB->AB", 1), ("AB->ABBB", 1), ("AB->ABBBBB", 1)]
    n = 0
    opt_time_cum = 0

    process = parse_process(process_str, ABCModel())
    graph = parse_dag("input/$process_str.txt", ABCModel())
    inputs = [gen_process_input(process) for _ in 1:n_inputs]

    get_compute_function(graph, process, machine)

    while true
        func_gen_time = @elapsed func = get_compute_function(graph, process, machine)
        res = bench(func, inputs)

        graph_properties = get_properties(graph)
        push!(
            df,
            (
                Threads.nthreads(),
                process_str,
                n,
                opt_time_cum,
                graph_properties.noNodes,
                graph_properties.noEdges,
                graph_properties.computeEffort,
                graph_properties.data,
                graph_properties.computeIntensity,
                func_gen_time,
                res.cpu_compile_time,
                res.cpu_single_thread_time,
                res.cpu_multi_thread_time,
                res.gpu_compile_time,
                res.gpu_time,
            ),
        )
        CSV.write(results_filename, df)

        if fixpoint_reached(optimizer, graph)
            break
        end

        opt_time_cum += @elapsed optimize!(optimizer, graph, STEPSIZE)
        n += STEPSIZE
    end
end

CSV.write(results_filename, df)
