#/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]*)$"

# 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):
    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:
                datapoints.append(m["timestamp"] - started_at)
                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_bar_chart(dataset, title, out_file):
    # calculate stats for each data series
    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)

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

def main():

    dataset = {}

    # dataset["performance"] = build_dataseries(measurements, "aos_performance:start", "aos_performance:done")

    # URPC
    # title = "URPC Micro Benchmark"
    # dataset["client_to_server"] = build_dataseries(measurements, "aos_urpc_nop:start", "aos_urpc_server:start")
    # dataset["server_schedule_task"] = build_dataseries(measurements, "aos_urpc_server:start", "aos_urpc_server:triggered_closure")
    # dataset["server_completed_task"] = build_dataseries(measurements, "aos_urpc_server:triggered_closure", "aos_urpc_server:done")
    # dataset["server_to_client"] = build_dataseries(measurements, "aos_urpc_server:done", "aos_urpc_nop:done")

    # BLOCK DRIVER
    # measurements = get_measurements_from_file("block_driver/no_optimization.log")
    # dataset = {}
    # # we skip the first datapoint here because it is reliably very high compared to all others
    # dataset["read_device"] = build_dataseries(measurements, "block_driver_read_object:start", "block_driver_read_object:memcpy")[1:]
    # dataset["read_memcpy"] = build_dataseries(measurements, "block_driver_read_object:memcpy", "block_driver_read_object:done")
    # dataset["write_read"] = build_dataseries(measurements, "block_driver_write_object:start", "block_driver_write_object:read")
    # dataset["write_memcpy"] = build_dataseries(measurements, "block_driver_write_object:read", "block_driver_write_object:memcpy")
    # dataset["write_device"] = build_dataseries(measurements, "block_driver_write_object:memcpy", "block_driver_write_object:done")
    # create_bar_chart(dataset, "Block Driver Micro Benchmark", "block_driver/no_optimization.jpg")
    # measurements = get_measurements_from_file("block_driver/no_sleep.log")
    # dataset_compare = {}
    # dataset_compare["read_unopt"] = dataset["read_device"]
    # dataset_compare["write_unopt"] = dataset["write_device"]
    # dataset_compare["read_nosleep"] = build_dataseries(measurements, "block_driver_read_object:start", "block_driver_read_object:memcpy")[1:]
    # dataset_compare["write_nosleep"] = build_dataseries(measurements, "block_driver_write_object:memcpy", "block_driver_write_object:done")
    # create_bar_chart(dataset_compare, "Block Driver Optimizations", "block_driver/optimizations.jpg")

    # create plots for data series
    # for key in dataset:
    #     d = dataset[key]
    #     p.clf()
    #     p.title(key)
    #     p.xlabel("datapoint index")
    #     p.ylabel("duration (cycles)")
    #     p.plot(range(len(d)), d)
    #     p.savefig(os.path.join(out_dir, f"{key}.jpg"))
    

if __name__ == "__main__":
    main()