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'",
    #"QED Process: 'ke->kkkkkke'",
    #"QED Process: 'ke->kkkkkkke'",
    "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 abc_proc_to_n(str::AbstractString)
    parts = split(str, "'")
    parts = split(parts[2], "->")
    b_count = count(c -> c == 'B', parts[2])
    return b_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]



title_string = "QED N-Photon Compton Scattering\nCalculate 10,000,000 Matrix Elements"

df_filt = filter(:process_name => x -> proc_to_n(x) >= 1, df)

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

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

@df df_red scatter(
    :process_size,
    :cpu_time,
    yerror = :cpu_std,
    label = "CPU execution time, 32 threads (s)",
    markersize = 6,
)
@df df_red scatter!(
    :process_size,
    :gpu_time,
    yerror = :gpu_std,
    label = "GPU execution time, A100 80GB (s)",
    markersize = 6,
)

plot!(
    title = title_string,
    yscale = :log10,
    legend = :outerbottom,
    legendcolumns = 2,
    legend_font_pointsize = 10,
    size = (800, 600),
    ylabel = "time (s)",
    xlabel = "process size (#)",
)

savefig("cpu_vs_gpu_qed.pdf")




title_string = "\$AB\\rightarrow AB^n\$ ABC Processes\nCalculate 10,000,000 Matrix Elements"

df_filt = filter(:process_name => x -> abc_proc_to_n(x) >= 1, df)

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

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

@df df_red scatter(
    :process_size,
    :cpu_time,
    yerror = :cpu_std,
    label = "CPU execution time, 32 threads (s)",
    markersize = 6,
)
@df df_red scatter!(
    :process_size,
    :gpu_time,
    yerror = :gpu_std,
    label = "GPU execution time, A100 80GB (s)",
    markersize = 6,
)

plot!(
    title = title_string,
    yscale = :log10,
    legend = :outerbottom,
    legendcolumns = 2,
    legend_font_pointsize = 10,
    size = (800, 600),
    ylabel = "time (s)",
    xlabel = "process size (#)",
)

savefig("cpu_vs_gpu_abc.pdf")
