from os import path
from argparse import Namespace
import glob
from typing import List, Tuple, Union
import re

import trappy
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt

from arguments import parser

def parse_basename_into_byte_count(basename: str) -> int:
    ret: int = -1

    regex = re.compile(r"(?P<size>[0-9]+)(?P<unit>[kmgKMG])")
    pat = regex.search(basename)
    if pat is not None:
        size = int(pat.group("size"))
        match pat.group("unit"):
            case "K": ret = size
            case "M": ret = size * 1024
            case "G": ret = size * 1024 * 1024
            case _:
                raise Exception("Unexpected unit \"{}\".".format(pat.group("unit")))

    return ret

def clip_to(
        dataset: Union[pd.Series, pd.DataFrame],
        percentile: float,
        axis: int = None
) -> Union[pd.Series, pd.DataFrame]:
    assert(0 <= percentile <= 1)
    if percentile == .0:
        return dataset

    lower_bound = dataset.quantile(1 - percentile)
    upper_bound = dataset.quantile(percentile)

    if isinstance(dataset, pd.DataFrame):
        clipped = dataset.clip(lower=lower_bound, upper=upper_bound, axis=axis)
    else:
        clipped = dataset.clip(lower=lower_bound, upper=upper_bound)

    return clipped

def parse_function_graph_ftrace(trace_fpath: str):
    trace = trappy.FTrace(trace_fpath)
    fngraph_df: pd.DataFrame = trace.funcgraph_exit.data_frame
    assert fngraph_df is not None

    # Sieve `__dcache_clean_poc` entries
    # Notably, the long ones are mostly preempted by e.g., softirq (rcu, etc.)
    # I should prob. identify them but whatever
    # Other long ones that are NOT preempted may be due to QEMU process
    # scheduling on host, not sure
    fngraph_df = fngraph_df.loc[fngraph_df["func"] == "__dcache_clean_poc"]

    # Compute runtime of function
    fngraph_df.loc[:, "rettime"] = (fngraph_df["rettime"]
        .astype(str)
        .apply(lambda x: int(x, base=16)))
    fngraph_df.loc[:, "calltime"] = (fngraph_df["calltime"]
        .astype(str)
        .apply(lambda x: int(x, base=16)))
    fngraph_df = fngraph_df.assign(
        runtime = fngraph_df["rettime"] - fngraph_df["calltime"])

    print("Shape: {}; head: \n{}".format(fngraph_df.shape, fngraph_df.head()))
    return fngraph_df

def main(args: Namespace):
    # Load dataset -- each should have good 1000 entries
    ftrace_dat_paths: List[str] = []
    ftrace_dfs: List[Tuple(int, pd.DataFrame)] = []
    if args.glob:
        ftrace_dat_paths = glob.glob(args.directory)
    else:
        ftrace_dat_paths = glob.glob(args.directory + "/*.dat")

    for ftrace_dat_path in ftrace_dat_paths:
        print("Loading \"{}\"...".format(ftrace_dat_path))
        fngraph_df = parse_function_graph_ftrace(ftrace_dat_path)
        size_in_kb = parse_basename_into_byte_count(
            str(path.basename(ftrace_dat_path).split(".")[:-1]))
        ftrace_dfs.append((size_in_kb, fngraph_df))

    ftrace_dfs.sort(key=lambda tp: tp[0])

    # Now we have filename-df pairs, combine into single df
    series = [
        clip_to(
            pd.to_numeric(df.reset_index()["runtime"]).apply(lambda x: x / 1e3),
            args.percentile
        ) for _, df in ftrace_dfs
    ]
    indices = [size for size, _ in ftrace_dfs]
    plot_df = pd.concat(series, axis=1, keys=indices) # In long-form
    print("Plot shape: {}; head: \n{}"
        .format(plot_df.shape, plot_df.head()))

    # plot violin plot
    fig, ax = plt.subplots(figsize=(12, 8))
    sns.violinplot(
        data=plot_df, orient="h", fill=False, log_scale=args.use_log_scale)
    ax.set_xlabel("Runtime (μs)")
    ax.set_ylabel("Write amount (KiB)")
    ax.set_title("Percentile: {}{}"
        .format(args.percentile, ", log scale" if args.use_log_scale else ""))
    fig.savefig(args.output)
    plt.show()

if __name__ == "__main__":
    args = parser.parse_args()
    sns.set_theme(context="talk", style="ticks", rc={"axes.grid": True})
    main(args)
