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 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]

# plotting with threads as x axis
for process_name in processes
    title_string = "$(beautify_title(process_name)), $n_inputs samples"
    println("$title_string")

    process_no_opt = process_name * " not optimized"
    process_red = process_name * " reduced"
    process_tape_no_opt = process_name * " not optimized tape"
    process_tape_red = process_name * " reduced tape"

    df_no_opt = filter(:process_name => x -> x == process_no_opt, df)
    df_red = filter(:process_name => x -> x == process_red, df)
    df_tape_no_opt = filter(:process_name => x -> x == process_tape_no_opt, df)
    df_tape_red = filter(:process_name => x -> x == process_tape_red, df)

    @df df_no_opt scatter(:cpu_threads, :cpu_time, label = "unoptimized function", markershape = :circle)
    @df df_red scatter!(:cpu_threads, :cpu_time, label = "reduced function", markershape = :rect)
    @df df_tape_no_opt scatter!(:cpu_threads, :cpu_time, label = "unoptimized tape", markershape = :utriangle)
    @df df_tape_red scatter!(:cpu_threads, :cpu_time, label = "reduced tape", markershape = :star)

    plot!(
        title = title_string,
        yscale = :linear,
        legend = :outerright,
        legendcolumns = 1,
        legend_font_pointsize = 10,
        size = (800, 600),
        ylabel = "execution time (s)",
        xlabel = "threads (#)",
    )

    savefig("$(process_name)_time.pdf")


    @df df_no_opt scatter(:cpu_threads, :cpu_rate, label = "unoptimized function", markershape = :circle)
    @df df_red scatter!(:cpu_threads, :cpu_rate, label = "reduced function", markershape = :rect)
    @df df_tape_no_opt scatter!(:cpu_threads, :cpu_rate, label = "unoptimized tape", markershape = :utriangle)
    @df df_tape_red scatter!(:cpu_threads, :cpu_rate, label = "reduced tape", markershape = :star)

    plot!(
        title = "Sample rate, " * title_string,
        yscale = :log10,
        legend = :outerright,
        legendcolumns = 1,
        legend_font_pointsize = 10,
        size = (800, 600),
        ylabel = "rate (" * L"s^{-1}" * ")",
        xlabel = "threads (#)",
    )

    savefig("$(process_name)_rate.pdf")


    @df df_no_opt scatter(:cpu_threads, :cpu_gflops, label = "unoptimized function", markershape = :circle)
    @df df_red scatter!(:cpu_threads, :cpu_gflops, label = "reduced function", markershape = :rect)
    @df df_tape_no_opt scatter!(:cpu_threads, :cpu_gflops, label = "unoptimized tape", markershape = :utriangle)
    @df df_tape_red scatter!(:cpu_threads, :cpu_gflops, label = "reduced tape", markershape = :star)

    plot!(
        title = "CPU performance, " * title_string,
        yscale = :linear,
        legend = :outerright,
        legendcolumns = 1,
        legend_font_pointsize = 10,
        size = (800, 600),
        ylabel = "CPU performance (GFLOPS)",
        xlabel = "threads (#)",
    )

    savefig("$(process_name)_performance.pdf")

end

# plotting with process size as x axis
THREADS = [1, 2, 4, 8, 16]

for threads in THREADS
    title_string = "$threads threads, $n_inputs samples"

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

    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_tape_no_opt = filter(:process_name => x -> match(r" not optimized tape$", x) !== nothing, df_filt)
    df_tape_red = filter(:process_name => x -> match(r" reduced tape$", x) !== nothing, df_filt)


    @df df_no_opt scatter(:process_size, :graph_gen_time, label = "graph generation time")
    @df df_red scatter!(:process_size, :optimization_time, label = "optimization time")
    @df df_no_opt scatter!(:process_size, :function_generation_time, label = "unoptimized function generation time")
    @df df_tape_no_opt scatter!(:process_size, :function_generation_time, label = "unoptimized tape generation time")
    @df df_red scatter!(:process_size, :function_generation_time, label = "reduced function generation time")
    @df df_tape_red scatter!(:process_size, :function_generation_time, label = "reduced tape generation time")

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

    savefig("gen_times_$(threads)_threads.pdf")
end