using CSV
using DataFrames
using Plots
using StatsPlots
using LaTeXStrings

if (length(ARGS) < 1)
    println("Please use with \"input_file.csv\"")
end

processes = [
    "QED Process: 'ke->ke'",
    "QED Process: 'ke->kke'",
    "QED Process: 'ke->kkke'",
    "QED Process: 'ke->kkkke'",
    "QED Process: 'ke->kkkkke'",
    "ABC Process: 'AB->AB'",
    "ABC Process: 'AB->ABBB'",
    "ABC Process: 'AB->ABBBBB'",
]

function proc_to_n(str::AbstractString)
    parts = split(str, "'")
    parts = split(parts[2], "->")
    k_count = count(c -> c == 'k', parts[2])
    return k_count
end

function beautify_title(str::AbstractString)
    parts = split(str, "'")

    preprefix = parts[1]
    infix = parts[2]
    sufsuffix = parts[3]

    parts = split(infix, "->")

    prefix = parts[1]
    suffix = parts[2]

    k_count = count(c -> c == 'k', suffix)
    B_count = count(c -> c == 'B', suffix)

    if k_count == 1 || B_count == 1
        new_suffix = suffix
    elseif k_count >= 1
        new_suffix = replace(suffix, r"k+" => "k^$k_count")
    elseif B_count >= 1
        new_suffix = replace(suffix, r"B+" => "B^$B_count")
    end

    return preprefix * L"%$prefix \rightarrow %$new_suffix" * sufsuffix
end

input_file = ARGS[1]
df = CSV.read(input_file, DataFrame)
n_inputs = df[:, "n_inputs"][1]
gpu_name = df[:, "gpu_name"][1]
if (gpu_name == "")
    println("Results file did not execute everything on GPU! (or didn't write gpu name)")
    exit(0)
end

# plotting with process size as x axis
title_string = "GPU $gpu_name, $n_inputs samples"

df_filt = filter(:process_name => x -> proc_to_n(x) >= 1, df)
df_filt.gpu_rate = df_filt.gpu_rate
df_filt.gpu_time = df_filt.gpu_time
df_filt.gpu_gflops = df_filt.gpu_gflops

df_filt.process_size = @. proc_to_n(df_filt.process_name)

df_no_opt = filter(:process_name => x -> match(r" not optimized$", x) !== nothing, df_filt)
df_red = filter(:process_name => x -> match(r" reduced$", x) !== nothing, df_filt)

@df df_no_opt scatter(:process_size, :gpu_rate, label = "unoptimized function execution rate", markersize = 7)
@df df_red scatter!(:process_size, :gpu_rate, label = "reduced function execution rate", markersize = 7)

plot!(
    #title = title_string * ", sample rate",
    yscale = :log10,
    legend = :outerbottom,
    xticks = [1, 2, 3, 4, 5],
    legendcolumns = 2,
    legend_font_pointsize = 10,
    size = (800, 600),
    ylabel = "rate (" * L"s^{-1}" * ")",
    xlabel = "process size (#)",
)

savefig("gpu_rate_$(gpu_name).pdf")



@df df_no_opt scatter(:process_size, :gpu_time, label = "unoptimized function execution time", markersize = 7)
@df df_red scatter!(:process_size, :gpu_time, label = "reduced function execution time", markersize = 7)

plot!(
    #title = title_string * ", execution time",
    yscale = :log10,
    legend = :outerbottom,
    xticks = [1, 2, 3, 4, 5],
    legendcolumns = 2,
    legend_font_pointsize = 10,
    size = (800, 600),
    ylabel = "time (s)",
    xlabel = "process size (#)",
)

savefig("gpu_times_$(gpu_name).pdf")



@df df_no_opt scatter(:process_size, :gpu_gflops, label = "unoptimized function", markersize = 7)
@df df_red scatter!(:process_size, :gpu_gflops, label = "reduced function", markersize = 7)

plot!(
    #title = title_string * ", GFLOPS",
    yscale = :linear,
    legend = :outerbottom,
    xticks = [1, 2, 3, 4, 5],
    legendcolumns = 2,
    legend_font_pointsize = 10,
    size = (800, 600),
    ylabel = "performance (GFLOPS)",
    xlabel = "process size (#)",
)

savefig("gpu_perf_$(gpu_name).pdf")