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 get_max_epochs(config_path: str) -> Optional[int]:
    """
    从config.py文件中解析train_cfg字典以获取max_epochs的值。

    Args:
        config_path (str): config.py文件的路径。

    Returns:
        Optional[int]: max_epochs的值，如果找不到则返回None。
    """
    if not os.path.exists(config_path):
        logging.error(f"配置文件不存在: {config_path}")
        return None
        
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            content = f.read()
        
        # 使用正则表达式查找 max_epochs
        match = re.search(r"train_cfg\s*=\s*dict\(.*?max_epochs\s*=\s*(\d+),.*?\)", content, re.DOTALL)
        
        if match:
            max_epochs = int(match.group(1))
            logging.info(f"从 {os.path.basename(config_path)} 中成功读取 max_epochs: {max_epochs}")
            return max_epochs
        else:
            logging.warning(f"在 {config_path} 中未找到 'max_epochs'。")
            return None
    except Exception as e:
        logging.error(f"解析 {config_path} 时出错: {e}")
        return None

def parse_log_metrics(log_path: str, max_epochs: int) -> List[Dict]:
    """
    解析日志文件，提取所有完整验证（validation）的结果，并计算其对应的Epoch。

    Args:
        log_path (str): 日志文件的路径。
        max_epochs (int): 从config.py中读取的最大epoch数。

    Returns:
        List[Dict]: 包含每次验证的 'epoch', 'summary', 'class_wise' 指标的字典列表。
    """
    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 []

    # 定义所有需要的正则表达式
    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
    )
    train_iter_pattern = re.compile(r"Iter\(train\)\s*\[\s*(\d+)\s*/\s*(\d+)\]")

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

    all_metrics = []

    # 遍历每一次的总结行 (summary)
    for i, summary_match in enumerate(summary_matches):
        # 寻找与总结行对应的类别表格
        # 表格应该出现在总结行之前
        last_table_match = None
        for table in reversed(table_matches):
            if table.end() < summary_match.start():
                # 确保这个表格没有被上一个总结行用过
                is_already_used = False
                if i > 0:
                    if table.end() < summary_matches[i-1].start():
                        is_already_used = True
                if not is_already_used:
                    last_table_match = table
                    break
        
        if not last_table_match:
            continue

        # 寻找表格前最近的 Iter(train) 行来计算epoch
        last_train_iter_match = None
        for train_iter in reversed(train_iter_matches):
            if train_iter.end() < last_table_match.start():
                last_train_iter_match = train_iter
                break
        
        epoch = "N/A"
        if last_train_iter_match and max_epochs:
            current_iter, total_iters = last_train_iter_match.groups()
            try:
                # 根据公式计算epoch
                epoch = int(int(current_iter) / int(total_iters) * max_epochs)
            except (ValueError, ZeroDivisionError) as e:
                logging.warning(f"Epoch 计算失败: {e}")

        # 解析总结指标
        summary_groups = summary_match.groups()
        results = {
            'epoch': epoch,
            'summary': {
                'aAcc': summary_groups[0],
                'mIoU': summary_groups[1],
                'mAcc': summary_groups[2]
            },
            'class_wise': []
        }

        # 解析每个类别的数据
        table_content = last_table_match.group(1)
        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']:
            all_metrics.append(results)

    if all_metrics:
        logging.info(f"✅ 成功从 {os.path.basename(log_path)} 中解析出 {len(all_metrics)} 组指标。")
    else:
        logging.warning(f"❌ 在日志 {os.path.basename(log_path)} 中未能找到完整的验证结果。")
        
    return all_metrics

# --- 主函数 ---
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]
            # ... (这部分路径映射逻辑保持不变)
            output_dataset_folder = f"{dataset_key}_outputs-MMSeg"

        all_logs_sorted = find_all_log_files_sorted(model_dir)
        
        if not all_logs_sorted:
            logging.warning(f"跳过算法 {model_name}，因为未找到任何日志文件。")
            continue
        
        # --- 新逻辑：获取max_epochs ---
        # 假设同一算法下所有训练的config是相同的，因此我们从最新的log对应的config读取
        latest_log_path = all_logs_sorted[0]
        config_path = os.path.join(os.path.dirname(latest_log_path), 'vis_data', 'config.py')
        max_epochs = get_max_epochs(config_path)
        if max_epochs is None:
            logging.error(f"无法为算法 {model_name} 找到 max_epochs，将跳过。")
            continue

        # --- 新逻辑：聚合所有日志的所有指标 ---
        all_metrics_for_model = []
        for log_file_path in all_logs_sorted:
            logging.info(f"正在解析: {os.path.relpath(log_file_path)}")
            # 传递max_epochs
            metrics_from_log = parse_log_metrics(log_file_path, max_epochs)
            if metrics_from_log:
                all_metrics_for_model.extend(metrics_from_log)

        if not all_metrics_for_model:
            logging.warning(f"❌❌❌跳过算法 {model_name}，因为在其所有日志文件中都未能找到有效的指标。❌❌❌")
            continue
            
        # --- 新逻辑：选择mIoU最高的记录 ---
        try:
            best_metric = max(all_metrics_for_model, key=lambda x: float(x['summary']['mIoU']))
            logging.info(f"找到了最佳指标: Epoch '{best_metric['epoch']}', mIoU: {best_metric['summary']['mIoU']}")
        except (ValueError, TypeError) as e:
            logging.error(f"为算法 {model_name} 寻找最佳mIoU时出错: {e}")
            continue

        # --- 创建一个 "宽" 格式的行 ---
        summary = best_metric['summary']
        short_model_name = model_name.split('Alg_', 1)[1] if 'Alg_' in model_name else model_name
        row_data = {
            'Algorithm': short_model_name,
            'Epoch': best_metric['epoch'],
            'mIoU': summary['mIoU'],
            'mAcc': summary['mAcc'],
            'aAcc': summary['aAcc']
        }
        
        for class_data in best_metric['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, os.path.basename(selected_dataset_dir))
    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}_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"\n=== 全部处理完成！最佳结果已成功保存到: {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)