import gradio as gr
import pandas as pd
from functools import reduce

CONFIG = safe_load(open("config.yaml"))

data = defaultdict(dict)

# 读取数据
for settings in CONFIG['settings']:
    for type in CONFIG['types']:
        # 根据配置文件中的路径读取数据
        data[settings][type] = pd.read_excel('data/' + CONFIG["settings_mapping"][settings] + f"-{type}.xlsx")
        
# 添加平均分列并四舍五入到小数点后两位
for settings in CONFIG['settings']:
    for type in CONFIG['types']:
        data[settings][type]["Average"] = data[settings][type].iloc[:,1:].mean(axis=1).round(2)
# 添加rank列
for settings in CONFIG['settings']:
    for type in CONFIG['types']:
        data[settings][type]["Rank"] = data[settings][type]["Average"].rank(ascending=False, method='min').astype(int)
# 将rank列放在第一列
for settings in CONFIG['settings']:
    for type in CONFIG['types']:
        cols = data[settings][type].columns.tolist()
        cols = cols[-1:] + cols[:-1]
        data[settings][type] = data[settings][type][cols]

# 四舍五入整个DataFrame到小数点后两位
for settings in CONFIG['settings']:
    for type in CONFIG['types']:
        data[settings][type] = data[settings][type].round(2)

css = """
table > thead {
    white-space: normal;
}

table {
    --cell-width-1: 250px;
}

table > tbody > tr > td:nth-child(2) > div {
    overflow-x: auto;
}

.filter-checkbox-group {
    max-width: max-content;
}

/* 确保第二列(Model)完全展开 */
table > tbody > tr > td:nth-child(2) {
    white-space: nowrap;
    width: auto;
}

/* 紧凑显示其他列 */
table > tbody > tr > td:not(:nth-child(2)) {
    white-space: normal;
    width: auto;
}
"""

# 定义模型类型和大小（占位符）
MODEL_TYPES = [
    "Open",
    "Proprietary",
    "Sentence Transformers",
    "Cross-Encoders",
    "Bi-Encoders",
    "Uses Instructions",
    "No Instructions",
]

NUMERIC_INTERVALS = {
    "<100M": pd.Interval(0, 100, closed="right"),
    "100M to 250M": pd.Interval(100, 250, closed="right"),
    "250M to 500M": pd.Interval(250, 500, closed="right"),
    "500M to 1B": pd.Interval(500, 1000, closed="right"),
    ">1B": pd.Interval(1000, 1_000_000, closed="right"),
}

# 定义过滤函数
def filter_data(search_query, model_types, model_sizes):
    output_df = df.copy()
    
    # Apply the search query
    if search_query:
        names = output_df.index.str.lower()
        masks = []
        for query in search_query.split(";"):
            masks.append(names.str.contains(query.lower()))
        output_df = output_df[reduce(lambda a, b: a | b, masks)]
    
    # Apply the model type filtering
    if set(model_types) != set(MODEL_TYPES):
        # Placeholder logic for model type filtering
        pass
    
    # Apply the model size filtering
    if model_sizes:
        # Placeholder logic for model size filtering
        pass
    
    return output_df.round(2)

# 创建示例 DataFrame (假设 df 是一个示例 DataFrame)
df = pd.DataFrame({
    'Model': ['ModelA', 'ModelB'],
    'Metric1': [0.123456, 0.789123],
    'Metric2': [0.654321, 0.321654],
    'Average': [0.0, 0.0],
    'Rank': [0, 0]
})

# 计算 Average 和 Rank 并四舍五入
df['Average'] = df.iloc[:, 1:-1].mean(axis=1).round(2)
df['Rank'] = df['Average'].rank(ascending=False, method='min').astype(int)
cols = df.columns.tolist()
cols = cols[-1:] + cols[:-1]
df = df[cols]

# Create the Gradio interface
with gr.Blocks(css=css) as demo:
    with gr.Box(className="container"):
        gr.Markdown("# Model Leaderboard")
        
        with gr.Row():
            search_box = gr.Textbox(
                label="Search Models (separate by ';')",
                placeholder="🔍 Search for a model and press enter...",
                scale=3
            )
            model_type_checkbox_group = gr.CheckboxGroup(
                label="Model types",
                choices=MODEL_TYPES, 
                value=MODEL_TYPES,
                interactive=True,
                elem_classes=["filter-checkbox-group"],
                scale=3
            )
            model_size_checkbox_group = gr.CheckboxGroup(
                label="Model sizes (in number of parameters)",
                choices=list(NUMERIC_INTERVALS.keys()),
                value=list(NUMERIC_INTERVALS.keys()),
                interactive=True,
                elem_classes=["filter-checkbox-group"],
                scale=2,
            )
        submit_button = gr.Button("Filter Data")
        
        with gr.Tabs() as result_table:
            for settings in CONFIG['settings']:
                with gr.Tab(label=settings):
                    for type in CONFIG['types']:
                        with gr.Box(className="container"):
                            gr.Markdown(f"### {type} Leaderboard")
                            gr.DataFrame(data[settings][type], type="pandas")
                


demo.launch()