using MetagraphOptimization
using LIKWID
using UUIDs
using DataFrames
using CSV
using Random
using BenchmarkTools
using Dates

function log(x...)
    println(now(), " ", join(x, " ")...)
    return flush(stdout)
end

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

df = DataFrame(
    process_name = String[],
    graph_gen_time = Float64[],
    optimization_time = Float64[],
    function_generation_time = Float64[],
    graph_nodes = Int[],
    graph_edges = Int[],
    graph_mem = Float64[],
    cpu_threads = Int[],
    n_inputs = Int[],
    nflops_likwid = Int[],
    cpu_time = Float64[],
    cpu_rate = Float64[],
    cpu_gflops = Float64[],
    cpu_std = Float64[],
    gpu_name = String[],
    gpu_time = Float64[],
    gpu_std = Float64[],
    gpu_rate = Float64[],
    gpu_gflops = Float64[],
)

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

nInputs = 1_000_000

# use "mock" machine that only uses cpu
machine = Machine(
    [
        MetagraphOptimization.NumaNode(
            0,
            1,
            MetagraphOptimization.default_strategy(MetagraphOptimization.NumaNode),
            -1.0,
            UUIDs.uuid1(),
        ),
    ],
    [-1.0;;],
)


function cpu_bench(tape, inputs)
    bench = @benchmark begin
        @inbounds Threads.@threads for i in eachindex($inputs)
            execute_tape($tape, $inputs[i])
        end
    end gcsample = true seconds = 300

    time = mean(bench.times) / 1e9
    s = std(bench.times) / 1e9
    rate = length(inputs) / time

    return (time, rate, s)
end

function bench_process(
    process::MetagraphOptimization.AbstractProcessDescription,
    process_name::String,
    graph::DAG,
    gen_time::Float64,
    opt_time::Float64,
    io::IO = stdout;
    use_likwid = false,
)
    log("\n--- Benchmarking $(process_name) ---")

    func_time = @elapsed tape = gen_tape(graph, process, machine)

    graph_props = GraphProperties(graph)
    NFLOPs = graph_props.computeEffort
    nflops_likwid = 0
    if use_likwid
        input = gen_process_input(process)

        # get rid of annoying output to console
        oldstd = stdout
        redirect_stdout(devnull)
        _, events = @perfmon "FLOPS_DP" execute_tape(tape, input)
        redirect_stdout(oldstd) # recover original stdout

        NFLOPs = first(events["FLOPS_DP"])["RETIRED_SSE_AVX_FLOPS_ALL"]
        nflops_likwid = NFLOPs
    end

    log("Generating $nInputs inputs with $(Threads.nthreads()) threads...")

    inputs = Vector{typeof(gen_process_input(process))}()
    resize!(inputs, nInputs)
    processes = Vector{typeof(process)}()
    for i in 1:Threads.nthreads()
        push!(processes, copy(process))
    end

    @inbounds Threads.@threads for i in eachindex(inputs)
        inputs[i] = gen_process_input(processes[Threads.nthreads()])
    end

    log("Benchmarking CPU with $(Threads.nthreads()) threads...")
    (time_cpu, rate_cpu, std_cpu) = cpu_bench(tape, inputs)
    flops_cpu = (rate_cpu * NFLOPs) / 10^9

    log("\nBenchmark Summary for $(process):")

    if use_likwid
        log("Measured FLOPS by LIKWID: $NFLOPs")
    else
        log("Total graph compute effort: $NFLOPs")
    end
    log("Total input size: $(bytes_to_human_readable(Base.summarysize(inputs)))")
    log("CPU, $(Threads.nthreads()) threads")
    log("  Time:  $time_cpu")
    log("  Rate:  $rate_cpu")
    log("  GFLOPS: $flops_cpu")

    if (process_name != "warmup")
        push!(
            df,
            Dict(
                :process_name => process_name,
                :graph_gen_time => gen_time,
                :optimization_time => opt_time,
                :function_generation_time => func_time,
                :graph_nodes => graph_props.noNodes,
                :graph_edges => graph_props.noEdges,
                :graph_mem => MetagraphOptimization.mem(graph),
                :cpu_threads => Threads.nthreads(),
                :n_inputs => nInputs,
                :nflops_likwid => nflops_likwid,
                :cpu_time => time_cpu,
                :cpu_std => std_cpu,
                :cpu_rate => rate_cpu,
                :cpu_gflops => flops_cpu,
                :gpu_name => "none",
                :gpu_time => 0.0,
                :gpu_std => 0.0,
                :gpu_rate => 0.0,
                :gpu_gflops => 0.0,
            ),
        )
    end

    return nothing
end

function bench_qed(process_string::String, skip_unoptimized = false)
    optimizer = ReductionOptimizer()

    process = parse_process(process_string, QEDModel())
    gen_time = @elapsed graph = gen_graph(process)
    opt_time = 0.0
    if !skip_unoptimized
        bench_process(process, "$process not optimized tape", graph, gen_time, opt_time)
    end

    opt_time = @elapsed optimize_to_fixpoint!(optimizer, graph)
    bench_process(process, "$process reduced tape", graph, gen_time, opt_time)

    return nothing
end

function bench_abc(process_string::String)
    optimizer = ReductionOptimizer()

    process = parse_process(process_string, ABCModel())
    gen_time = @elapsed graph = parse_dag("input/$process_string.txt", ABCModel())
    bench_process(process, "$process not optimized tape", graph, gen_time, 0.0)

    opt_time = @elapsed optimize_to_fixpoint!(optimizer, graph)
    bench_process(process, "$process reduced tape", graph, gen_time, opt_time)

    return nothing
end

# sadly cannot put these in functions because the world age must increase after the function is created which happens only in the global scope

## -- WARMUP TO COMPILE FUNCTIONS first
optimizer = ReductionOptimizer()

process = parse_process("ke->kke", QEDModel())
gen_time = @elapsed graph = gen_graph(process)
opt_time = @elapsed optimize_to_fixpoint!(optimizer, graph)
bench_process(process, "warmup", graph, gen_time, opt_time)

# AB->AB^3
process = parse_process("AB->ABBB", ABCModel())
gen_time = @elapsed graph = parse_dag("input/AB->ABBB.txt", ABCModel())
opt_time = @elapsed optimize_to_fixpoint!(optimizer, graph)
bench_process(process, "warmup", graph, gen_time, opt_time)

## -- WARMUP END

# compton
bench_qed("ke->ke")
CSV.write(results_filename, df)
bench_qed("ke->kke")
CSV.write(results_filename, df)
bench_qed("ke->kkke")
CSV.write(results_filename, df)
bench_qed("ke->kkkke")
CSV.write(results_filename, df)
bench_qed("ke->kkkkke")
CSV.write(results_filename, df)
bench_qed("ke->kkkkkke")
CSV.write(results_filename, df)
bench_qed("ke->kkkkkkke")
CSV.write(results_filename, df)

bench_abc("AB->AB")
CSV.write(results_filename, df)
bench_abc("AB->ABBB")
CSV.write(results_filename, df)
bench_abc("AB->ABBBBB")
CSV.write(results_filename, df)
