using Plots
using StatsPlots
using CSV
using DataFrames
using LaTeXStrings


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


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, "->")

    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 L"%$prefix \rightarrow %$new_suffix"
end

processes = ["ke->ke", "ke->kke", "ke->kkke", "ke->kkkke", "ke->kkkkke", "AB->AB", "AB->ABBB", "AB->ABBBBB"]

input_file = ARGS[1]
input_file_gpu = ARGS[2]
df = CSV.read(input_file, DataFrame)
df_gpu = CSV.read(input_file_gpu, DataFrame)
n_inputs = 50_000

for process in processes
    df_filt = filter(:process => x -> x == process, df)
    df_filt_gpu = filter(:process => x -> x == process, df_gpu)

    # add dummy factors to get the numbers in similar orders of magnitude
    df_filt.cumulative_optimization_time = df_filt.cumulative_optimization_time .* 1e4
    df_filt_gpu.cumulative_optimization_time = df_filt_gpu.cumulative_optimization_time .* 1e4
    df_filt_gpu.gpu_t = df_filt_gpu.gpu_t .* 1e1

    cpu = !isempty(df_filt)
    gpu = !isempty(df_filt_gpu)

    ymax = 0.0

    if cpu
        @df df_filt scatter(
            :operations,
            :cumulative_optimization_time,
            label = "Cumulative Optimization Time (x10000) (s)",
            markersize = 4,
        )
        ymax = max(df_filt[!, :cpu_st_t]..., df_filt[!, :cumulative_optimization_time]...) * 1.1 * 1e4
        @df df_filt scatter!(
            :operations,
            :cpu_st_t,
            label = "Single-Threaded Execution (s)",
            markersize = 4,
            markershape = :square,
        )
    end

    if gpu
        if !cpu
            @df df_filt_gpu scatter(
                :operations,
                :cumulative_optimization_time,
                label = "Cumulative Optimization Time (x10000) (s)",
                markersize = 4,
            )
            ymax = max(df_filt_gpu[!, :gpu_t]..., df_filt_gpu[!, :cumulative_optimization_time]...) * 1.1 * 1e4
        end

        @df df_filt_gpu scatter!(
            :operations,
            :gpu_t,
            label = "GPU Execution (x10) (s)",
            markersize = 4,
            markershape = :diamond,
        )
    end

    if cpu || gpu
        plot!(
            title = ("$(beautify_title(process)) Reduction Progression ($(n_inputs) Inputs)"),
            xscale = :linear,
            yscale = :linear,
            #ylim = (0, ymax),
            legend = :outerbottom,
            minorgrid = true,
            xticks = :process_size,
            #yticks = [1e-3, 1e-1, 1e1, 1e3],
            xgrid = false,
            xminorticks = false,
            legendcolumns = 1,
            legend_font_pointsize = 12,
            fontsize = 12,
            size = (800, 600),
            xlabel = "optimizer steps (#)",
            ylabel = "time (s)",
        )

        savefig("$(String(process))_reduction_bench.pdf")
    end
end

# ABC vs QED
AB_process = "AB->ABBB"
QED_process = "ke->kkkkke"

df_filt_AB = filter(:process => x -> x == AB_process, df)
df_filt_QED = filter(:process => x -> x == QED_process, df)

max_AB = max(df_filt_AB[!, :operations]...)
max_QED = max(df_filt_QED[!, :operations]...)

df_filt_AB.reduction_progress = df_filt_AB.operations ./ max_AB .* 100.0
df_filt_QED.reduction_progress = df_filt_QED.operations ./ max_QED .* 100.0

df_filt_AB.relative_performance = df_filt_AB.cpu_st_t ./ df_filt_AB[!, :cpu_st_t][1] .* 100.0
df_filt_QED.relative_performance = df_filt_QED.cpu_st_t ./ df_filt_QED[!, :cpu_st_t][1] .* 100.0

@df df_filt_AB scatter(:reduction_progress, :relative_performance, label = AB_process, markersize = 4)
@df df_filt_QED scatter!(:reduction_progress, :relative_performance, label = QED_process, markersize = 4)

plot!(
    title = ("Relative Performance of $(beautify_title(QED_process)) versus $(beautify_title(AB_process)) on CPU"),
    xscale = :linear,
    yscale = :linear,
    #ylim = (0, ymax),
    legend = :outerbottom,
    minorgrid = true,
    xticks = :process_size,
    #yticks = [1e-3, 1e-1, 1e1, 1e3],
    xgrid = false,
    xminorticks = false,
    legendcolumns = 1,
    legend_font_pointsize = 12,
    fontsize = 12,
    size = (800, 600),
    xlabel = "reduction progress (%)",
    ylabel = "relative time taken (%)",
)

savefig("reduction_bench_relative.pdf")


# ABC vs QED on GPU
AB_process = "AB->ABBB"
QED_process = "ke->kkke"

df_filt_AB = filter(:process => x -> x == AB_process, df_gpu)
df_filt_QED = filter(:process => x -> x == QED_process, df_gpu)

max_AB = max(df_filt_AB[!, :operations]...)
max_QED = max(df_filt_QED[!, :operations]...)

df_filt_AB.reduction_progress = df_filt_AB.operations ./ max_AB .* 100.0
df_filt_QED.reduction_progress = df_filt_QED.operations ./ max_QED .* 100.0

df_filt_AB.relative_performance = df_filt_AB.gpu_t ./ df_filt_AB[!, :gpu_t][1] .* 100.0
df_filt_QED.relative_performance = df_filt_QED.gpu_t ./ df_filt_QED[!, :gpu_t][1] .* 100.0

@df df_filt_AB scatter(:reduction_progress, :relative_performance, label = AB_process, markersize = 4)
@df df_filt_QED scatter!(:reduction_progress, :relative_performance, label = QED_process, markersize = 4)

plot!(
    title = (
        "Relative Performance of $(beautify_title(QED_process)) versus $(beautify_title(AB_process)) on GPU (A100)"
    ),
    xscale = :linear,
    yscale = :linear,
    #ylim = (0, ymax),
    legend = :outerbottom,
    minorgrid = true,
    xticks = :process_size,
    #yticks = [1e-3, 1e-1, 1e1, 1e3],
    xgrid = false,
    xminorticks = false,
    legendcolumns = 1,
    legend_font_pointsize = 12,
    fontsize = 12,
    size = (800, 600),
    xlabel = "reduction progress (%)",
    ylabel = "relative time taken (%)",
)

savefig("reduction_bench_relative_gpu.pdf")
