import pandas as pd
from datasets import load_dataset

from src.display.utils import COLS, METRICS_COLS, BenchRawColumn, BenchTableModel, fields
from src.display.utils import COLS_TO_AVERAGE, COLS_DEPEND_ON_MODEL, TASK_TYPE_MAP


def get_leaderboard_df_from_hf_dataset(path: str) -> pd.DataFrame:
    dataset = load_dataset(path, split="train")
    df = pd.DataFrame(dataset)
    headers = {e.raw_name: e.name for e in fields(BenchRawColumn)}
    df = df.rename(columns=headers)
    cparams = BenchRawColumn.model_params.name
    df[cparams] = (df[cparams] / 1000000).round()  # mandatory
    cinfer = BenchRawColumn.infer_time.name
    df[cinfer] = (df[cinfer] * 1000).round(2)  # infer time per 1k samples
    BenchTableModel.infer_time.name = "Infer Time (s) - 1k samples"
    df = df[COLS]
    df = df.round(5)
    df[METRICS_COLS] = df[METRICS_COLS].round(2)

    return df

def summarize_model_task_type_performance(df):
    """Summarizes model performance across task types"""
    df = df.copy()
    df['Task'] = df['Task'].map(TASK_TYPE_MAP)
    grouped = df.groupby(['Model', 'Task'])

    avg_aggs = {col: ['mean', 'std'] for col in COLS_TO_AVERAGE}
    model_aggs = {col: 'first' for col in COLS_DEPEND_ON_MODEL}
    agg_dict = {**avg_aggs, **model_aggs}
    agg_df = grouped.agg(agg_dict).reset_index()

    agg_df.columns = [
        f"{col[0]}_{col[1]}" if isinstance(col, tuple) and col[1] in ['mean', 'std']
        else col[0] if isinstance(col, tuple)
        else col
        for col in agg_df.columns
    ]

    for col in COLS_TO_AVERAGE:
        mean_col = f"{col}_mean"
        std_col = f"{col}_std"

        agg_df[col] = agg_df[mean_col].round(3).astype(str) + ' ± ' + agg_df[std_col].round(3).astype(str)
        agg_df.drop(columns=[mean_col, std_col], inplace=True)


    final_cols = ['Model', 'Task'] + COLS_TO_AVERAGE + COLS_DEPEND_ON_MODEL
    agg_df = agg_df[final_cols]

    return agg_df