#/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_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 = []

    started_at = None
    for m in measurements:
        if started_at is None:
            if m["tag"] == start_tag:
                started_at = m["timestamp"]
        else:
            if m["tag"] == end_tag:
                datapoints.append(m["timestamp"] - started_at)
                started_at = None
            elif m["tag"] == start_tag:
                print("[ERROR] Invalid sequence, ignoring last start tag")
                started_at = m["timestamp"]
    
    return datapoints


def main():
    raw_log = read_data(data_file)

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

    dataset = {}

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

    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")

    # 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_median, d_std)

    # create output directory
    os.makedirs(out_dir, exist_ok=True)

    # create plots for data series
    for key in dataset:
        d = dataset[key]
        d.sort()
        p.clf()
        p.title(key)
        p.xlabel("datapoint index")
        p.ylabel("duration (cycles)")
        p.scatter(range(len(d)), d)
        p.savefig(os.path.join(out_dir, f"{key}.jpg"))
    
    p.clf()
    fig, ax = p.subplots(figsize=(12,5))
    p.title("Metrics")
    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, "metrics.jpg"))

if __name__ == "__main__":
    main()