#/bin/env python3
from distutils.log import error
from statistics import mean, median
import sys
from matplotlib import pyplot as p
from mailbox import linesep
import os
import re

from numpy import std

MEASUREMENT_TAG = "MEASUREMENT"
# searches for a measurement with #1 name, #2 tag, #3 time
MEASUREMENT_REGEX = r"^.* " + MEASUREMENT_TAG + r" ([a-zA-Z_:-]*) ([0-9]*)$"
FREQUENCY = 8000000

# TODO change this for the actual measurement
# data_basedir = os.environ.get("BFBUILD_QEMU")
# data_basedir = os.environ.get("BFBUILD")
data_basedir = "./plots"
# data_file = os.path.join(data_basedir, "full_output.log")
out_dir = os.path.join(os.path.dirname(sys.argv[0]), "plots")

def extract_measurement(line):
    match = re.match(MEASUREMENT_REGEX, line)
    if not match:
        print(f"[ERROR] Malformed measurement: '{line}'")
        exit(1)
    return {
        "tag": match.group(1),
        "timestamp": int(match.group(2)),
    }

def extract_measurements(lines):
    filtered = filter(lambda l: MEASUREMENT_TAG in l, lines)
    return list(map(extract_measurement, filtered))

def read_data(file):
    with open(file) as f:
        return f.readlines()
        
def build_dataseries(measurements, start_tag, end_tag, divider=(FREQUENCY / 1000)):
    datapoints = []
    measurements = list(filter(lambda m: start_tag == m["tag"] or end_tag == m["tag"], measurements))

    started_at = None
    for m in measurements:
        if started_at is None:
            if m["tag"] == start_tag:
                started_at = m["timestamp"]
            elif m["tag"] == end_tag:
                print(f"[ERROR] Got end tag without start {end_tag}")
        else:
            if m["tag"] == end_tag:
                duration = (m["timestamp"] - started_at) / divider
                datapoints.append(duration)
                started_at = None
            elif m["tag"] == start_tag:
                print(f"[ERROR] Invalid sequence, ignoring last start tag {start_tag}")
                started_at = m["timestamp"]
    
    return datapoints

def get_measurements_from_file(file):
    raw_log = read_data(os.path.join(data_basedir, file))

    measurements = extract_measurements(raw_log)
    measurements.sort(key=lambda m: m["timestamp"])

    return measurements

def create_metrics(dataset, unit="ms"):
    metrics = {}
    for key in dataset:
        d_mean = mean(dataset[key])
        d_std = std(dataset[key])
        d_median = median(dataset[key])
        metrics[key] = (d_mean, d_std)
        print(f"{key} - mean: {d_mean} {unit}, std: {d_std}")

    return metrics
    
def create_timeline(dataset, title, file, unit="cycles", loc="center right"):
    p.clf()
    p.title(title)
    p.xlabel("i-th measurement")
    p.ylabel(f"duration ({unit})")
    for key in dataset:
        p.plot(range(len(dataset[key])), dataset[key], label=key)
    
    p.legend(loc=loc)
    p.savefig(os.path.join(out_dir, file))

def create_bar_chart(metrics, title, out_file, figsize=(8,4), unit="ms"):
    # calculate stats for each data series

    # create the plot
    p.clf()
    fig, ax = p.subplots(figsize=figsize)
    p.title(title)
    p.ylabel(f"median duration ({unit})")
    p.bar(
        metrics.keys(), 
        [x[0] for x in metrics.values()], 
        yerr=[x[1] for x in metrics.values()],
        capsize=5
    )
    ax.set_ylim(0)
    p.savefig(os.path.join(out_dir, out_file))

def create_multi_bar_chart(metrics_groups, title, out_file, figsize=(8,4), width=0.4, unit="ms"):
    # calculate stats for each data series

    labels = list(metrics_groups[list(metrics_groups.keys())[0]].keys())
    x = list(range(len(labels)))

    # create the plot
    p.clf()
    fig, ax = p.subplots(figsize=figsize)
    p.title(title)
    p.ylabel(f"median duration ({unit})")
    for i, key in enumerate(metrics_groups.keys()):
        offset = - width / 2 + 2 * i * (width / len(metrics_groups))
        b = ax.bar(
            list(map(lambda a: a + offset, x)), 
            [x[0] for x in metrics_groups[key].values()],
            width,
            label=key,
            yerr=[x[1] for x in metrics_groups[key].values()],
            capsize=5
        )
        # ax.bar_label(b)

    #draw grouped bar chart
    ax.set_xticks(x, labels)
    ax.legend()

    fig.tight_layout()
    p.savefig(os.path.join(out_dir, out_file))


def main():

    # dataset = {}
    # measurements = get_measurements_from_file("perf/perf.log")
    # dataset["measurements"] = build_dataseries(measurements, "aos_performance:start", "aos_performance:done", divider=1)
    # create_timeline(dataset, "Performance Measurement System Latency", "perf/perf.jpg")

    # URPC
    # dataset = {}
    # yielded_measurements = get_measurements_from_file("urpc/yield.log")
    # dataset["yield_to_server"] = build_dataseries(yielded_measurements, "aos_urpc_nop:start", "aos_urpc_server:start", divider=1)
    # dataset["yield_to_client"] = build_dataseries(yielded_measurements, "aos_urpc_server:done", "aos_urpc_nop:done", divider=1)
    # measurements = get_measurements_from_file("urpc/empty_loop.log")
    # dataset["empty_to_server"] = build_dataseries(measurements, "aos_urpc_nop:start", "aos_urpc_server:start", divider=1)
    # dataset["empty_to_client"] = build_dataseries(measurements, "aos_urpc_server:done", "aos_urpc_nop:done", divider=1)
    # create_bar_chart(create_metrics(dataset, unit="cycles"), "URPC Yield vs Empty Loop", "urpc/performance.jpg", unit="cycles")
    # dataset = {}
    # dataset["client_to_server"] = build_dataseries(yielded_measurements, "aos_urpc_nop:start", "aos_urpc_server:start", divider=1)
    # dataset["server_scheduled"] = build_dataseries(yielded_measurements, "aos_urpc_server:start", "aos_urpc_server:triggered_closure", divider=1)
    # dataset["server_completed"] = build_dataseries(yielded_measurements, "aos_urpc_server:triggered_closure", "aos_urpc_server:done", divider=1)
    # dataset["server_to_client"] = build_dataseries(yielded_measurements, "aos_urpc_server:done", "aos_urpc_nop:done", divider=1)
    # create_bar_chart(create_metrics(dataset, unit="cycles"), "URPC Micro Benchmark", "urpc/micro.jpg", unit="cycles")

    # BLOCK DRIVER
    # measurements = get_measurements_from_file("block_driver/no_optimization.log")
    # # disk latency without optimization
    # dataset = {}
    # # we only take 200 measurements to not overload the plot and to have the same number of read and write measurements
    # dataset["read unoptimized"] = build_dataseries(measurements, "block_driver_read_object:start", "block_driver_read_object:memcpy")[10:210]
    # dataset["write unoptimized"] = build_dataseries(measurements, "block_driver_write_object:memcpy", "block_driver_write_object:done")[10:210]
    # # disk latency with optimization
    # measurements = get_measurements_from_file("block_driver/no_sleep.log")
    # # we only take 200 measurements to not overload the plot and to have the same number of read and write measurements
    # dataset["read optimized"] = build_dataseries(measurements, "block_driver_read_object:start", "block_driver_read_object:memcpy")[10:210]
    # dataset["write optimized"] = build_dataseries(measurements, "block_driver_write_object:memcpy", "block_driver_write_object:done")[10:210]
    # create_metrics(dataset)
    # create_timeline(dataset, "Block Driver Disk Latency", "block_driver/disk_latency.jpg", unit="ms")

    # FILESYSTEM
    # throughput
    # metrics = {}
    # measurements = get_measurements_from_file("file_system/makro_no_optimization.log")
    # dataset_unopt = {}
    # dataset_unopt["read"] = list(map(lambda d: 10 * 4096 / d, build_dataseries(measurements, "measure_read_throughput:pre", "measure_read_throughput:post", divider=FREQUENCY)))[1:]
    # dataset_unopt["write"] = list(map(lambda d: 10 * 4096 / d, build_dataseries(measurements, "measure_write_throughput:pre", "measure_write_throughput:post", divider=FREQUENCY)))[1:]
    # metrics["unoptimized"] = create_metrics(dataset_unopt)
    # measurements = get_measurements_from_file("file_system/makro_multi_cache.log")
    # dataset_opt = {}
    # dataset_opt["read"] = list(map(lambda d: 10 * 4096 / d, build_dataseries(measurements, "measure_read_throughput:pre", "measure_read_throughput:post", divider=FREQUENCY)))[1:]
    # dataset_opt["write"] = list(map(lambda d: 10 * 4096 / d, build_dataseries(measurements, "measure_write_throughput:pre", "measure_write_throughput:post", divider=FREQUENCY)))[1:]
    # metrics["optimized"] = create_metrics(dataset_opt)
    # create_multi_bar_chart(metrics, "Filesystem Throughput", "file_system/throughput.jpg", unit="bytes/s")
    
    # latencies of other functions
    # metrics = {}
    # dataset_noopt = {}
    # measurements = get_measurements_from_file("file_system/makro_no_optimization.log")
    # dataset_noopt["fopen"] = build_dataseries(measurements, "measure_fopen_fclose:pre_open", "measure_fopen_fclose:post_open")
    # dataset_noopt["fclose"] = build_dataseries(measurements, "measure_fopen_fclose:pre_close", "measure_fopen_fclose:post_close")
    # dataset_noopt["fcreate"] = build_dataseries(measurements, "measure_fcreate_and_rm_single:pre_create", "measure_fcreate_and_rm_single:post_create")
    # dataset_noopt["rm"] = build_dataseries(measurements, "measure_fcreate_and_rm_single:pre_rm", "measure_fcreate_and_rm_single:post_rm")
    # dataset_noopt["mkdir"] = build_dataseries(measurements, "measure_mkdir_and_rmdir_single:pre_mkdir", "measure_mkdir_and_rmdir_single:post_mkdir")
    # dataset_noopt["rmdir"] = build_dataseries(measurements, "measure_mkdir_and_rmdir_single:pre_rmdir", "measure_mkdir_and_rmdir_single:post_rmdir")
    # metrics["unoptimized"] = create_metrics(dataset_noopt)
    # dataset_opt = {}
    # measurements = get_measurements_from_file("file_system/makro_multi_cache.log")
    # dataset_opt["fopen"] = build_dataseries(measurements, "measure_fopen_fclose:pre_open", "measure_fopen_fclose:post_open")
    # dataset_opt["fclose"] = build_dataseries(measurements, "measure_fopen_fclose:pre_close", "measure_fopen_fclose:post_close")
    # dataset_opt["fcreate"] = build_dataseries(measurements, "measure_fcreate_and_rm_single:pre_create", "measure_fcreate_and_rm_single:post_create")
    # dataset_opt["rm"] = build_dataseries(measurements, "measure_fcreate_and_rm_single:pre_rm", "measure_fcreate_and_rm_single:post_rm")
    # dataset_opt["mkdir"] = build_dataseries(measurements, "measure_mkdir_and_rmdir_single:pre_mkdir", "measure_mkdir_and_rmdir_single:post_mkdir")
    # dataset_opt["rmdir"] = build_dataseries(measurements, "measure_mkdir_and_rmdir_single:pre_rmdir", "measure_mkdir_and_rmdir_single:post_rmdir")
    # metrics["optimized"] = create_metrics(dataset_opt)
    # create_multi_bar_chart(metrics, "Filesystem Operation Latency", "file_system/latency.jpg")

    # PAGING
    # makro
    # dataset = {}
    # measurements = get_measurements_from_file("paging/no_optimization_core_1.log")
    # dataset["core 1"] = build_dataseries(measurements, "pt_exception:start", "pt_exception:unlocked")
    # measurements = get_measurements_from_file("paging/no_optimization_core_0.log")
    # dataset["core 0"] = build_dataseries(measurements, "pt_exception:start", "pt_exception:unlocked")
    # create_timeline(dataset, "Pagefault handling latency", "paging/paging_makro_latency.jpg", unit="ms")
    # # micro
    # dataset = {}
    # metrics = {}
    # measurements = get_measurements_from_file("paging/no_optimization_core_0.log")
    # dataset["locking"] = build_dataseries(measurements, "pt_exception:start", "pt_exception:locked")
    # dataset["resolve vaddr"] = build_dataseries(measurements, "pt_exception:locked", "pt_exception:vaddr")
    # dataset["request memory"] = build_dataseries(measurements, "pt_exception:vaddr", "pt_exception:got_frame")
    # dataset["mapping"] = build_dataseries(measurements, "pt_exception:got_frame", "pt_exception:mapped")
    # dataset["unlocking"] = build_dataseries(measurements, "pt_exception:mapped", "pt_exception:unlocked")
    # metrics["core 0"] = create_metrics(dataset)
    # dataset = {}
    # measurements = get_measurements_from_file("paging/no_optimization_core_1.log")
    # dataset["locking"] = build_dataseries(measurements, "pt_exception:start", "pt_exception:locked")
    # dataset["resolve vaddr"] = build_dataseries(measurements, "pt_exception:locked", "pt_exception:vaddr")
    # dataset["request memory"] = build_dataseries(measurements, "pt_exception:vaddr", "pt_exception:got_frame")
    # dataset["mapping"] = build_dataseries(measurements, "pt_exception:got_frame", "pt_exception:mapped")
    # dataset["unlocking"] = build_dataseries(measurements, "pt_exception:mapped", "pt_exception:unlocked")
    # metrics["core 1"] = create_metrics(dataset)
    # create_multi_bar_chart(metrics, "Pagefault handling micro benchmark", "paging/paging_micro_latency.jpg")

    # UMP
    # dataset = {}
    # measurements = get_measurements_from_file("ump/different_cores.log")
    # dataset["different cores"] = build_dataseries(measurements, "block_driver_client:dispatching", "block_driver_client:payload", 1)[10:]
    # measurements = get_measurements_from_file("ump/same_core.log")
    # dataset["same core"] = build_dataseries(measurements, "block_driver_client:dispatching", "block_driver_client:payload", 1)[10:]
    # create_timeline(dataset, "UMP roundtrip latency", "ump/latency.jpg")

    pass

if __name__ == "__main__":
    main()