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

# TODO 这个是获取最后一次结果的 TODO

# --- 配置日志记录 ---
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# --- 辅助函数 ---
def find_all_log_files_sorted(algorithm_dir: str) -> List[str]:
    """
    查找给定算法目录中所有的.log文件，并按从新到旧的顺序排列。

    Args:
        algorithm_dir (str): 算法的根目录。

    Returns:
        List[str]: 按时间倒序排列的日志文件路径列表。
    """
    try:
        subdirs = [d for d in os.listdir(algorithm_dir) if os.path.isdir(os.path.join(algorithm_dir, d))]
    except FileNotFoundError:
        logging.error(f"算法目录不存在: {algorithm_dir}")
        return []

    if not subdirs:
        logging.warning(f"在目录 {algorithm_dir} 中未找到任何时间戳子目录。")
        return []

    # 按名称倒序排序，最新的目录会排在最前面
    sorted_subdirs = sorted(subdirs, reverse=True)

    log_files = []
    for subdir_name in sorted_subdirs:
        subdir_path = os.path.join(algorithm_dir, subdir_name)
        logs_in_subdir = glob.glob(os.path.join(subdir_path, '*.log'))
        if logs_in_subdir:
            # 假设每个子目录只有一个log文件
            log_files.append(logs_in_subdir[0])
    
    return log_files

def parse_log_metrics(log_path: str) -> Optional[Dict]:
    """
    解析日志文件，提取最后一次完整验证（validation）的结果及其对应的Epoch。

    Args:
        log_path (str): 日志文件的路径。

    Returns:
        Optional[Dict]: 包含 'epoch', 'summary', 'class_wise' 指标的字典，如果解析失败则返回 None。
    """
    try:
        with open(log_path, 'r', encoding='utf-8') as f:
            content = f.read()
    except IOError as e:
        logging.error(f"无法读取日志文件: {log_path}。错误: {e}")
        return None
    # 定义正则表达式
    summary_pattern = re.compile(
        r"Iter\(val\) \[\d+/\d+\]\s+aAcc:\s*([\d.]+)\s+mIoU:\s*([\d.]+)\s+mAcc:\s*([\d.]+)"
    )
    class_table_pattern = re.compile(
        # --- 使用新的模式匹配可变长度的顶部边框 ---
        r"\+(?:-+\+)+\s*\n"
        r"\|.*?Class.*?\|.*?IoU.*?\|.*?Acc.*?\|\n"
        # --- 匹配中间边框 ---
        r"\+(?:-+\+)+\s*\n"
        # --- 捕获表格主体 ---
        r"((?:\|.*?\|.*?\|.*?\|\n)+)"
        # --- 匹配底部边框 ---
        r"\+(?:-+\+)+\s*\n",
        re.MULTILINE
    )
    epoch_pattern = re.compile(r"Saving checkpoint at (\d+) epochs|resumed epoch: (\d+)")

    # 查找所有匹配项
    summary_matches = list(re.finditer(summary_pattern, content))
    table_matches = list(re.finditer(class_table_pattern, content))
    epoch_matches = list(re.finditer(epoch_pattern, content))

    if not summary_matches or not table_matches:
        logging.warning(f"❌ 在日志 {os.path.basename(log_path)} 中未能找到完整的验证结果。")
        return None

    last_summary_match = summary_matches[-1]
    last_table_match = None
    for table in reversed(table_matches):
        if table.end() < last_summary_match.start():
            last_table_match = table
            break
            
    if not last_table_match:
        logging.warning(f"❌ 在日志 {os.path.basename(log_path)} 中找到总结行但未能匹配到对应的类别表格。")
        return None

    # 寻找关联的最新Epoch
    last_epoch = "N/A"
    latest_epoch_num = -1
    for epoch_match in epoch_matches:
        if epoch_match.end() < last_table_match.start():
            epoch_str = epoch_match.group(1) or epoch_match.group(2)
            if epoch_str:
                epoch_num = int(epoch_str)
                if epoch_num > latest_epoch_num:
                   latest_epoch_num = epoch_num
                   last_epoch = f"epoch_{epoch_num}"

    # 解析数据
    summary_groups = last_summary_match.groups()
    results = {
        'epoch': last_epoch,
        'summary': {
            'aAcc': summary_groups[0],
            'mIoU': summary_groups[1],
            'mAcc': summary_groups[2]
        },
        'class_wise': []
    }

    table_content = last_table_match.group(0)
    row_pattern = re.compile(r"\|\s*([\w\s]+?)\s*\|\s*([\d.]+)\s*\|\s*([\d.]+)\s*\|")
    for line in table_content.strip().split('\n'):
        row_match = row_pattern.match(line)
        if row_match:
            class_name, iou, acc = row_match.groups()
            results['class_wise'].append({
                'Class': class_name.strip(),
                'IoU': iou,
                'Acc': acc
            })
    
    if results['class_wise']:
        logging.info(f"✅ 成功从 {os.path.basename(log_path)} 中解析出 Epoch '{last_epoch}' 的指标。")
        return results
    else:
        logging.warning(f"❌ 在 {os.path.basename(log_path)} 中未能解析出任何类别行。")
        return None

# --- 主函数 ---
def main(args):
    """
    脚本主入口，负责编排整个自动化分析流程。
    """
    input_root = args.input_dir
    output_root = args.output_dir
    
    if not os.path.isdir(input_root):
        logging.error(f"输入目录不存在: {input_root}")
        return

    # --- 交互式菜单 ---
    all_dataset_dirs = sorted([
        d for d in glob.glob(os.path.join(input_root, '*_outputs-MMSeg')) if os.path.isdir(d)
    ])

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

    dataset_map = {str(i + 1): path for i, path in enumerate(all_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 = []
    selected_dataset_dir = None
    
    if choice1 in dataset_map:
        selected_dataset_dir = dataset_map[choice1]
        logging.info(f"您已选择数据集: [{os.path.basename(selected_dataset_dir)}]")

        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)} 中未发现任何算法子文件夹。")
            return
            
        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()

        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("无效的算法选择，程序已退出。")
            return
    else:
        logging.error("无效的数据集选择，程序已退出。")
        return

    # --- 开始处理选定的算法 (逻辑已修改) ---
    csv_rows = []
    output_dataset_folder = ""

    for model_dir in model_dirs:
        model_name = os.path.basename(model_dir)
        logging.info(f"\n--- 开始处理算法: {model_name} ---")

        # (路径构建代码保持不变)
        if not output_dataset_folder:
            dataset_key = model_name.split('-')[0]
            dataset_folder_map = {
                '1_cholecseg8k': '1_CholecSeg8k-13Type-1920x1080_outputs-MMSeg',
                '2_autolaparo': '2_AutoLaparo-10Type-1920x1080_outputs-MMSeg',
                '3_1_endovis_2017': '3_1_Endovis_2017-8Type-512x512_outputs-MMSeg',
                '3_2_endovis_2018': '3_2_Endovis_2018-8Type-512x512_outputs-MMSeg',
                '4_dresden': '4_Dresden-11Type-512x512_outputs-MMSeg'
            }
            output_dataset_folder = dataset_folder_map.get(dataset_key, f"{dataset_key}_outputs-MMSeg")
        
        # --- 新的循环查找逻辑 ---
        # 1. 获取所有按时间倒序排列的日志文件
        all_logs_sorted = find_all_log_files_sorted(model_dir)
        
        if not all_logs_sorted:
            logging.warning(f"跳过算法 {model_name}，因为未找到任何日志文件。")
            continue

        # 2. 循环尝试解析，直到成功或全部失败
        metrics = None
        for log_file_path in all_logs_sorted:
            logging.info(f"正在尝试解析: {os.path.relpath(log_file_path)}")
            metrics = parse_log_metrics(log_file_path)
            if metrics:
                logging.info(f"在 {os.path.basename(log_file_path)} 中成功找到并解析了指标。")
                break  # 找到后立即跳出循环

        # 3. 如果所有日志都尝试失败，则跳过此算法
        if not metrics:
            logging.warning(f"❌❌❌跳过算法 {model_name}，因为在其所有日志文件中都未能找到有效的指标。❌❌❌")
            continue
            
        # --- 创建一个 "宽" 格式的行 ---
        summary = metrics['summary']
        short_model_name = model_name.split('Alg_', 1)[1]
        row_data = {
            'Algorithm': short_model_name,
            'Epoch': metrics['epoch'],
            'mIoU': summary['mIoU'],
            'mAcc': summary['mAcc'],
            'aAcc': summary['aAcc']
        }
        
        # 将每个类别的IoU和Acc作为新列添加到行数据中
        for class_data in metrics['class_wise']:
            class_name = class_data['Class'].replace(' ', '_') # 清理类名以用作表头
            row_data[f'{class_name}_IoU'] = class_data['IoU']
            row_data[f'{class_name}_Acc'] = class_data['Acc']
        
        csv_rows.append(row_data)

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

    # --- 动态生成并排序表头 ---
    # 基础列保持固定顺序
    base_fieldnames = ['Algorithm', 'Epoch', 'mIoU', 'mAcc', 'aAcc']
    # 从第一个结果中获取所有与类别相关的列名，并按字母排序
    first_row_keys = csv_rows[0].keys()
    class_fieldnames = sorted([key for key in first_row_keys if key not in base_fieldnames])
    # 最终的完整表头
    final_fieldnames = base_fieldnames + class_fieldnames
    
    # --- 构建输出路径并写入文件 ---
    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]
    # 在文件名中加入 "_wide" 以区分格式
    output_csv_path = os.path.join(final_output_dir, f'{dataset_name}_metrics_summary_wide.csv')
    
    try:
        with open(output_csv_path, 'w', newline='', encoding='utf-8-sig') as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=final_fieldnames)
            writer.writeheader()
            writer.writerows(csv_rows)
        
        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 最终指标提取脚本 (V2)")
    parser.add_argument(
        '--input_dir',
        type=str,
        default='../Hardisk',
        help="包含数据集输出文件夹 (例如 '..._outputs-MMSeg') 的根目录。"
    )
    parser.add_argument(
        '--output_dir',
        type=str,
        default='../BestMode_Predict_Results_DataSet_Public',
        help="用于存储所有分析结果的根目录。"
    )
    
    args = parser.parse_args()
    main(args)