import os
import glob
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import re
import argparse
from collections import defaultdict

def get_model_family(model_name):
    """
    根据模型名称提取模型族。
    例如: 'my_bisenetv1_r50' -> 'my_bisenetv1'
          'my_fast_scnn' -> 'my_fast_scnn'
    """
    # 使用正则表达式匹配，将 _rXX 或 _dXX 等后缀去掉
    match = re.match(r'^(.*?)_r\d+$', model_name)
    if match:
        return match.group(1)
    return model_name

def select_dataset(results_dir):
    """
    扫描目录，对数据集进行分组，让用户交互式选择一个数据集进行合并分析。
    """
    print("正在扫描可用的数据集...")
    try:
        # 查找所有匹配后缀的目录
        all_dirs = glob.glob(os.path.join(results_dir, '*_outputs-MMSeg')) + \
                   glob.glob(os.path.join(results_dir, '*_outputs-SegModel'))

        if not all_dirs:
            print(f"在 '{results_dir}' 中未找到任何数据集目录 (以 '_outputs-MMSeg' 或 '_outputs-SegModel' 结尾)。")
            return None, None

        # --- 新增逻辑：按基本数据集名称对目录进行分组 ---
        datasets_map = defaultdict(list)
        for dir_path in all_dirs:
            if os.path.isdir(dir_path):
                # 提取基本名称，例如 '1_CholecSeg8k-13Type-1920x1080'
                base_name = re.sub(r'_outputs-(MMSeg|SegModel)$', '', os.path.basename(dir_path))
                datasets_map[base_name].append(dir_path)

        sorted_dataset_names = sorted(datasets_map.keys())

    except Exception as e:
        print(f"扫描目录 '{results_dir}' 时出错: {e}")
        return None, None

    print("\n请选择要合并分析的数据集:")
    for i, name in enumerate(sorted_dataset_names):
        # 显示每个数据集包含的源文件夹数量
        source_count = len(datasets_map[name])
        print(f"  [{i+1}] {name} ({source_count}个源)")

    while True:
        try:
            choice = input(f"\n请输入选项编号 (1-{len(sorted_dataset_names)}): ")
            choice_idx = int(choice) - 1
            if 0 <= choice_idx < len(sorted_dataset_names):
                selected_name = sorted_dataset_names[choice_idx]
                selected_dirs = datasets_map[selected_name] # 获取与所选数据集关联的所有目录
                return selected_dirs, selected_name
            else:
                print("无效的选项，请输入列表中的编号。")
        except (ValueError, IndexError):
            print("无效的输入，请输入一个数字编号。")
        except (KeyboardInterrupt, EOFError):
            print("\n操作已取消。")
            return None, None

def F1_plot_performance_speed(selected_dirs, dataset_name, output_base_dir):
    """
    根据选定的数据集目录列表，加载并合并数据、生成图表和表格，并保存到指定的输出目录。

    Args:
        selected_dirs (list): 用户选择的原始数据所在的所有目录的列表。
        dataset_name (str): 从目录名中提取的数据集名称。
        output_base_dir (str): 保存所有输出文件的根目录。
    """
    print(f"\n正在为数据集 '{dataset_name}' 合并数据并生成图表...")

    # 在指定的输出根目录下，为当前数据集创建一个专属的输出文件夹
    dataset_output_dir = os.path.join(output_base_dir, dataset_name)
    os.makedirs(dataset_output_dir, exist_ok=True)
    print(f"所有输出文件将被保存到: {dataset_output_dir}")

    # --- 修改逻辑：从多个目录加载并合并数据 ---
    all_metrics = []
    all_fps = []

    print("正在读取以下来源的数据:")
    for selected_dir in selected_dirs:
        print(f" - {os.path.basename(selected_dir)}")
        metrics_file = os.path.join(selected_dir, f"{dataset_name}_metrics_summary_wide.csv")
        fps_file = os.path.join(selected_dir, f"{dataset_name}_flops_params_fps_summary.csv")

        # 检查文件是否存在
        if not os.path.exists(metrics_file) or not os.path.exists(fps_file):
            print(f"  -> 警告: 在目录 '{os.path.basename(selected_dir)}' 中缺少数据文件，已跳过。")
            continue

        try:
            metrics_df_part = pd.read_csv(metrics_file)
            all_metrics.append(metrics_df_part)

            fps_df_part = pd.read_csv(fps_file)
            all_fps.append(fps_df_part)
        except Exception as e:
            print(f"  -> 错误: 读取CSV文件时出错: {e}")
            continue

    if not all_metrics or not all_fps:
        print("\n错误: 未能从任何有效的源目录中加载数据，无法继续生成报告。")
        return

    # 合并来自所有源的数据
    metrics_df = pd.concat(all_metrics, ignore_index=True)
    fps_df = pd.concat(all_fps, ignore_index=True)
    print("\n数据合并完成。")

    # 对合并后的数据进行去重处理
    if 'Epoch' in metrics_df.columns:
        metrics_df = metrics_df.sort_values('Epoch', ascending=False).drop_duplicates('Algorithm')
    else:
        metrics_df = metrics_df.drop_duplicates('Algorithm')
        
    fps_df = fps_df.drop_duplicates('Model')

    # 合并两个DataFrame
    merged_df = pd.merge(metrics_df, fps_df, left_on='Algorithm', right_on='Model', how='inner')

    if merged_df.empty:
        print("错误: 数据合并失败。请检查 'Algorithm' 和 'Model' 列中的模型名称是否完全匹配。")
        print(f"  - 指标文件中的模型: {metrics_df['Algorithm'].unique()}")
        print(f"  - 性能文件中的模型: {fps_df['Model'].unique()}")
        return

    # 调用函数创建并保存摘要表格到新的输出目录
    T1_create_and_save_summary_table(merged_df, dataset_output_dir, dataset_name)

    # 调用函数来提取和保存所有IoU数据到新的输出目录
    T2_extract_and_save_iou_data(metrics_df, dataset_output_dir, dataset_name)

    # 提取模型族
    merged_df['Family'] = merged_df['Model'].apply(get_model_family)

    # --- 绘图 ---
    plt.style.use('seaborn-v0_8-whitegrid')
    fig, ax = plt.subplots(figsize=(16, 10))

    # 定义颜色和标记
    families = sorted(merged_df['Family'].unique())
    palette = sns.color_palette("husl", len(families))
    markers = ['o', 's', 'X', 'D', '^', 'P', '*', 'v', '<', '>']

    # 循环绘制每个模型族
    for i, family in enumerate(families):
        family_df = merged_df[merged_df['Family'] == family].sort_values('Average_FPS')
        color = palette[i]
        marker = markers[i % len(markers)]

        # 绘制散点
        ax.scatter(family_df['Average_FPS'], family_df['mIoU'],
                   color=color, marker=marker, s=150, label=family, zorder=3)

        # 如果族内有多个模型，则用线连接
        if len(family_df) > 1:
            ax.plot(family_df['Average_FPS'], family_df['mIoU'],
                    color=color, linestyle='--', linewidth=1.5, zorder=2)

        # 在每个点旁边添加模型全名注释
        for j, row in family_df.iterrows():
            ax.text(row['Average_FPS'] * 1.01, row['mIoU'], row['Model'],
                    fontsize=9, verticalalignment='center')

    # 设置图表属性
    ax.set_title(f'Model Performance vs. Inference Speed ({dataset_name})', fontsize=18, pad=20)
    ax.set_xlabel('Inference Speed (FPS)', fontsize=14)
    ax.set_ylabel('Mean IoU (%)', fontsize=14)
    ax.legend(title='Model Family', bbox_to_anchor=(1.02, 1), loc='upper left', borderaxespad=0.)

    plt.tight_layout(rect=[0, 0, 0.88, 1]) # 调整布局为图例留出空间
    plt.grid(True, which='both', linestyle='--', linewidth=0.5)

    # 保存图表到新的输出目录
    output_filename_png = f"F1_{dataset_name}_mIoU_vs_FPS.png"
    save_file_path_png = os.path.join(dataset_output_dir, output_filename_png)
    plt.savefig(save_file_path_png, dpi=600)

    output_filename_svg = f"F1_{dataset_name}_mIoU_vs_FPS.svg"
    save_file_path_svg = os.path.join(dataset_output_dir, output_filename_svg)
    plt.savefig(save_file_path_svg)

    print(f"\n图表已成功生成并保存为: {save_file_path_svg} 和 {save_file_path_png}")
    plt.close(fig) # 关闭图形，避免在循环中使用时重复显示

def T1_create_and_save_summary_table(merged_df, output_dir, dataset_name):
    """
    根据合并后的数据创建、格式化并保存性能摘要表格。
    """
    print("正在创建摘要表格...")

    # 检查所需列是否存在
    required_columns = ['Model', 'mIoU', 'mAcc', 'aAcc', 'Average_FPS', 'FLOPs', 'Params']
    if not all(col in merged_df.columns for col in required_columns):
        print("错误: DataFrame中缺少必要的列。请检查CSV文件内容。")
        print(f"  - 需要的列: {required_columns}")
        print(f"  - 实际的列: {merged_df.columns.tolist()}")
        return

    # 提取并复制数据，避免修改原始DataFrame
    summary_df = merged_df[required_columns].copy()

    # 清理和转换数据
    summary_df['FLOPs'] = summary_df['FLOPs'].astype(str).str.replace(r'\s*G', '', regex=True).astype(float)
    summary_df['Params'] = summary_df['Params'].astype(str).str.replace(r'\s*M', '', regex=True).astype(float)

    # 按照用户的要求重命名列
    summary_df.rename(columns={
        'Average_FPS': 'FPS',
        'FLOPs': 'FLOPs(G)',
        'Params': 'Params(M)'
    }, inplace=True)

    # 按 mIoU 降序排序
    summary_df = summary_df.sort_values(by='mIoU', ascending=False)

    # 保存表格到CSV文件
    summary_filename = f"T1_{dataset_name}_performance_summary.csv"
    summary_save_path = os.path.join(output_dir, summary_filename)

    try:
        summary_df.to_csv(summary_save_path, index=False, float_format='%.3f')
        print(f"摘要表格已成功保存到: {summary_save_path}")
    except Exception as e:
        print(f"保存摘要表格时出错: {e}")

def T2_extract_and_save_iou_data(metrics_df, output_dir, dataset_name):
    """
    从 metrics DataFrame 中提取所有 mIoU 和 Class_IoU，并保存到新的CSV文件。
    """
    print("正在提取所有 mIoU 和 Class_IoU 数据...")

    # 检查'Algorithm'列是否存在
    if 'Algorithm' not in metrics_df.columns:
        print("错误: 'Algorithm' 列未找到，无法继续。")
        return

    # 找出所有与IoU相关的列
    iou_columns = ['Algorithm', 'mIoU'] + [col for col in metrics_df.columns if col.endswith('_IoU') and col != 'mIoU']

    # 移除重复的列名（以防万一）
    iou_columns = list(dict.fromkeys(iou_columns))

    # 提取数据
    iou_df = metrics_df[iou_columns].copy()

    # 按 mIoU 降序排序，便于查看
    if 'mIoU' in iou_df.columns:
        iou_df = iou_df.sort_values(by='mIoU', ascending=False)

    # 定义并保存文件
    iou_filename = f"T2_{dataset_name}_all_iou_summary.csv"
    iou_save_path = os.path.join(output_dir, iou_filename)

    try:
        iou_df.to_csv(iou_save_path, index=False, float_format='%.2f')
        print(f"所有IoU数据已成功保存到: {iou_save_path}")
    except Exception as e:
        print(f"保存IoU数据时出错: {e}")

if __name__ == '__main__':
    # --- 设置命令行参数解析 ---
    parser = argparse.ArgumentParser(description="从模型评估结果生成性能与速度对比图和摘要表。")
    parser.add_argument(
        '--input_dir',
        type=str,
        default='../BestMode_Predict_Results_DataSet_Public',
        help="包含所有数据集结果的根目录 (例如 '..._outputs-MMSeg' 或 '..._outputs-SegModel' 的父目录)。"
    )
    parser.add_argument(
        '--output_dir',
        type=str,
        default='./',
        help="用于存储所有生成的图表和表格的根目录。"
    )
    args = parser.parse_args()

    # 确保输出目录存在
    os.makedirs(args.output_dir, exist_ok=True)

    # 启动交互式选择
    selected_directories, selected_dataset_name = select_dataset(args.input_dir)

    # 如果用户成功选择，则生成图表和表格
    if selected_directories and selected_dataset_name:
        F1_plot_performance_speed(selected_directories, selected_dataset_name, args.output_dir)