import subprocess
import tempfile
import re
import matplotlib.pyplot as plt
from collections import defaultdict
import argparse
from math import log2
import pandas as pd



parser = argparse.ArgumentParser(description='Read data from a data file and plot this.')

parser.add_argument('--data_file', action='store', dest='data_file', default='plot_data.txt',
                    help='The desired data to be plotted.')
parser.add_argument('--static_n', action='store_true',
                    help="Is the n variable static in the data file?")
parser.add_argument('--static_step', action='store_true',
                    help="Is the step variable static in the data file?")


args = parser.parse_args()

data = open(args.data_file).read()


# d = {'static variable' : {step : time, step : time}}


def prep_data(static_var, data):
    non_static_var = (static_var + 1) % 2 # if 0 then 1, if 1 then 0
    #non_static_var = 1 - static_var #rip :(

    lines = data.split('\n')
    static_vars = set()

    data = defaultdict(dict)
    
    for line in lines[0:-1]:
        all_vars = re.findall(r"\d+", line)
        
        static = int(all_vars[static_var])
        non_static = int(all_vars[non_static_var])
        time = int(all_vars[2])
        
        static_vars.add(static)
        data[static][non_static] = time

    return data
        
        
    

static_var = 0 if args.static_n else 1
plot_data = prep_data(static_var, data)
df = pd.DataFrame(data=plot_data)
print(df)

plt.imshow(df, cmap='hot')
plt.show()


#first = plot_data[list(plot_data.keys())[2]]

#plt.plot(first[0], first[1])

#for static, plotting_data in plot_data.items():
#    x_axis = [log2(x) for x in plotting_data[0]]
#    plt.plot(x_axis, plotting_data[1],  marker="o", label = static)


#plt.xticks([2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26], ["$log2(2^2)$", "log2(2^4)", "log2(2^6)", "log2(2^8)", "log2(2^10)", "log2(2^12)", "log2(2^14)", "log2(2^16)", "log2(2^18)", "log2(2^20)", "log2(2^22)", "log2(2^24)", "log2(2^26)"])
#plt.legend(loc=1, borderaxespad=0.)
#plt.ylabel('time/ms')
#plt.xlabel('log(step)')
#plt.show()