import os
import glob
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import re

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_dataset_dirs = sorted([
            d for d in glob.glob(os.path.join(results_dir, '*_outputs-MMSeg')) if os.path.isdir(d)
        ])
    except Exception as e:
        print(f"扫描目录 '{results_dir}' 时出错: {e}")
        return None, None

    if not all_dataset_dirs:
        print(f"在 '{results_dir}' 中未找到任何数据集目录 (以 '_outputs-MMSeg' 结尾)。")
        print("请确保脚本与 'BestMode_Predict_Results_DataSet_Public' 文件夹在同一级目录下。")
        return None, None

    print("\n请选择要可视化的数据集:")
    for i, dir_path in enumerate(all_dataset_dirs):
        dataset_name = os.path.basename(dir_path).replace('_outputs-MMSeg', '')
        print(f"  [{i+1}] {dataset_name}")

    while True:
        try:
            choice = input(f"\n请输入选项编号 (1-{len(all_dataset_dirs)}): ")
            choice_idx = int(choice) - 1
            if 0 <= choice_idx < len(all_dataset_dirs):
                selected_dir = all_dataset_dirs[choice_idx]
                dataset_name = os.path.basename(selected_dir).replace('_outputs-MMSeg', '')
                return selected_dir, dataset_name
            else:
                print("无效的选项，请输入列表中的编号。")
        except (ValueError, IndexError):
            print("无效的输入，请输入一个数字编号。")
        except (KeyboardInterrupt, EOFError):
            print("\n操作已取消。")
            return None, None

def plot_performance_speed(selected_dir, dataset_name):
    """
    根据选定的数据集目录，加载数据并生成图表。
    """
    print(f"\n正在为数据集 '{dataset_name}' 生成图表...")

    # 构建文件路径
    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"错误: 在目录 '{selected_dir}' 中缺少所需的数据文件。")
        print(f"  - 检查是否存在: {os.path.basename(metrics_file)}")
        print(f"  - 检查是否存在: {os.path.basename(fps_file)}")
        return

    # 加载数据
    try:
        metrics_df = pd.read_csv(metrics_file)
        # 只保留最新的epoch结果，避免重复
        metrics_df = metrics_df.sort_values('Epoch', ascending=False).drop_duplicates('Algorithm')

        fps_df = pd.read_csv(fps_file)
    except FileNotFoundError as e:
        print(f"错误: 无法找到文件 {e.filename}")
        return
    except Exception as e:
        print(f"读取CSV文件时出错: {e}")
        return

    # 合并两个DataFrame
    # metrics_df中的'Algorithm'列对应fps_df中的'Model'列
    merged_df = pd.merge(metrics_df, fps_df, left_on='Algorithm', right_on='Model')

    if merged_df.empty:
        print("错误: 数据合并失败。请检查 'Algorithm' 和 'Model' 列中的模型名称是否匹配。")
        return
    
    # 调用新函数来创建并保存摘要表格
    T1_create_and_save_summary_table(merged_df, selected_dir, dataset_name)

    # 调用新函数来提取和保存所有IoU数据
    T2_extract_and_save_iou_data(metrics_df, selected_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(selected_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(selected_dir, output_filename_svg)
    plt.savefig(save_file_path_svg)
    print(f"\n图表已成功生成并保存为: {save_file_path_svg} 和 {save_file_path_png}")
    plt.show()

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

    Args:
        merged_df (pd.DataFrame): 包含所有模型指标和性能数据的DataFrame。
        output_dir (str): 保存CSV文件的目标目录。
        dataset_name (str): 数据集的名称，用于生成文件名。
    """
    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文件内容。")
        return

    # 提取并复制数据，避免修改原始DataFrame
    summary_df = merged_df[required_columns].copy()
    
    # 清理和转换数据
    # 将 '118 G' -> 118.0
    summary_df['FLOPs'] = summary_df['FLOPs'].astype(str).str.replace(' G', '', regex=False).astype(float)
    # 将 '13.274M' -> 13.274
    summary_df['Params'] = summary_df['Params'].astype(str).str.replace('M', '', regex=False).astype(float)

    # 按照用户的要求重命名列
    summary_df.rename(columns={
        'Average_FPS': 'FPS',
        'FLOPs': 'G(GFLOPS)',
        'Params': 'Para(Params)'
    }, 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文件。

    Args:
        metrics_df (pd.DataFrame): 包含所有指标的原始DataFrame。
        output_dir (str): 保存CSV文件的目标目录。
        dataset_name (str): 数据集的名称，用于生成文件名。
    """
    print("正在提取所有 mIoU 和 Class_IoU 数据...")

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

    # 找出所有与IoU相关的列
    # 包括 'mIoU' 以及所有以 '_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 降序排序，便于查看
    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__':
    # 设置包含所有数据集结果的根目录
    results_root_dir = '../BestMode_Predict_Results_DataSet_Public'
    
    # 启动交互式选择
    selected_directory, selected_dataset_name = select_dataset(results_root_dir)
    
    # 如果用户成功选择，则生成图表
    if selected_directory and selected_dataset_name:
        plot_performance_speed(selected_directory, selected_dataset_name)