import math
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
import seaborn as sns


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)

hm = sns.heatmap(df, annot=False, square=False, cmap=sns.cm.rocket_r)
hm.invert_yaxis()

# fix xtics
locs, labels = plt.xticks()
plt.xticks(locs, [f"$2^{{{int(math.log2(int(i.get_text())))}}}$" for i in labels])
locs, labels = plt.yticks()
plt.yticks(locs, [f"$2^{{{int(math.log2(int(i.get_text())))}}}$" for i in labels])

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()