# -*- coding: utf-8 -*-
"""
MMSegmentation 自动化推理、评估与分析脚本

本脚本旨在为基于 MMSegmentation 训练的语义分割模型提供一个全面、自动化的后处理流水线。
它能够系统地发现指定目录下的训练产物，并为每个模型执行一系列详尽的分析任务，最终生成结构化的报告。

核心功能：
1.  模型发现与检查点选择：
    - 自动扫描指定输入目录（默认为 './Outputs_mmseg'）下的所有子文件夹。
    - 智能选择用于推理的权重文件：优先选择 'best.pth'，若不存在，则选择周期数（epoch）最大的检查点文件。
    - 定位与模型对应的配置文件（.py）和训练日志文件（.json）。

2.  验证集推理与结果生成：
    - 对每个模型的验证集进行推理，并使用 tqdm 显示进度条。
    - 生成并保存原始的、未经着色的预测掩码图（predicted_raw_masks）。
    - 生成并保存包含“原始图像-预测图（着色后）-真值图标注（着色后）”的三图对比分析图像（prediction_analysis）。

3.  综合性能评估：
    - 计算并记录多项关键评估指标，包括 IoU (mIoU), F1-Score (mFscore/mDice), Accuracy (aAcc/mAcc), Recall (mRecall), 和 Precision (mPrecision)。
    - 深入计算每个类别的混淆矩阵基本元素：真正例（TP）、真负例（TN）、假正例（FP）、假负例（FN）。这一功能通过对 IoUMetric 内部数据的再处理实现，提供了标准评估流程之外的精细化分析能力。

4.  模型复杂度分析：
    - 计算模型的参数量（Parameters）和计算量（FLOPs），以评估其资源消耗。
    - 该计算基于 MMEngine 提供的底层分析工具，确保了与 MMSegmentation 官方工具的一致性。

5.  训练过程可视化：
    - 解析训练日志文件（.json），提取损失（loss）、准确率（accuracy）、交并比（mIoU）等关键指标在训练过程中的变化。
    - 使用 Matplotlib 和 Seaborn 绘制指标变化曲线图，并保存为图像文件，便于直观分析模型的收敛情况。

6.  结构化结果输出：
    - 所有生成的分析结果（指标、图像、日志图表）均被保存在一个结构化的输出目录中（默认为 './BestMode_Predict_Results_DataSet_Public'）。
    - 每个模型的输出都存放在以其原始文件夹命名的子目录中，保持清晰的对应关系。
    - 最终的量化指标（包括总体指标、模型复杂度、逐类TP/TN/FP/FN）被整合并保存到一个名为 'test_set_metrics.csv' 的文件中，采用长格式（long format）存储，便于后续的数据处理和比较。

使用方法：
    直接在终端运行此脚本。可以通过命令行参数指定输入和输出的根目录。
    python predict.py --input_dir./Outputs_mmseg --output_dir./BestMode_Predict_Results_DataSet_Public

依赖：
    - torch, torchvision
    - mmengine
    - mmcv
    - mmsegmentation>=1.0.0
    - numpy
    - pandas
    - matplotlib
    - seaborn
    - tqdm
    - opencv-python
"""

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

import cv2
import numpy as np
import pandas as pd
import torch
import mmcv
import matplotlib.pyplot as plt
import seaborn as sns
from tqdm import tqdm

from mmengine.config import Config
from mmengine.runner import Runner
from mmengine.analysis import get_model_complexity_info
from mmseg.apis import init_model, inference_model
from mmseg.evaluation.metrics import IoUMetric
from mmseg.visualization import SegLocalVisualizer
from mmseg.registry import MODELS
from mmseg.structures import SegDataSample

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

# 自动选择设备
DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
logging.info(f"使用设备: {DEVICE}")

# --- 辅助函数 ---
def find_model_files(model_dir: str) -> Optional[Dict[str, 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


    # log_files = glob.glob(os.path.join(model_dir, 'vis_data', '*.json'))
    # if not log_files:
    #     logging.warning(f"在目录 {model_dir}/vis_data/ 中未找到日志文件 (.json)。")
    #     return None
    # # 通常只有一个json日志文件，取第一个
    # log_path = log_files

    # return {'config': config_path, 'checkpoint': checkpoint_path, 'log': log_path}
    return {'config': config_path, 'checkpoint': checkpoint_path}

def calculate_model_complexity(cfg: Config) -> Dict[str, float]:
    """
    计算模型的 FLOPs 和参数量。

    Args:
        cfg (Config): 加载的模型配置对象。

    Returns:
        Dict[str, float]: 包含 'flops_G' 和 'params_M' 的字典。
    """
    try:
        model = MODELS.build(cfg.model)
        model.eval()
        
        # 从配置中获取输入尺寸，如果不存在则使用默认值
        if hasattr(cfg, 'crop_size'):
            input_shape = (3, *cfg.crop_size)
        else:
            # 尝试从test_pipeline中查找
            input_shape = None
            if hasattr(cfg, 'test_pipeline'):
                for transform in cfg.test_pipeline:
                    if transform['type'] == 'Resize':
                        input_shape = (3, transform['scale'][1], transform['scale'])
                        break
            if input_shape is None:
                input_shape = (3, 512, 512)
                logging.info(f"配置中未找到 'crop_size' 或 'Resize'，使用默认输入尺寸 {input_shape} 进行 FLOPs 计算。")

        # 使用 MMEngine 的内置工具进行分析
        result = get_model_complexity_info(model, input_shape)
        
        flops_str = result.get('flops_str', '0.0G')
        params_str = result.get('params_str', '0.0M')
        
        # 解析字符串以获取数值
        flops_g = float(re.search(r'(\d+\.?\d*)', flops_str).group(1))
        params_m = float(re.search(r'(\d+\.?\d*)', params_str).group(1))
        
        # 单位转换
        if 'G' not in flops_str.upper():
            flops_g /= 1e3 # 假设是 MMac
        if 'M' not in params_str.upper():
            params_m /= 1e6 # 假设是 K

        return {'flops_G': flops_g, 'params_M': params_m}
    except Exception as e:
        logging.error(f"计算模型复杂度时出错: {e}")
        return {'flops_G': 0.0, 'params_M': 0.0}

def analyze_training_log(log_path: str, output_dir: str):
    """
    解析训练日志文件并绘制指标曲线。

    Args:
        log_path (str): 训练日志文件 (.json) 的路径。
        output_dir (str): 保存绘图的目录。
    """
    try:
        log_data = []
        with open(log_path, 'r') as f:
            for line in f:
                log_data.append(json.loads(line))

        df = pd.DataFrame(log_data)
        
        # 提取关键指标
        metrics_to_plot = {
            'mIoU': 'mIoU',
            'mAcc': 'mAcc',
            'aAcc': 'aAcc',
            'loss': 'loss'
        }
        
        # 筛选出验证和训练周期的行
        val_df = df[df['mode'] == 'val'].dropna(subset=['step'])
        train_df = df[df['mode'] == 'train'].dropna(subset=['step'])

        sns.set_theme(style="whitegrid")
        
        for key, name in metrics_to_plot.items():
            plt.figure(figsize=(12, 6))
            
            if key in val_df.columns:
                sns.lineplot(data=val_df, x='step', y=key, label=f'Validation {name}')
            
            if key in train_df.columns and key == 'loss':
                # 训练 loss 通常波动较大，可以进行平滑处理
                train_df[f'{key}_smooth'] = train_df[key].rolling(window=50, min_periods=1).mean()
                sns.lineplot(data=train_df, x='step', y=f'{key}_smooth', label=f'Training {name} (Smoothed)')

            plt.title(f'{name} Curve during Training')
            plt.xlabel('Training Steps')
            plt.ylabel(name)
            plt.legend()
            plt.tight_layout()
            save_path = os.path.join(output_dir, f'{name}_curve.png')
            plt.savefig(save_path)
            plt.close()
            logging.info(f"已保存日志图表: {save_path}")

    except Exception as e:
        logging.error(f"分析训练日志 {log_path} 时出错: {e}")

def run_inference_and_evaluation(cfg: Config, checkpoint_path: str, output_dir: str) -> Dict:
    """
    执行模型推理、评估和可视化。

    该函数通过一次性遍历验证数据集来高效地完成所有任务，避免了重复的数据加载和模型前向传播。

    Args:
        cfg (Config): 加载的模型配置对象。
        checkpoint_path (str): 模型检查点文件的路径。
        output_dir (str): 保存所有输出的根目录。

    Returns:
        Dict: 包含评估指标和逐类 TP/TN/FP/FN 计数的字典。
    """
    # --- 1. 初始化 ---
    model = init_model(cfg, checkpoint_path, device=DEVICE)
    
    # 以编程方式构建数据加载器，确保与训练时的数据预处理流程一致
    val_dataloader = Runner.build_dataloader(cfg.val_dataloader)
    
    # 实例化评估器和可视化器
    metric = IoUMetric(iou_metrics=['mIoU'])
    visualizer = SegLocalVisualizer(save_dir=os.path.join(output_dir, 'prediction_analysis'))
    
    # 关键步骤：为可视化器设置数据集元信息（类别名和调色板）
    visualizer.dataset_meta = val_dataloader.dataset.metainfo

    metric.dataset_meta = val_dataloader.dataset.metainfo
    
    # --- 2. 创建输出子目录 ---
    raw_mask_dir = os.path.join(output_dir, 'predicted_raw_masks')
    viz_dir = os.path.join(output_dir, 'prediction_analysis')
    os.makedirs(raw_mask_dir, exist_ok=True)
    os.makedirs(viz_dir, exist_ok=True)

    # --- 3. 推理循环 ---
    model.eval()
    metric.results = [] # 清空历史结果
    with torch.no_grad():
        for data in tqdm(val_dataloader, desc=f"推理与评估: {os.path.basename(output_dir)}"):
            # 将数据移动到指定设备
            # inputs = data['inputs'].to(DEVICE) # 有误
            inputs = data['inputs'][0].to(DEVICE).float().unsqueeze(0)
            data_samples =data['data_samples']
            
            # 执行推理
            result = model(inputs, data_samples=data_samples, mode='predict')
            
            # 将预测结果合并到 GT 样本中，以便于评估和可视化
            for i in range(len(data_samples)):
                data_samples[i].pred_sem_seg = result[i].pred_sem_seg
            
            # 使用评估器的 process 方法处理一个批次的结果
            predictions_as_dicts = [r.to_dict() for r in result]
            metric.process(data, predictions_as_dicts)
            
            # 保存原始预测掩码和可视化结果
            for sample in data_samples:
                pred_mask = sample.pred_sem_seg.data.squeeze().cpu().numpy().astype(np.uint8)
                img_filename = os.path.basename(sample.img_path)
                cv2.imwrite(os.path.join(raw_mask_dir, img_filename), pred_mask)
                
                # 生成并保存三图对比的可视化结果
                image = mmcv.imread(sample.img_path)
                visualizer.add_datasample(
                    name=os.path.splitext(img_filename),
                    image=image,
                    data_sample=sample,
                    show=False
                )

    # --- 4. 计算最终指标 ---
    metrics_summary = metric.compute_metrics(metric.results)
    
    # --- 5. 派生 TP, TN, FP, FN ---
    if not metric.results:
        logging.error("评估结果为空，无法计算 TP/TN/FP/FN。")
        return {'summary': metrics_summary, 'per_class': {}}

    # 聚合所有批次的结果
    total_area_intersect = torch.stack([res for res in metric.results]).sum(0)
    total_area_union = torch.stack([res[1] for res in metric.results]).sum(0)
    total_area_pred_label = torch.stack([res[2] for res in metric.results]).sum(0)
    total_area_label = torch.stack([res[3] for res in metric.results]).sum(0)

    # TP: 预测为正类，实际也为正类 (交集) [4]
    tp = total_area_intersect.numpy()
    # FP: 预测为正类，实际为负类 (预测区域 - 交集) [4]
    fp = (total_area_pred_label - total_area_intersect).numpy()
    # FN: 预测为负类，实际为正类 (真值区域 - 交集) [4]
    fn = (total_area_label - total_area_intersect).numpy()
    
    # TN: 预测为负类，实际也为负类
    # TN_i = total_valid_pixels - (TP_i + FP_i + FN_i)
    # total_valid_pixels = sum of all elements in confusion matrix = sum of all ground truth pixels
    total_valid_pixels = total_area_label.sum().item() * len(val_dataloader.dataset) / sum(len(b['inputs']) for b in val_dataloader) # 估算
    union = total_area_pred_label + total_area_label - total_area_intersect
    tn = total_valid_pixels - union.numpy()

    per_class_metrics = {'tp': tp, 'tn': tn, 'fp': fp, 'fn': fn}
    
    return {'summary': metrics_summary, 'per_class': per_class_metrics}

def consolidate_and_save_results(metrics_data: Dict, complexity_data: Dict, metainfo: Dict, output_dir: str):
    """
    整合所有量化指标并保存为 CSV 文件。

    Args:
        metrics_data (Dict): 来自 run_inference_and_evaluation 的指标数据。
        complexity_data (Dict): 来自 calculate_model_complexity 的复杂度数据。
        metainfo (Dict): 数据集的元信息，包含类别名。
        output_dir (str): 保存 CSV 文件的目录。
    """
    num_classes = len(metainfo['classes'])
    class_names = metainfo['classes']
    
    rows = []
    
    # 添加总体指标
    summary = metrics_data.get('summary', {})
    rows.append({'metric': 'iou_score', 'class_id': -1, 'class_name': 'Overall', 'value': summary.get('mIoU', 0)})
    rows.append({'metric': 'f1_score', 'class_id': -1, 'class_name': 'Overall', 'value': summary.get('mFscore', summary.get('mDice', 0))})
    rows.append({'metric': 'accuracy', 'class_id': -1, 'class_name': 'Overall', 'value': summary.get('aAcc', 0)})
    rows.append({'metric': 'recall', 'class_id': -1, 'class_name': 'Overall', 'value': summary.get('mRecall', 0)})
    rows.append({'metric': 'precision', 'class_id': -1, 'class_name': 'Overall', 'value': summary.get('mPrecision', 0)})
    
    # 添加模型复杂度
    rows.append({'metric': 'params_M', 'class_id': -1, 'class_name': 'N/A', 'value': complexity_data.get('params_M', 0)})
    rows.append({'metric': 'flops_G', 'class_id': -1, 'class_name': 'N/A', 'value': complexity_data.get('flops_G', 0)})
    
    # 添加逐类 TP/TN/FP/FN
    per_class = metrics_data.get('per_class', {})
    for metric_name, values in per_class.items():
        if len(values) == num_classes:
            for i in range(num_classes):
                rows.append({
                    'metric': metric_name,
                    'class_id': i,
                    'class_name': class_names[i],
                    'value': values[i]
                })

    df = pd.DataFrame(rows)
    save_path = os.path.join(output_dir, 'prediction_analysis', 'test_set_metrics.csv')
    df.to_csv(save_path, index=False)
    logging.info(f"所有量化指标已保存至: {save_path}")

# --- 主函数 ---
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("无效的数据集选择，程序已退出。")
    
    # --- 交互式选择修改结束 ---

    # 修改后的循环，将遍历经过用户筛选后的 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")
        
        final_output_dir = os.path.join(output_root, output_dataset_folder, model_name)
        os.makedirs(final_output_dir, exist_ok=True)
        logging.info(f"结果将保存至: {final_output_dir}")

        try:
            # 加载配置
            cfg = Config.fromfile(files['config'])
            
            # 1. 执行推理、评估和可视化
            metrics_data = run_inference_and_evaluation(cfg, files['checkpoint'], final_output_dir)
            
            # 2. 分析训练日志
            analyze_training_log(files['log'], os.path.join(final_output_dir, 'prediction_analysis'))
            
            # 3. 计算模型复杂度
            complexity_data = calculate_model_complexity(cfg)
            
            # 4. 整合并保存所有量化结果
            val_dataloader = Runner.build_dataloader(cfg.val_dataloader)
            metainfo = val_dataloader.dataset.metainfo
            
            consolidate_and_save_results(metrics_data, complexity_data, metainfo, final_output_dir)

            logging.info(f"--- 模型 {model_name} 处理完成 ---")

        except Exception as e:
            logging.error(f"处理模型 {model_name} 时发生严重错误: {e}", exc_info=True)


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)