using CSV
using DataFrames
using Plots
using StatsPlots
using LaTeXStrings

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

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

input_file = ARGS[1]
df = CSV.read(input_file, DataFrame)

# plotting with process size as x axis
THREADS = [1]

for threads in THREADS
    title_string = "n-photon Compton diagram generation"

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

    # ns -> s
    df_filt.graph_gen_mean = @. df_filt.graph_gen_mean / 1e9
    df_filt.graph_gen_std = @. df_filt.graph_gen_std / 1e9

    # B -> MB (not MiB since the log scale is base 10)
    df_filt.graph_mem = @. df_filt.graph_mem / 1e6
    df_filt.graph_mem_reduced = @. df_filt.graph_mem_reduced / 1e6

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

    println(df_filt[!, :process_size])
    println(df_filt[!, :graph_mem])
    println(df_filt[!, :graph_mem_reduced])


    @df df_filt scatter(:process_size, :graph_mem, label = "unreduced graph", markersize = 7)
    scatter!(
        df_filt[!, :process_size],
        df_filt[!, :graph_mem_reduced],
        label = "reduced graph",
        markershape = :square,
        markersize = 7,
    )

    plot!(
        title = "n-photon Compton diagram memory footprint",
        yscale = :log10,
        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),
        ylabel = "memory footprint (MB)",
        xlabel = "process size (#)",
    )

    savefig("gen_memory_$(threads).pdf")

    @df df_filt scatter(
        :process_size,
        :graph_gen_mean,
        yerror = :graph_gen_std,
        label = "graph generation time",
        markersize = 7,
    )
    scatter!(
        df_filt[!, :process_size],
        df_filt[!, :graph_elapsed_reduce],
        label = "graph reduction time",
        markershape = :square,
        markersize = 7,
    )

    plot!(
        title = title_string,
        yscale = :log10,
        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),
        ylabel = "time (s)",
        xlabel = "process size (#)",
    )

    savefig("gen_times_$(threads)_threads.pdf")
    exit(0)

    # graph size
    title_string = "n-photon Compton unreduced graph size"

    @df df_filt scatter(:process_size, :graph_nodes, label = "nodes", markershape = :circle)
    @df df_filt scatter!(:process_size, :graph_edges, label = "edges", markershape = :square)
    @df df_filt scatter!(:process_size, :graph_u_nodes, label = "U-nodes", markershape = :star)
    @df df_filt scatter!(:process_size, :graph_v_nodes, label = "V-nodes", markershape = :utriangle)
    @df df_filt scatter!(:process_size[2:end], :graph_s1_nodes[2:end], label = "S1-nodes", markershape = :x)
    @df df_filt scatter!(:process_size, :graph_s2_nodes, label = "S2-nodes", markershape = :diamond)

    plot!(
        title = title_string,
        yscale = :log10,
        legend = :outerbottom,
        yminorgrid = true,
        xticks = :process_size,
        yticks = [1e1, 1e3, 1e5, 1e7],
        xgrid = false,
        xminorticks = false,
        legendcolumns = 2,
        legend_font_pointsize = 10,
        size = (800, 600),
        ylabel = "(#)",
        xlabel = "process size (#)",
    )

    savefig("compton_graph_size_unreduced.pdf")


    # graph size
    title_string = "n-photon Compton reduced graph size"

    @df df_filt scatter(:process_size, :graph_nodes_reduced, label = "nodes", markershape = :circle)
    @df df_filt scatter!(:process_size, :graph_edges_reduced, label = "edges", markershape = :square)
    @df df_filt scatter!(:process_size, :graph_u_nodes_reduced, label = "U-nodes", markershape = :star)
    @df df_filt scatter!(:process_size, :graph_v_nodes_reduced, label = "V-nodes", markershape = :utriangle)
    @df df_filt scatter!(:process_size[2:end], :graph_s1_nodes_reduced[2:end], label = "S1-nodes", markershape = :x)
    @df df_filt scatter!(:process_size, :graph_s2_nodes_reduced, label = "S2-nodes", markershape = :diamond)

    plot!(
        title = title_string,
        yscale = :log10,
        legend = :outerbottom,
        yminorgrid = true,
        xticks = :process_size,
        yticks = [1e1, 1e2, 1e3, 1e4, 1e5, 1e6],
        xgrid = false,
        xminorticks = false,
        legendcolumns = 2,
        legend_font_pointsize = 10,
        size = (800, 600),
        ylabel = "(#)",
        xlabel = "process size (#)",
    )

    savefig("compton_graph_size_reduced.pdf")


    # graph size versus
    title_string = "n-photon Compton graph sizes"

    @df df_filt scatter(:process_size, :graph_nodes, label = "nodes", markershape = :circle)
    @df df_filt scatter!(:process_size, :graph_edges, label = "edges", markershape = :square)
    @df df_filt scatter!(:process_size, :graph_nodes_reduced, label = "nodes (after reduction)", markershape = :star)
    @df df_filt scatter!(
        :process_size,
        :graph_edges_reduced,
        label = "edges (after reduction)",
        markershape = :utriangle,
    )

    plot!(
        title = title_string,
        yscale = :log10,
        legend = :outerbottom,
        yminorgrid = true,
        xticks = :process_size,
        yticks = [1e1, 1e2, 1e3, 1e4, 1e5, 1e6],
        xgrid = false,
        xminorticks = false,
        legendcolumns = 2,
        legend_font_pointsize = 10,
        size = (800, 600),
        ylabel = "(#)",
        xlabel = "process size (#)",
    )

    savefig("compton_graph_size_versus.pdf")

end

# for a specific process, plot times with threads as x
process = "ke->kkkkkkkke"
title_string = "n-photon Compton diagram generation times, $process"

df_filt = filter(:process_name => x -> x == process, df)
df_filt.graph_gen_mean = @. df_filt.graph_gen_mean / 1e9
df_filt.graph_gen_std = @. df_filt.graph_gen_std / 1e9

@df df_filt scatter(
    :cpu_threads,
    :graph_gen_mean,
    yerror = :graph_gen_std,
    label = "graph generation time",
    markersize = 7,
)

plot!(
    title = title_string,
    yscale = :linear,
    legend = :outerbottom,
    minorgrid = true,
    xticks = :cpu_threads,
    #yticks = [1e-3, 1e-2, 1e-1, 1e-0, 1e1],
    ylim = (0, max(df_filt[!, :graph_gen_mean]...) * 1.1),
    xgrid = false,
    xminorticks = false,
    legendcolumns = 1,
    legend_font_pointsize = 10,
    size = (800, 600),
    ylabel = "time (s)",
    xlabel = "process size (#)",
)

savefig("gen_times_$(process).pdf")