#!/usr/bin/env python3
import matplotlib.pyplot as plt
import numpy as np
import sys

if len(sys.argv) < 2:
    print("Usage:", sys.argv[0], "<data file name>")
    exit()

# Index to name lookup table
id_name_map = ("field", "matrix", "neighborhood", "pointer", "update", "transitive")
# Name to index lookup hash table
name_id_map = {id_name_map[idx]: idx for idx in range(len(id_name_map))}

# Extracts one benchmark's result from a file with multiple
def multibench_file(filename):
    data = np.loadtxt(filename, usecols=(1,2,4,0), converters={0: lambda n: name_id_map[n.decode("utf_8")]})
    names = np.unique(data[:,3])
    for name in names:
        yield (data[data[:,3] == name][:,:-1], id_name_map[int(name)])

# Subsamples one benchmark at a specific WSS
def plot_exec_vs_cache_data(data, ax, wss, **kwargs):
    wsss = np.unique(data[:,0])
    # Get exec time for this working set size with all ways allocated
    divisor = data[np.logical_and(data[:,0] == wss, data[:,1] == 16.0)][0,2]
    # Normalize relative to that exec time
    one_wss_data = data[data[:,0] == wss] / [1,1,divisor]
    # Plot max exec time vs cache allocation
    return ax.plot(np.flip(one_wss_data[:,1],0), np.flip(one_wss_data[:,2],0), **kwargs)

# Wrapper function for backwards compatibility
def plot_exec_vs_cache(filename, ax, wss, line='-'):
    data = np.loadtxt(filename, usecols=(1,2,4))
    plot_exec_vs_cache_data(data, ax, wss, linestyle=line, label=filename.split("-")[-4])

# Styles and colors which match those used in the paper
name_to_label = {"matrix":"-", "transitive":"-.", "neighborhood":"--", "pointer":":", "update":"-", "field":"-."}
name_to_color = {"matrix":"C0", "transitive":"C1", "neighborhood":"C2", "pointer":"C3", "update":"C4", "field":"C5"}
plt.rcParams["font.family"] = "serif"
plt.rcParams['figure.figsize'] = [6, 4]
legend_labels = []
wss = 4 * 2**20

fig, ax = plt.subplots()
ax.set_xlabel("Cache Half Ways Allocated")
ax.set_ylabel("Worst-Case Execution Time Multiplier")
for (data, name) in multibench_file(sys.argv[1]):
    line, = plot_exec_vs_cache_data(data, ax, wss, linestyle=name_to_label[name], label=name, color=name_to_color[name])
    # Save the max plotted value and legend info
    legend_labels.append((max(line.get_ydata()), line, name))

# Sort the legend by the max plotted value
legend_labels.sort(key=lambda k: k[0], reverse=True)
ax.legend([x[1] for x in legend_labels], [x[2] for x in legend_labels])

# Plot from most to least ways allocated
ax.set_xlim(ax.get_xlim()[::-1])

fig.savefig("exec_vs_cache_4mb_ae.pdf")
print("Plot saved as exec_vs_cache_4mb_ae.pdf")