"""
Compute statistics from CSV files generated by ddlog. See README.md for elaboration l
"""
import sys
from os import listdir

import pandas as pd
from pandas import DataFrame


def main():
    args = sys.argv
    if len(args) != 4:
        print("Usage: python generate_stats.py path/to/csv_files/ command output_file_name", file=sys.stderr)
        sys.exit(1)

    path = args[1]
    command = args[2]
    output_file = args[3]
    try:
        files = [f for f in listdir(path)]
    except:
        print("Exception: ", sys.exc_info()[0], file=sys.stderr)

    # Contains tuples of (filename, worker_number)
    worker_files = []
    # Files expected to be of the form: name_wN.csv where N is the number of workers represented by this file.
    for f in files:
        # Parse file name to get number. Yuck.
        worker_num = int(f.split('_')[1].split('.')[0][1:])
        worker_files.append((f, worker_num))

    # Sort by worker number.
    worker_files.sort(key=lambda x: x[1])

    # Split into two lists
    [files, workers] = list(zip(*worker_files))

    categories = {"worker_id": "category", "event_type": "category", "operator_name": "category",
                          "operator_addr": "category"}

    commands = ["exe_invocation_table", "exe_time_summary", "progress_table", "data_sizes_table"]
    if command == commands[0]:
        generate_execution_invocation_statistics(files, output_file, path, workers, categories)
    elif command == commands[1]:
        execution_times(files, path, workers, categories)
    elif command == commands[2]:
        progress_events(files, output_file, path, workers, categories)
    elif command == commands[3]:
        data_message_events(files, output_file, path, workers, categories)
    else:
        print("Unknown command \"{}\". Possible commands: {}".format(command, commands), file=sys.stderr)

def progress_events(files, output_file, path, workers, categories):
    push_events = DataFrame()

    for (file, i) in zip(files, workers):
        df = pd.read_csv(path + '/' + file, dtype=categories)

        # Push progress events
        push = df[df["event_type"] == "PushProgress"]
        push_groups = push.groupby("operator_id").size()

        push_events["progress_events_{}w".format(i)] = push_groups

    schedule_events = df[df["event_type"] == "Schedule"]
    op_names = schedule_events.groupby("operator_id")["operator_name"].unique()
    op_address = schedule_events.groupby("operator_id")["operator_addr"].unique()
    push_events["operator_name"] = op_names.apply(lambda e: str(e).split()[0])
    push_events["operator_address"] = op_address.apply(lambda e: str(e).split()[0])
    push_events.to_csv(output_file)

def data_message_events(files, output_file, path, workers, categories):
        data_messages = DataFrame()

        for (file, i) in zip(files, workers):
            df = pd.read_csv(path + '/' + file, dtype=categories)

            # Push progress events
            push = df[df["event_type"] == "Message"]
            push_groups = push.groupby("channel")
            mean = push_groups["data_length"].mean()
            median = push_groups["data_length"].median()
            data_messages["mean_{}w".format(i)] = mean
            data_messages["median_{}w".format(i)] = median

        data_messages.to_csv(output_file)


def execution_times(files, path, workers, categories):
    for (file, i) in zip(files, workers):
        df = pd.read_csv(path + '/' + file, dtype=categories)

        schedule_events = df[df["event_type"] == "Schedule"]

        only_dataflow_op = schedule_events[schedule_events["operator_id"] == 0]
        total_dataflow_execution = only_dataflow_op["elapsed_time"].sum()

        # Find and omit scoped operators.
        scoped_operators = find_nested_operators(schedule_events["operator_addr"], schedule_events["operator_id"])
        no_scoped_operators = schedule_events[~schedule_events["operator_id"].isin(scoped_operators)]

        total_non_scoped_execution = no_scoped_operators["elapsed_time"].sum()

        #  Process guarded message times:
        guarded_message_events = df[df["event_type"] == "GuardedMessage"]
        total_guarded_time = guarded_message_events["elapsed_time"].sum()

        dataflow_seconds = total_dataflow_execution / 10 ** 9
        operators_seconds = total_non_scoped_execution / 10 ** 9
        guarded_seconds = total_guarded_time / 10 ** 9

        #  Print!
        print("W{} dataflow operator execution: {:.3f}s".format(i, dataflow_seconds))
        print("W{} Non-nested operators execution: {:.3f}s".format(i, operators_seconds))
        print("W{} Guarded seconds: {:.3f}s".format(i, guarded_seconds))
        print("") # new line


def generate_execution_invocation_statistics(files, output_file, path, workers, categories):
    '''
    Generate CSV file with per-operator-id data. Including execution time, and number of invocations, both total and
    normalized values.
    :param files: List of files to process
    :param output_file: File to write final csv to
    :param path: path to directory containing `files`
    :param workers: sorted list of worker numbers. Should correspond to the same order as the files in `files`
    :return: None
    '''
    worker_execution_times = DataFrame()
    worker_invocations = DataFrame()
    avg_time_per_invocation = DataFrame()
    for (file, i) in zip(files, workers):
        df = pd.read_csv(path + '/' + file, dtype=categories)

        schedule_events = df[df["event_type"] == "Schedule"]

        scoped_operators = find_nested_operators(schedule_events["operator_addr"], schedule_events["operator_id"])

        # Omit scoped operators like dataflow and iter.
        schedule_events = schedule_events[~schedule_events["operator_id"].isin(scoped_operators)]
        by_operators = schedule_events.groupby("operator_id")

        # Total operator time in microseconds.
        sums = by_operators["elapsed_time"].sum() / 10 ** 3
        worker_execution_times[i] = sums

        # Get invocation executions.
        invocation_times = by_operators.size()
        worker_invocations[i] = invocation_times
        avg_time_per_invocation["{}w avg time per invocation".format(i)] = sums / invocation_times
    # Get name and address of each opeartor.
    # df indexed by operator id mapping to operator name and address.
    op_names = schedule_events.groupby("operator_id")["operator_name"].unique()
    op_address = schedule_events.groupby("operator_id")["operator_addr"].unique()
    # normalized execution times.
    normalized_times = DataFrame()
    highest_workers_total_execution = worker_execution_times[workers[-1]]
    for (col_name, col_data) in worker_execution_times.iteritems():
        normalized_times["{}w normalized execution".format(col_name)] = col_data / worker_execution_times[workers[0]]
    total_execution_column_name = "{} workers total execution".format(workers[-1])
    normalized_times[total_execution_column_name] = highest_workers_total_execution
    # Normalized invocations
    normalized_invocations = DataFrame()
    highest_workers_total_invocations = worker_invocations[workers[-1]]
    for (col_name, col_data) in worker_invocations.iteritems():
        normalized_times["{}w normalized invocations".format(col_name)] = col_data / worker_invocations[1]
    normalized_invocations[total_execution_column_name] = highest_workers_total_invocations
    # Create final data.
    final_data = DataFrame()
    final_data["operator_name"] = op_names
    final_data["operator_addrs"] = op_address
    # There seems to be no single method to do df concats. So we iterate over all columns.
    for (col_name, col_data) in normalized_times.iteritems():
        final_data[col_name] = col_data
    for (col_name, col_data) in worker_invocations.iteritems():
        # Use a better name than just its integer i number.
        final_data["{}w total invocations".format(col_name)] = col_data
    for (col_name, col_data) in normalized_invocations.iteritems():
        final_data[col_name] = col_data
    for (col_name, col_data) in avg_time_per_invocation.iteritems():
        final_data[col_name] = col_data
    # Append Average data and total number of invocations
    for i in workers:
        mean = final_data["{}w avg time per invocation".format(i)].mean()
        print("Workers {}. Execution time per invocation mean: {}".format(i, mean))
    for i in workers:
        mean = final_data["{}w total invocations".format(i)].mean()
        print("Workers {}. # of invocations mean: {}".format(i, mean))
    # Keep only rows with a good amount of operator invocations.
    # refined_view = final_data[final_data["eight_workers_total_invocations"] > 100]
    refined_view = final_data
    # Clean up look of categorical data.
    # Creates garbage warning which according to this answer nothing can really be done about it:
    # https://stackoverflow.com/questions/42105859/pandas-map-to-a-new-column-settingwithcopywarning
    # WTF. Searching around for way too long verifies that this seemingly trivial task cannot be
    # performed without provoking this warning...
    refined_view["operator_name"] = refined_view["operator_name"].apply(lambda e: str(e).split()[0])
    refined_view["operator_addrs"] = refined_view["operator_addrs"].apply(lambda e: str(e).split()[0])
    refined_view = refined_view.sort_index()
    # sorted_operators = refined_view[["operator_name", "operator_addrs"]]
    # sorted_operators["operator_addrs"] = sorted_operators["operator_addrs"].map(address_to_vec)
    # sorted_operators = sorted_operators.sort_values(by="operator_addrs")
    refined_view.to_csv(output_file)


def address_to_vec(addr):
    """
    Convert an address into a list of numbers.
    """

    # Strip of first and last char, then create list of nums.
    # -2 to also ignore trailing dash.
    numbers = addr[1:-2].split('-')
    v = []
    for e in numbers:
        v.append(int(e))
    return v


def find_nested_operators(addresses, op_ids):
    """
    Given a list of string of the form "n1-n2-...-" (the address strings from operator_addr).
    Return all the nested operators. The way we do it is pretty inefficient but that's okay...
    """
    nested_operators = set()

    for addr in addresses:
        # Ignore trailing dash and split into numbers.
        nums = addr[:-1].split('-')
        # Get all except last number. Which is the address of the nested operator.
        nested_op_addr = '-'.join(str(n) for n in nums[:-1])
        # Add back dash for EZ comparison later.
        nested_operators.add(nested_op_addr + '-')

    # We have the nested operators as addresses, map them back to their operator_id
    # Start going down the list until we find the address that ends with our operator_id. Stop as soon
    # as we find it.
    nested_op_ids = set()
    for nested_addr in nested_operators:
        for (current_addr, op_id) in zip(addresses, op_ids):
            # Found it!
            if current_addr == nested_addr:
                nested_op_ids.add(op_id)
                break

    return nested_op_ids


if __name__ == "__main__":
    main()
