data_engineering.py

"""
Performs data cleaning on aggregated profile data in csv format.
Capability to fill missing values, add indicator columns, detect complete,
partial, or unique features by model, filter on system data only, GPU data,
or indicator columns only.

Operates from the aggregated csv generated by format_profiles.py
"""
from pathlib import Path

import numpy as np
import pandas as pd
from format_profiles import read_csv
import config
import json

def subsample(df: pd.DataFrame, num: int, col: str = "model") -> pd.DataFrame:
    """
    Return a dataframe with the first <num> rows from each value in column <col>.
    For example, if col="model" and num=10, then samples the first 10 profiles per model.
    """
    return df.groupby(col).head(num).reset_index(drop=True).copy(deep=True)


def removeColumnsFromOther(keep_cols, remove_df):
    """
    Given two dataframes keep_cols and remove_df,
    remove each column of remove_df if that column is
    not in keep_cols.
    """
    to_remove = [x for x in remove_df.columns if x not in keep_cols.columns]
    return remove_cols(remove_df, to_remove)


def softmax(x):
    return(np.exp(x - np.max(x)) / np.exp(x - np.max(x)).sum())

def remove_cols(df: pd.DataFrame, substrs: list, endswith: bool = False, verbose=False):
    """Removes columns of the dataframe which include (or end with) any of the substrings."""

    if not isinstance(substrs, list):
        substrs = [substrs]

    def remove_col(col_name: str):
        for substr in substrs:
            if not endswith and substr in col_name:
                return True
            if endswith and col_name.endswith(substr):
                return True
        return False

    remove_columns = [col_name for col_name in df.columns if remove_col(col_name)]

    if verbose:
        print("\nRemoving columns:")
        for i in remove_columns:
            print(i)

    return df.drop(remove_columns, axis=1)


def filter_cols(df: pd.DataFrame, substrs: list, keep: list = None, verbose: bool=False):
    """Only keeps columns containing substrings"""
    if keep is None:
        keep = ["model", "model_family", "file"]

    def keep_col(col_name: str):
        for substr in substrs:
            if substr in col_name:
                return True
        for x in keep:
            if col_name == x:
                return True
        return False
    
    remove_columns = [col_name for col_name in df.columns if not keep_col(col_name)]
    if verbose:
        print("\nRemoving columns:")
        for i in remove_columns:
            print(i)

    return df.drop(remove_columns, axis=1)


def get_csv(aggregated_csv_folder, remove_nans=True, gpu_activities_only=False, api_calls_only=False):
    df = read_csv(aggregated_csv_folder, gpu_activities_only=gpu_activities_only, api_calls_only=api_calls_only)
    if remove_nans:
        df = remove_cols(df, "nan", verbose=True)
    return df


def missing_data(aggregated_csv_folder, gpu_activities_only=False, api_calls_only=False):
    """
    Returns the columns and number of missing datapoints by model.  Missing data are denoted by NaN.

    :param aggregated_csv_folder: path to the folder under ./profiles/ which contains the aggregated.csv file.
    :return: A dict keyed by model with the columns with missing data and the number of missing datapoints.
    """
    df = get_csv(aggregated_csv_folder, gpu_activities_only=gpu_activities_only, api_calls_only=api_calls_only)
    model_nans = {}

    for model in df["model"].unique():
        model_df = df.loc[df["model"] == model]
        n_model_profiles = len(model_df.index)
        model_nans[model] = model_df.isna().sum() / n_model_profiles

    # model_nans is now a dict of {model_type: Pandas series of {column_name: % of that column that is empty}}

    return model_nans


def mutually_exclusive_data(aggregated_csv_folder, gpu_activities_only=False, api_calls_only=False):
    """
    Returns information about the data as a dict with elements described below:

    mutually_exclusive_attributes: the attributes which are exclusive to each model.  Only include an attribute
        if each model profile includes this attribute and none of the others include it.  Organized by model

    partial_attributes: attributes which models have for some profiles but not others, organized by model.

    no_data_attributes: list of attributes for which there is no data from any model.

    complete_attributes: A list of attributes which each model has completely.

    partially_exclusive_attributes: A mapping from attributes to models for attributes which some but
        not all models have

    :param aggregated_csv_folder: path to the folder under ./profiles/ which contains the aggregated.csv file.
    """

    model_nans = missing_data(aggregated_csv_folder, gpu_activities_only=gpu_activities_only, api_calls_only=api_calls_only)
    num_models = len(list(model_nans.keys()))

    # dict of model: exclusive attributes to that model (this model has it completely and no other model has any)
    mutually_exclusive = {model: [] for model in model_nans}

    # dict of model: attributes which that model has partially (some missing values)
    partial_attribute = {model: [] for model in model_nans}

    # array of attributes on which no model has data
    attributes_with_no_data = []

    # array of attributes on which all models have complete data
    complete_attributes = []

    # dict of attribute: models with that attribute completely.
    # Only includes attributes that at least 1 model doesnt have and at least 1 model has.
    partially_exclusive = {}

    for attribute in model_nans[next(iter(model_nans))].axes[0]:    # gets the first element
        models_with_complete_attribute = []
        models_with_partial_attribute = []
        models_without_attribute = []
        for model in model_nans:
            percent_empty = model_nans[model][attribute]
            if percent_empty == 0.0:
                models_with_complete_attribute.append(model)
            if 0.0 < percent_empty < 1.0:
                models_with_partial_attribute.append(model)
                # this model has this attribute partially
                partial_attribute[model].append(attribute)
            if percent_empty == 1.0:
                models_without_attribute.append(model)
        if len(models_with_complete_attribute) == 1 and len(models_with_partial_attribute) == 0:
            # only one model has this attribute
            mutually_exclusive[models_with_complete_attribute[0]].append(attribute)
        if len(models_without_attribute) == num_models:
            # no model has this attribute
            attributes_with_no_data.append(attribute)
        if len(models_with_complete_attribute) == num_models:
            # all models have this attribute completly
            complete_attributes.append(attribute)
        if 0 < len(models_with_complete_attribute) < num_models and len(models_without_attribute) > 0:
            # at least 1 model has this and at least 1 model doesn't have this attribute
            partially_exclusive[attribute] = models_with_complete_attribute

    res = {
        "mutually_exclusive_attributes": mutually_exclusive,
        "partial_attributes": partial_attribute,
        "no_data_attributes": attributes_with_no_data,
        "complete_attributes": complete_attributes,
        "partially_exclusive_attributes": partially_exclusive
    }

    return res


def shared_data(agg_csv_folder, system_data_only=False, no_system_data=False, gpu_activities_only=False, api_calls_only=False):
    """
    Return a dataframe containing only complete features that can be used for machine learning.

    :param agg_csv_folder: the folder under ./profiles/ where the aggregated csv lives.
    :param system_data_only: If true, only return system data (clock, temp, power, fan).
    :param no_system_data: if true, excludes system data.
    :return: a dataframe
    """
    if system_data_only and no_system_data:
        raise ValueError("system_data_only and no_system_data cannot both be true.")

    df = get_csv(agg_csv_folder)
    complete_attributes = mutually_exclusive_data(agg_csv_folder, gpu_activities_only=gpu_activities_only, api_calls_only=api_calls_only)["complete_attributes"]
    df = df[complete_attributes]    # only consider complete data

    if not system_data_only and not no_system_data:
        return df

    def system_column(col):
        for sys_signal in config.SYSTEM_SIGNALS:
            if col.endswith(sys_signal):
                return True
        return False

    system_cols = [col_name for col_name in complete_attributes if system_column(col_name)]

    if system_data_only:
        system_cols.append('model')
        system_cols.append('model_family')
        system_cols.append('file')
        return df[system_cols]

    # else no_system_data is true
    return df.drop(system_cols, axis=1)


def train_test_split(df, ratio=0.8):
    """Splits the df into a train and test set with equal representation over all the model classes."""
    test_df = pd.DataFrame()
    train_df = pd.DataFrame()
    for model in config.MODELS:
        model_rows = df[df["model"] == model]
        num_rows = len(model_rows.index)

        train_rows = int(num_rows * ratio)
        train_df = pd.concat([train_df, model_rows.head(train_rows)], ignore_index=True)

        test_rows = num_rows - train_rows
        test_df = pd.concat([test_df, model_rows.tail(test_rows)], ignore_index=True)

    return train_df, test_df


def get_data_and_labels(df, shuffle=True, label=None):
    """Splits a dataframe into data points and their associated labels (model)."""
    if shuffle:
        df = df.sample(frac=1)

    if not label:
        y = df["model"]
    else:
        y = df[label]

    x = df.drop(columns=["file", "model_family", "model"], axis=1)

    return x, y


def add_indicator_columns(df):
    """
    For any column that includes NaNs, add a binary indicator column.
    Missing values are replaced by the mean.

    Since GPU data is split into 6 features (min, max, avg, num_calls, time_ms, time_percent),
    only 1 indicator column will be added for each of these 6 features.
    
    Indicators have the name indicator_<meta_feature>.  For example, a meta
    feature is [CUDA memcpy DtoD] which is associated with features
    avg_us_[CUDA memcpy DtoD], max_ms_[CUDA memcpy DtoD], min_us_[CUDA memcpy DtoD]
    time_percent_[CUDA memcpy DtoD], time_ms_[CUDA memcpy DtoD], num_calls_[CUDA memcpy DtoD].

    There will only be one indicator per meta feature.
    """
    prefixes = ['min_us_', 'max_ms_', 'avg_us_', 'num_calls_', 'time_ms_', 'time_percent_']
    
    def stripPrefix(col_name: str) -> str:
        for prefix in prefixes:
            if col_name.startswith(prefix):
                return col_name[len(prefix):]
        raise ValueError
    
    for col in df.columns:
        column = df[col]
        if column.isna().sum() > 0:
            indicator_col = column.notna().astype(int)
            indicator_col.name = f"indicator_{stripPrefix(col)}"
            mean = column.mean()
            df[col] = df[col].fillna(mean)
            # before adding the column, check if the indicator for this feature has
            # already been added
            if indicator_col.name in df.columns:
                continue
            df = pd.concat([df, indicator_col], axis=1)
        if df[col].isna().sum() > 0:
            raise RuntimeError(f"Still NaNs in column {col}")
    return df


def add_indicator_cols_to_input(df, x: pd.Series, exclude: list = []) -> pd.Series:
    """
    For an input x, add indicator columns to x so that it has the same 
    columns as df.  Replace missing values in x with mean.  Return
    new x.  Drop columns of x that are not in df.
    Excluded columns are skipped.

    Indicators have the name indicator_<meta_feature>.  For example, a meta
    feature is [CUDA memcpy DtoD] which is associated with features
    avg_us_[CUDA memcpy DtoD], max_ms_[CUDA memcpy DtoD], min_us_[CUDA memcpy DtoD]
    time_percent_[CUDA memcpy DtoD], time_ms_[CUDA memcpy DtoD], num_calls_[CUDA memcpy DtoD].

    There will only be one indicator per meta feature.

    Therefore to check if an input x has a meta_feature, you have to see if any of the
    features avg_ms_{meta_feature}, max_ms_(meta_feature), ... exist in x.
    """

    # first add indicators
    for col in df.columns:
        # only look at indicator columns in the reference df
        if col in exclude or not col.startswith("indicator_"):
            continue
        meta_feature = col.split("indicator_")[1]
        found = False
        for col_name in x.keys():
            if col_name.find(meta_feature) >= 0:
                x[col] = 1
                found = True
                break
        if not found:
            x[col] = 0
    
    # now all indicator columns are added, fill in missing
    # values with the mean
    for col in df.columns:
        if col in exclude or col.startswith("indicator_"):
            continue
        if col not in x.keys():
            x[col] = df[col].mean()

    for i in x.keys():
        if i not in df.columns:
            x = x.drop(i)
    
    # sort x in order of the keys of the df.
    x = pd.DataFrame([x])
    new_df = pd.concat((df, x), ignore_index=True)

    result = new_df.iloc[-1]
    result = result.drop(exclude)
    return result


def all_data(agg_csv_folder, system_data_only=False, no_system_data=False, indicators_only=False, gpu_activities_only=False, api_calls_only=False):
    """
    Return a dataframe containing all features.  Creates indicator columns for incomplete
    features, and fills NaNs with the mean of that feature.

    :param agg_csv_folder: the folder where the aggregated csv lives.
    :param system_data_only: If true, only return system data (clock, temp, power, fan).
    :param no_system_data: if true, excludes system data.
    :param indicators_only: only return indicator columns
    :return: a dataframe
    """
    if system_data_only and no_system_data:
        raise ValueError("system_data_only and no_system_data cannot both be true.")

    df = get_csv(agg_csv_folder, gpu_activities_only=gpu_activities_only, api_calls_only=api_calls_only)

    df = add_indicator_columns(df)

    if indicators_only and system_data_only:
        raise ValueError

    if indicators_only:
        indicator_cols = [col for col in df.columns if col.startswith("indicator")]
        indicator_cols.append('model')
        indicator_cols.append('model_family')
        indicator_cols.append('file')
        return df[indicator_cols]

    if not system_data_only and not no_system_data:
        return df

    def system_column(col):
        for sys_signal in config.SYSTEM_SIGNALS:
            if col.endswith(sys_signal):
                return True
        return False

    system_cols = [col_name for col_name in df.columns if system_column(col_name)]

    if system_data_only:
        system_cols.append('model')
        system_cols.append('model_family')
        system_cols.append('file')
        return df[system_cols]

    # else no_system_data is true
    return df.drop(system_cols, axis=1)


if __name__ == '__main__':
    # test = shared_data("zero_noexe", system_data_only=True)
    # print(test)
    test = mutually_exclusive_data("zero_noexe_lots_models")
    print(json.dumps(test, indent=4))
    all_data("zero_noexe_lots_models")
    # test = mutually_exclusive_data("zero_noexe")
    # print(json.dumps(test, indent=4))

    # test = missing_data("debug_profiles")
    # print(test)
    exit(0)