#!/usr/bin/env python3

from sys import argv, exit
import os
from typing import Tuple, List
import re

import numpy as np
import pandas as pd
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}")]


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


frames = list()

for i, file in enumerate(files):
    time, pops = parse_file(f"{ path }/{ file }")

    parts = file.split(".")[0].split("_")

    data = {
            "name": parts[0],
            "time": time,
            "pops": int(sum(pops)),
            }

    if len(parts) > 1:
        data["num_landmarks"] = int(parts[1])
    if len(parts) > 2:
        data["best_size"] = int(parts[2])

    frames.append(pd.DataFrame(data, index=[i]))

results = pd.concat(frames)


# dijkstra vs A* vs greedy and random ALT

#condition = ((results["name"] == "astar")
#        | (results["name"] == "dijkstra")
#        | ((results["num_landmarks"] == 32) & (results["best_size"] == 4)))
#
# basic = results[condition].sort_values(by=["time"], ascending=False)
# names = basic["name"].to_list()
# for i in range(len(names)):
#     if names[i] == "greedy":
#         names[i] = "ALT 32 greedy"
#     if names[i] == "random":
#         names[i] = "ALT 32 random"
# times = basic["time"] * 1000
# 
# fig, ax = plt.subplots()
# 
# bar = ax.bar(names, times)
# ax.set_title("Average query time by algorithm")
# ax.set_ylabel("time in ms")
# ax.bar_label(bar)
# fig.savefig("basic.svg")
# plt.show()


# comparison greedy vs. random, best_size = 4

condition = (((results["name"] == "greedy") | (results["name"] == "random")) & (results["best_size"] == 4))

data = results[condition]
greedy = data[data["name"] == "greedy"]
random = data[data["name"] == "random"]

print(greedy)
print(random)

labels = greedy["num_landmarks"].to_list()
labels = [int(label) for label in labels]
greedy_times = greedy["time"].to_list()
random_times = random["time"].to_list()

width = 0.35

fig, ax = plt.subplots()

x = np.arange(len(labels))

greedy_bars = ax.bar(x - width/2, greedy_times, width, label="Greedy")
random_bars = ax.bar(x + width/2, random_times, width, label="Random")
ax.set_title("Comparison of ALT with random and greedy landmark selection")
ax.set_ylabel("time in ms")
ax.set_xlabel("number of landmarks")
ax.set_xticks(x, labels)
ax.legend()
# ax.bar_label(greedy_bars)
# ax.bar_label(random_bars)

fig.save("random_greedy_comp.svg")

plt.show()


# comparison greedy best_size
# comparison random best_size