import pandas as pd
import gradio as gr
import plotly.graph_objects as go
import plotly.express as px
import numpy as np

type_emoji = {
    "RTL-Specific": "🔴",
    "General": "🟢",
    "Coding": "🔵"
}

def model_hyperlink(link, model_name):
    return f'<a target="_blank" href="{link}" style="color: var(--link-text-color); text-decoration: underline;text-decoration-style: dotted;">{model_name}</a>'

def handle_special_cases(benchmark, metric):
    if metric == 'Exact Matching (EM)':
        benchmark = 'RTL-Repo'
    elif benchmark == 'RTL-Repo':
        metric = 'Exact Matching (EM)'
    return benchmark, metric

def filter_RTLRepo(subset: pd.DataFrame) -> pd.DataFrame:
    details = subset[['Model', 'Model URL', 'Model Type', 'Params']].drop_duplicates('Model')
    filtered_df = subset[['Model', 'Score']].rename(columns={'Score': 'Exact Matching (EM)'})
    filtered_df = pd.merge(filtered_df, details, on='Model', how='left')
    filtered_df['Model'] = filtered_df.apply(lambda row: model_hyperlink(row["Model URL"], row["Model"]), axis=1)
    filtered_df['Type'] = filtered_df['Model Type'].map(lambda x: type_emoji.get(x, ""))
    filtered_df = filtered_df[['Type', 'Model', 'Params', 'Exact Matching (EM)']]
    filtered_df = filtered_df.sort_values(by='Exact Matching (EM)', ascending=False).reset_index(drop=True)
    return filtered_df

def filter_bench(subset: pd.DataFrame, df_agg=None, agg_column=None) -> pd.DataFrame:
    details = subset[['Model', 'Model URL', 'Model Type', 'Params']].drop_duplicates('Model')
    pivot_df = subset.pivot_table(index='Model', columns='Metric', values='Score', aggfunc='mean').reset_index()
    
    if df_agg is not None and agg_column is not None and agg_column in df_agg.columns:
        agg_data = df_agg[['Model', agg_column]].rename(columns={agg_column: 'Aggregated ⬆️'})
        pivot_df = pd.merge(pivot_df, agg_data, on='Model', how='left')
    else:# fallback
        pivot_df['Aggregated ⬆️'] = pivot_df.mean(axis=1, numeric_only=True).round(2)
    
    pivot_df = pd.merge(pivot_df, details, on='Model', how='left')
    pivot_df['Model'] = pivot_df.apply(lambda row: model_hyperlink(row["Model URL"], row["Model"]), axis=1)
    pivot_df['Type'] = pivot_df['Model Type'].map(lambda x: type_emoji.get(x, ""))
    pivot_df.rename(columns={'Syntax (STX)': 'STX', 'Functionality (FNC)': 'FNC', 'Synthesis (SYN)': 'SYN', 'Performance': 'Perf'}, inplace=True)
    
    columns_order = ['Type', 'Model', 'Params', 'Aggregated ⬆️', 'STX', 'FNC', 'SYN', 'Power', 'Perf', 'Area']
    pivot_df = pivot_df[[col for col in columns_order if col in pivot_df.columns]]
    pivot_df = pivot_df.sort_values(by='Aggregated ⬆️', ascending=False).reset_index(drop=True)
    return pivot_df

def filter_bench_all(subset: pd.DataFrame, df_agg=None, agg_column=None) -> pd.DataFrame:
    details = subset[['Model', 'Model URL', 'Model Type', 'Params']].drop_duplicates('Model')
    pivot_df = subset.pivot_table(index='Model', columns='Metric', values='Score', aggfunc='mean').reset_index().round(2)
    
    if df_agg is not None:
        if agg_column is not None and agg_column in df_agg.columns:
            agg_data = df_agg[['Model', agg_column]].rename(columns={agg_column: 'Aggregated ⬆️'})
            pivot_df = pd.merge(pivot_df, agg_data, on='Model', how='left')
        else:
            agg_columns = [col for col in df_agg.columns if col.startswith('Agg ')]
            if agg_columns:
                df_agg['Average_Agg'] = df_agg[agg_columns].mean(axis=1)
                agg_data = df_agg[['Model', 'Average_Agg']].rename(columns={'Average_Agg': 'Aggregated ⬆️'})
                pivot_df = pd.merge(pivot_df, agg_data, on='Model', how='left')
            else: # fallback
                pivot_df['Aggregated ⬆️'] = pivot_df.mean(axis=1, numeric_only=True).round(2)
    else: # fallback
        pivot_df['Aggregated ⬆️'] = pivot_df.mean(axis=1, numeric_only=True).round(2)
    
    pivot_df = pd.merge(pivot_df, details, on='Model', how='left')
    pivot_df['Model'] = pivot_df.apply(lambda row: model_hyperlink(row["Model URL"], row["Model"]), axis=1)
    pivot_df['Type'] = pivot_df['Model Type'].map(lambda x: type_emoji.get(x, ""))
    
    pivot_df.rename(columns={
        'Exact Matching (EM)': 'EM', 
        'Syntax (STX)': 'Avg STX', 
        'Functionality (FNC)': 'Avg FNC', 
        'Synthesis (SYN)': 'Avg SYN', 
        'Power': 'Avg Power',
        'Performance': 'Avg Perf',
        'Area': 'Avg Area',
    }, inplace=True)
    
    columns_order = ['Type', 'Model', 'Params', 'Aggregated ⬆️', 'Avg STX', 'Avg FNC', 'Avg SYN', 'Avg Power', 'Avg Perf', 'Avg Area']
    pivot_df = pivot_df[[col for col in columns_order if col in pivot_df.columns]]
    pivot_df = pivot_df.sort_values(by='Aggregated ⬆️', ascending=False).reset_index(drop=True)
    return pivot_df