import os
import glob
import logging
import argparse
import re
import subprocess
import csv
from typing import Dict, Optional, Tuple, List

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# --- 辅助函数 ---
def find_model_files(model_dir: str):
    """
    在给定的模型目录中查找配置文件、最佳检查点和日志文件。

    Args:
        model_dir (str): 模型的根目录。

    Returns:
        Optional]: 包含 'config', 'checkpoint', 'log' 路径的字典，
                                  如果缺少任何必要文件，则返回 None。
    """
    config_files = glob.glob(os.path.join(model_dir, '*.py'))
    if not config_files:
        logging.warning(f"在目录 {model_dir} 中未找到配置文件 (.py)。")
        return None
    config_path = config_files[0]

    checkpoint_path = os.path.join(model_dir, 'best.pth')
    if not os.path.exists(checkpoint_path):
        epoch_files = glob.glob(os.path.join(model_dir, 'epoch_*.pth'))
        if not epoch_files:
            logging.warning(f"在目录 {model_dir} 中未找到 'best.pth' 或 'epoch_*.pth' 检查点文件。")
            return None
        
        # 通过正则表达式从文件名中提取周期数并找到最大的
        latest_epoch = -1
        latest_file = None
        for f in epoch_files:
            match = re.search(r'epoch_(\d+)\.pth', os.path.basename(f))
            if match:
                epoch_num = int(match.group(1))
                if epoch_num > latest_epoch:
                    latest_epoch = epoch_num
                    latest_file = f
        
        if latest_file:
            checkpoint_path = latest_file
        else:
            logging.warning(f"在目录 {model_dir} 中无法确定最新的检查点文件。")
            return None

    return {'config': config_path, 'checkpoint': checkpoint_path}

def find_model_config(model_dir: str):
    """
    在给定的模型目录中查找配置文件 (.py)。

    Args:
        model_dir (str): 模型的根目录。

    Returns:
        Optional[str]: 配置文件的路径，如果未找到则返回 None。
    """
    config_files = glob.glob(os.path.join(model_dir, '*.py'))
    if not config_files:
        logging.warning(f"在目录 {model_dir} 中未找到配置文件 (.py)。")
        return None
    return config_files[0]

def get_shape_from_path(path: str):
    """
    从文件夹路径中通过正则表达式提取分辨率 (宽x高)。

    Args:
        path (str): 数据集文件夹的路径。

    Returns:
        Optional[Tuple[int, int]]: 一个包含 (高度, 宽度) 的元组，如果未找到则返回 None。
                                    注意：工具需要 H W 格式。
    """
    match = re.search(r'(\d+)x(\d+)', os.path.basename(path))
    if match:
        width, height = int(match.group(1)), int(match.group(2))
        return (height, width) # 返回 H, W
    return None

def get_flops_and_params(config_path: str, shape: Tuple[int, int]) -> Optional[Dict[str, str]]:
    """
    运行 mmsegmentation 的 get_flops.py 工具并解析其输出。
    此版本适配了新版的直接输出格式 (例如 "Flops: 0.118T")。

    Args:
        config_path (str): 模型的 .py 配置文件路径。
        shape (Tuple[int, int]): 输入图像的 (H, W) 元组。

    Returns:
        Optional[Dict[str, str]]: 包含 'params' 和 'flops' 的字典，如果失败则返回 None。
    """
    # 检查工具脚本是否存在
    tool_script = 'tools/analysis_tools/get_flops.py'
    if not os.path.exists(tool_script):
        logging.error(f"错误: '{tool_script}' 未找到。请确保在 MMSegmentation 项目的根目录下运行此脚本。")
        return None

    # 构建命令行
    command = [
        'python', tool_script, config_path,
        '--shape', str(shape[0]), str(shape[1])
    ]

    logging.info(f"执行命令: {' '.join(command)}")

    try:
        # 执行命令并捕获输出
        result = subprocess.run(command, capture_output=True, text=True, check=True, encoding='utf-8')
        output = result.stdout
        
        # 使用新的正则表达式来匹配更新后的输出格式
        flops_match = re.search(r"Flops:\s*([0-9.]+\s*[TGMK]?)", output)
        params_match = re.search(r"Params:\s*([0-9.]+\s*[TGMK]?)", output)

        if flops_match and params_match:
            raw_flops_str = flops_match.group(1).strip()
            params = params_match.group(1).strip()
            # --- 开始单位换算 ---
            value_str = raw_flops_str.rstrip('TGMKtgmk').strip()
            unit = raw_flops_str[-1].upper() if raw_flops_str[-1].isalpha() else 'G'
            try:
                value = float(value_str)
                if unit == 'T':
                    value_in_g = value * 1000
                elif unit == 'M':
                    value_in_g = value / 1000
                elif unit == 'K':
                    value_in_g = value / 1_000_000
                else: # 默认单位是 G
                    value_in_g = value
                # 使用 :g 格式化可以去除末尾多余的0
                flops = f"{value_in_g:g} G"
            except ValueError:
                flops = raw_flops_str # 如果转换失败，则使用原始值
            # --- 单位换算结束 ---

            logging.info(f"✅ 解析成功: FLOPs={flops} (原始值: {raw_flops_str}), Params={params}")
            return {'flops': flops, 'params': params}
        else:
            logging.warning(f"❌ 无法从命令输出中解析 FLOPs 或 Params。请检查以下输出内容:\n---\n{output}\n---")
            return None

    except subprocess.CalledProcessError as e:
        logging.error(f"执行 get_flops.py 时出错。返回码: {e.returncode}")
        logging.error(f"错误输出 (stderr):\n---\n{e.stderr}\n---")
        return None
    except FileNotFoundError:
        logging.error("错误: 'python' 命令未找到。请确保 Python 环境已正确配置。")
        return None

# --- 主函数 ---
def main(args):
    """
    脚本主入口，负责编排整个自动化分析流程。
    """
    input_root = args.input_dir
    output_root = args.output_dir
    # --- 开始交互式选择修改 (V2 - 两级菜单) ---
    if not os.path.isdir(input_root):
        logging.error(f"输入目录不存在: {input_root}")
        return

    # 1. 定义有效的数据集文件夹白名单
    VALID_DATASET_FOLDERS = [
        '1_CholecSeg8k-13Type-1920x1080_outputs-MMSeg',
        '2_AutoLaparo-10Type-1920x1080_outputs-MMSeg',
        '3_1_Endovis_2017-8Type-512x512_outputs-MMSeg',
        '3_2_Endovis_2018-8Type-512x512_outputs-MMSeg',
        '4_Dresden-11Type-512x512_outputs-MMSeg'
    ]
    
    # 2. 查找存在的、有效的数据集目录
    existing_dataset_dirs = [
        os.path.join(input_root, d) for d in VALID_DATASET_FOLDERS 
        if os.path.isdir(os.path.join(input_root, d))
    ]

    if not existing_dataset_dirs:
        logging.error(f"在输入目录 {input_root} 中未找到任何有效的数据集文件夹。")
        return

    # 3. 第一级菜单：选择数据集
    dataset_map = {str(i + 1): path for i, path in enumerate(existing_dataset_dirs)}
    
    print("\n" + "="*50)
    print("--- 步骤 1: 请选择要处理的数据集 ---")
    for key, path in dataset_map.items():
        print(f"{key}: {os.path.basename(path)}")
    print("="*50)
    
    choice1 = input("请输入数据集编号并按回车键: ").strip()

    model_dirs = [] # 初始化最终要处理的目录列表
    
    if choice1 in dataset_map:
        selected_dataset_dir = dataset_map[choice1]
        logging.info(f"您已选择数据集: [{os.path.basename(selected_dataset_dir)}]")

        # 4. 查找选定数据集下的所有算法子目录
        alg_dirs = sorted([
            d for d in glob.glob(os.path.join(selected_dataset_dir, '*')) if os.path.isdir(d)
        ])

        if not alg_dirs:
            logging.warning(f"在 {os.path.basename(selected_dataset_dir)} 中未发现任何算法子文件夹。程序退出。")
        else:
            # 5. 第二级菜单：选择算法
            alg_map = {str(i + 1): path for i, path in enumerate(alg_dirs)}
            print("\n" + "="*50)
            print("--- 步骤 2: 请选择要处理的算法 ---")
            print("0: 批量处理当前数据集下的【全部】算法")
            for key, path in alg_map.items():
                print(f"{key}: {os.path.basename(path)}")
            print("="*50)

            choice2 = input("请输入算法编号 (或输入 '0' 处理全部) 并按回车键: ").strip()

            # 6. 根据第二级选择，最终确定 model_dirs
            if choice2 == '0':
                model_dirs = alg_dirs
                logging.info(f"您选择了批量处理全部 {len(model_dirs)} 个算法。")
            elif choice2 in alg_map:
                model_dirs = [alg_map[choice2]] # 将单个路径放入列表中
                logging.info(f"您选择了处理单个算法: {os.path.basename(model_dirs[0])}")
            else:
                logging.error("无效的算法选择，程序已退出。")
    else:
        logging.error("无效的数据集选择，程序已退出。")
    
    # --- 交互式选择修改结束 ---
    results: List[Dict[str, str]] = []
    # 修改后的循环，将遍历经过用户筛选后的 model_dirs 列表
    for model_dir in model_dirs:
        model_name = os.path.basename(model_dir)
        logging.info(f"--- 开始处理模型: {model_name} ---")
        files = find_model_files(model_dir)
        if not files:
            logging.warning(f"跳过目录 {model_dir}，因为缺少必要文件。")
            continue

        # 构建输出目录
        # 从模型名中提取数据集标识作为Key
        dataset_key = model_name.split('-')[0] 
        dataset_map = {
            '1_cholecseg8k': '1_CholecSeg8k-13Type-1920x1080_outputs-MMSeg',
            '2_autolaparo': '2_AutoLaparo-10Type-1280x1024_outputs-MMSeg',
            '3_1_endovis_2017': '3_1_EndoVis_2017-7Type-1280x1024_outputs-MMSeg',
            '3_2_endovis_2018': '3_2_EndoVis_2018-11Type-1280x1024_outputs-MMSeg',
            '4_dresden': '4_Dresden-6Type-1920x1080_outputs-MMSeg'
        }
        # 使用提取的Key（字符串）进行查询，并为默认值也使用该Key
        output_dataset_folder = dataset_map.get(dataset_key, f"{dataset_key}_outputs-MMSeg")

        # 尝试从数据集文件夹名称中获取分辨率
        input_shape = get_shape_from_path(selected_dataset_dir)
        if not input_shape:
            # 如果无法自动提取，要求用户输入
            logging.warning(f"无法从文件夹 '{os.path.basename(selected_dataset_dir)}' 名称中自动检测分辨率。")
            try:
                h_str = input("请输入默认测试高度 (H)，例如 512: ").strip()
                w_str = input("请输入默认测试宽度 (W)，例如 512: ").strip()
                input_shape = (int(h_str), int(w_str))
            except ValueError:
                logging.error("输入无效，必须是整数。程序退出。")
                return
        logging.info(f"将使用输入形状 (H, W): {input_shape} 进行计算。")

        # 加载配置
        config_file = find_model_config(model_dir)

        # 获取 FLOPs 和 Params
        stats = get_flops_and_params(config_file, input_shape)
        if stats:
            short_model_name = model_name.split('Alg_', 1)[1]
            results.append({
                'Model': short_model_name,
                'Params': stats['params'],
                'FLOPs': stats['flops'],
                'Input_Shape (HxW)': f"{input_shape[0]}x{input_shape[1]}"
            })
        else:
            logging.warning(f"未能获取模型 {model_name} 的统计信息。")

    # --- 将结果写入 CSV 文件 ---
    if not results:
        logging.info("没有成功获取任何模型的统计数据，不生成 CSV 文件。")
        return

    # 新建文件夹并保存 CSV
    final_output_dir = os.path.join(output_root, output_dataset_folder)
    os.makedirs(final_output_dir, exist_ok=True)
    dataset_name = os.path.basename(selected_dataset_dir).split('_outputs-MMSeg')[0]
    output_csv_path = os.path.join(final_output_dir, f'{dataset_name}_flops_params_summary.csv')
    
    try:
        with open(output_csv_path, 'w', newline='', encoding='utf-8') as csvfile:
            fieldnames = ['Model', 'Params', 'FLOPs', 'Input_Shape (HxW)']
            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
            
            writer.writeheader()
            writer.writerows(results)
        
        logging.info(f"=== 全部处理完成！结果已成功保存到: {output_csv_path} ===")
    except IOError as e:
        logging.error(f"无法写入 CSV 文件: {output_csv_path}。错误: {e}")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="MMSegmentation 自动化评估脚本")
    parser.add_argument(
        '--input_dir',
        type=str,
        default='../Hardisk',
        help="包含已训练模型文件夹的根目录。"
    )
    parser.add_argument(
        '--output_dir',
        type=str,
        default='../BestMode_Predict_Results_DataSet_Public',
        help="用于存储所有分析结果的根目录。"
    )
    
    args = parser.parse_args()
    main(args)