#!/usr/bin/env python3

from sys import argv, exit
import os
from csv import writer
from typing import Tuple, List
import re
import numpy as np
from pprint import pprint

import matplotlib.pyplot as plt

if len(argv) != 2:
    print(f"Usage: { argv[0] } <results_dir>")
    exit(1)

path = argv[1]

files = [f for f in os.listdir(path) if os.path.isfile(f"{ path }/{f}")]

# files = [f for f in files if re.match(r"greedy_64_.+", f) is not None ]


def parse_file(file: str) -> Tuple[float, List[int]]:

    pops = list()
    time = None
    with open(file) as f:
        for line in f.readlines():
            m = re.match(r"popped\s(?P<pops>\d+)\s.*", line)
            if m is not None:
                pops.append(int(m.groupdict()["pops"]))
                continue

            m = re.match(r"It took\s(?P<time>\S+).*", line)
            if m is not None:
                time = float(m.groupdict()["time"])

    if len(pops) == 0:
        raise Exception(f"Parsing { file } failed, no heap pops found!")

    return time, pops

results = dict()

with open("times.csv", "w+") as times_file:
    times = writer(times_file)

    with open("pops.csv", "w+") as pops_file:
        pops = writer(pops_file)

        for file in files:
            name = file.split(".")[0]
            full_path = f"{ path }/{ file }"
            time, pop = parse_file(full_path)

            total_pops = sum(pop)/len(pop)

            results[name] = (total_pops, time)

            times.writerow([name, time])

            pops.writerow([name, *pop])
rel_pops = list()
abs_pops = list()
rel_time = list()
abs_time = list()
labels = list()



pprint(results)
baseline = "dijkstra"
base_pops = results[baseline][0]
base_time = results[baseline][1]

for name, values in results.items():
    pops, time = values
    labels.append(name)
    rel_pops.append(pops/base_pops)
    rel_time.append(time/base_time)
    abs_pops.append(pops)
    abs_time.append(time)



x = np.arange(len(labels))  # the label locations
width = 0.35  # the width of the bars

fig, ax = plt.subplots()
rects1 = ax.bar(x - width/2, abs_time , width, label='time')
# rects2 = ax.bar(x + width/2, abs_pops, width, label='pops')

ax.legend()
ax.set_xticks(x, labels)

plt.show()