import os
import pandas as pd
from src.leaderboard.read_evals import load_model_data


def get_leaderboard_df(results_path):
    """获取完整的排行榜数据，包含所有指标"""
    all_models = []

    if not os.path.exists(results_path) or not os.path.isdir(results_path):
        print(f"结果路径 {results_path} 不存在或不是目录")
        return pd.DataFrame()

    # 遍历所有模型文件夹
    for model_folder in os.listdir(results_path):
        print(f"加载模型: {model_folder}")
        model_data = load_model_data(model_folder, results_path)

        if model_data:
            all_models.append(model_data)
        else:
            print(f"模型 {model_folder} 加载失败")

    # 转换为DataFrame并返回
    df = pd.DataFrame(all_models)
    print(f"成功加载 {len(df)} 个模型数据")
    return df


def get_filtered_leaderboard(results_path, model_type="All", strategies=[], filter_mode="交集 (满足所有选中策略)"):
    """
    获取筛选后的排行榜数据

    参数:
    - results_path: 结果目录路径
    - model_type: 模型类型筛选
    - strategies: 策略筛选列表
    - filter_mode: 筛选模式，"交集 (满足所有选中策略)" 或 "并集 (满足任一选中策略)"
    """
    # 获取完整数据
    df = get_leaderboard_df(results_path)
    if df.empty:
        return df

    # 应用模型类型筛选
    if model_type and model_type != "All":
        # 筛选包含该模型类型的行
        df = df[df['model_type'].str.contains(model_type)]

    # 应用策略筛选
    if strategies:
        # 获取对应的策略列
        strategy_columns = [f'strategy_{s}' for s in strategies if f'strategy_{s}' in df.columns]

        if strategy_columns:
            # 根据筛选模式应用不同的筛选逻辑
            if filter_mode == "交集 (满足所有选中策略)":
                # 所有选中策略都必须为True
                df = df[df[strategy_columns].all(axis=1)]
            else:
                # 至少一个选中策略为True
                df = df[df[strategy_columns].any(axis=1)]
        else:
            print(f"没有找到与选中策略匹配的列: {strategies}")
            return pd.DataFrame()

    return df