import math
import statistics
from collections import defaultdict

import matplotlib.pyplot as plt


def load_results(filename) -> dict:
    results = defaultdict(list)
    with open(filename, "r") as file:
        for line in file:
            if "error" in line:
                continue
            values = line.split()
            keys = [
                ("num_servers", int),
                ("database_size", int),
                ("block_size", int),
                ("protocol_name", str),
                ("total_cpu_time", int),
                ("bytes_sent", int),
                ("bytes_received", int)
            ]
            d = dict([(a[0], a[1](b)) for a, b in zip(keys, values)])
            results[(d["num_servers"], d["database_size"], d["block_size"], d["protocol_name"])].append(d)
    return results


def clean_results(results) -> dict:
    cleaned_results = defaultdict(list)
    for test, result in results.items():
        cpu_time = statistics.mean(sorted([int(r["total_cpu_time"]) for r in result]))  # [1:-1]
        cleaned_results[result[0]["protocol_name"]].append({**result[0], "total_cpu_time": cpu_time})
    return cleaned_results


def plot(all_results: dict, y_func: callable, x_func: callable, title=None, y_label=None, x_label=None):
    fig, ax = plt.subplots()
    for protocol_name, results in all_results.items():
        sorted_results = sorted(results, key=lambda r: x_func(r))
        ax.plot(
            [x_func(r) for r in sorted_results],
            [y_func(r) for r in sorted_results],
        )

    ax.set_xscale("log", basex=2)
    #ax.set_yscale("log", basey=10)

    if x_label is not None:
        plt.xlabel(x_label)
    if y_label is not None:
        plt.ylabel(y_label)

    plt.legend(all_results.keys(), loc="upper left")
    if title is not None:
        plt.title(title)

    plt.savefig(title.replace(" ", "_").replace("/", "") + ".pdf")
    #plt.show()


if __name__ == '__main__':
    plot(
        clean_results(load_results("results_only_with_8_bs.log")),
        y_label="Time",
        x_label="Total Database Size",
        title="Fixed 8-bit block size",
        y_func=lambda r: max(1, r["total_cpu_time"]),
        x_func=lambda r: r["database_size"] * r["block_size"]
    )

    plot(
        clean_results(load_results("results_only_with_8_bs.log")),
        y_label="Time",
        x_label="Total Database Size",
        title="Fixed 8-bit block size - simulated 1MiB/s tx, 5MiB/s rx",
        y_func=lambda r: r["total_cpu_time"] + (r["bytes_sent"]/(1*1024)) + (r["bytes_received"]/(5*1024)),  # 1024 B/ms = 1MiB/s
        x_func=lambda r: r["database_size"] * r["block_size"]
    )

    #plot(
    #    {n: [r for r in rs if r["block_size"] == 2] for n, rs in clean_results(load_results("results_overflow.log")).items()},
    #    y_label="Time",
    #    x_label="Total Database Size",
    #    title="lol",
    #    y_func=lambda r: r["total_cpu_time"] + (r["bytes_sent"]/(1*1024)) + (r["bytes_received"]/(10*1024)),  # 1024 B/ms = 1MiB/s
    #    x_func=lambda r: r["database_size"] * r["block_size"]
    #)