"""
Author: Weisen Pan
Date: 2023-10-24
"""
import re
import pandas as pd

def camel_to_snake(name):
    """Convert a CamelCase string to snake_case."""
    s1 = re.sub('(.)([A-Z][a-z]+)', r'\1_\2', name)
    return re.sub('([a-z0-9])([A-Z])', r'\1_\2', s1).lower()

TAG_LIST = ["InitSchedule", "PostEviction", "PostDeschedule", "ScheduleInflation", "DescheduleInflation"]
TAG_SNAKE_LIST = [camel_to_snake(tag) for tag in TAG_LIST]

POLICY_ABBR_DICT = {
    '01-Random':        'Random',
    '02-DotProd':       'DotProd',
    '03-GpuClustering': 'Clustering',
    '04-GpuPacking':    'Packing',
    '05-BestFit':       'BestFit',
    '06-FGD':           'FGD',
}

def move_tag_to_new_column(df, tag_list=TAG_SNAKE_LIST):
    """Reformat dataframe by moving tags to new columns."""
    meta_cols = [col for col in df.columns if not any(col.endswith("_" + tag) for tag in tag_list)]
    data_cols = [col for col in df.columns if col not in meta_cols]

    output_rows = []

    for _, row in df.iterrows():
        meta_data = {col: row[col] for col in meta_cols}

        for tag in tag_list:
            data_dict = {"tag": tag, **meta_data}

            for col in data_cols:
                if col.endswith("_" + tag):
                    key = col.rstrip("_" + tag)
                    data_dict[key] = row[col]
                    output_rows.append(pd.DataFrame.from_dict(data_dict, orient='index').T)

    return pd.concat(output_rows)

TOTAL_NUM_GPU_DICT = {
    "openb": 6212,
}

def get_total_num_gpu(workload):
    """Retrieve the total number of GPUs based on the workload."""
    return TOTAL_NUM_GPU_DICT.get(workload, 1)

def parse_workload_name(workload_name):
    """Placeholder function to parse the workload name."""
    return workload_name

def shorten_self_policy(policy):
    """Shorten policy names where self reference is found."""
    parts = policy.split('_')
    if len(parts) > 1 and parts[0].split('-')[1] == parts[1]:
        parts[1] = 'self'
        return '_'.join(parts)
    return policy