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

using Base.Threads


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

results_filename = "full_node_bench.csv"

df = DataFrame(
    process_name = String[],
    cpu_threads = Int[],
    gpu_devices = Int[],
    n_inputs = Int[],
    chunk_size = Int[],
    time = Float64[],
    std = Float64[],
    rate = Float64[],
    cpu_chunks = Float64[],
    gpu_chunks = Float64[],
    memory_est = Float64[],
)

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

nInputs = 2^26

lck = ReentrantLock()

progress = 1
cpu_chunks = 0
gpu_chunks = 0

chunkSizes = [1024, 4096, 16384, 65536, 262144, 1048576]    # 2^10 to 2^20

function cpu_worker(compute_func, inputs, chunk_size)
    global progress
    global cpu_chunks
    global lck
    quit = false
    work_start = 0
    work_end = 0
    while true
        lock(lck) do
            if progress >= nInputs
                quit = true
            else
                work_start = progress
                progress = progress + chunk_size
                work_end = min(progress - 1, nInputs)
                cpu_chunks = cpu_chunks + 1
                #log("CPU Worker $(Threads.threadid()) computing $(cpu_chunks)th cpu chunk ($work_start, $work_end)")
            end
        end
        if quit
            break
        end

        for i in work_start:work_end
            compute_func(inputs[i])
        end
    end

    #log("CPU Worker on $(Threads.threadid()) finished!")

    return nothing
end

# called with a specific device selected
function gpu_worker(kernel!, inputs, chunk_size)
    global progress
    global gpu_chunks
    global lck
    cuOutputs = CuVector{ComplexF64}()
    resize!(cuOutputs, chunk_size)

    quit = false
    work_start = 0
    work_end = 0
    while true
        lock(lck) do
            if progress >= nInputs
                quit = true
            else
                work_start = progress
                progress = progress + chunk_size
                work_end = min(progress - 1, nInputs)
                gpu_chunks = gpu_chunks + 1
                #log("GPU Worker $(CUDA.device()) computing $(gpu_chunks)th gpu chunk ($work_start, $work_end)")
            end
        end
        if quit
            break
        end

        cuInputs = CuVector(inputs[work_start:work_end])
        ts = 32
        bs = Int(chunk_size / 32)
        @cuda threads = ts blocks = bs always_inline = true kernel!(cuInputs, cuOutputs, chunk_size)
        CUDA.device_synchronize()
    end

    #log("GPU Worker on Device $(CUDA.device()) finished!")

    return nothing
end

cpu_gpu_ratio = Vector{Tuple{Int, Int}}()

function full_compute(compute_func, kernel!, inputs, chunk_size)
    global progress
    progress = 1
    global cpu_chunks
    cpu_chunks = 0
    global gpu_chunks
    gpu_chunks = 0

    tasks = Vector()

    for dev in CUDA.devices()
        t = Threads.@spawn device!(dev) do
            gpu_worker(kernel!, inputs, chunk_size)
            return nothing
        end
        push!(tasks, t)
    end

    for i in 1:(Threads.nthreads() - length(CUDA.devices()))
        t = Threads.@spawn cpu_worker(compute_func, inputs, chunk_size)
        push!(tasks, t)
    end

    for t in tasks
        wait(t)
    end

    push!(cpu_gpu_ratio, (cpu_chunks, gpu_chunks))
    return nothing
end

function bench(compute_function, kernel!, inputs, chunk_size)
    global cpu_gpu_ratio
    empty!(cpu_gpu_ratio)

    bench = @benchmark begin
        full_compute($compute_function, $kernel!, $inputs, $chunk_size)
    end gcsample = true seconds = 60

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

    med_cpu_chunks = median(getindex.(cpu_gpu_ratio, 1))
    med_gpu_chunks = median(getindex.(cpu_gpu_ratio, 2))
    mem_estimate = bench.memory

    log("CPU/GPU ratios: $(cpu_gpu_ratio)")

    return (time, rate, s, med_cpu_chunks, med_gpu_chunks, mem_estimate)
end

function full_node_bench(process::MetagraphOptimization.AbstractProcessDescription, func, kernel!, chunk_size, inputs)
    process_name = string(process)
    log("\n--- Benchmarking $(process_name) on $(nInputs) with chunk size $(chunk_size) ---")

    log("Available Cuda Devices:")
    display.(CUDA.devices())

    log("Benchmarking full node...")
    (time, rate, s, med_cpu_chunks, med_gpu_chunks, mem_estimate) = bench(func, kernel!, inputs, chunk_size)
    log(
        "Benchmarking complete with median time $(time), $(med_cpu_chunks) cpu chunks, and $(med_gpu_chunks) gpu chunks.",
    )

    push!(
        df,
        Dict(
            :process_name => process_name,
            :cpu_threads => Threads.nthreads() - length(CUDA.devices()),
            :gpu_devices => length(CUDA.devices()),
            :n_inputs => nInputs,
            :chunk_size => chunk_size,
            :time => time,
            :std => s,
            :rate => rate,
            :cpu_chunks => med_cpu_chunks,
            :gpu_chunks => med_gpu_chunks,
            :memory_est => mem_estimate,
        ),
    )

    return nothing
end

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

optimizer = ReductionOptimizer()
processes = ["ke->ke", "ke->kke", "ke->kkke", "ke->kkkke", "ke->kkkkke"]

for proc in processes
    process = parse_process(proc, QEDModel())
    graph = gen_graph(process)
    optimize_to_fixpoint!(optimizer, graph)
    compute_func = get_compute_function(graph, process, machine)
    kernel! = get_cuda_kernel(graph, process, machine)

    log("Generating $nInputs inputs with $(Threads.nthreads()) threads...")
    inputs = Vector{typeof(gen_process_input(process))}()
    resize!(inputs, nInputs)
    procs = Vector{typeof(process)}()
    for i in 1:Threads.nthreads()
        push!(procs, copy(process))
    end

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

    for chunk_size in chunkSizes
        full_node_bench(process, compute_func, kernel!, chunk_size, inputs)
        CSV.write(results_filename, df)
    end
end;
