import logging, argparse, sys, utils
import shutil, tempfile
from pathlib import Path
from typing import Dict, List, Tuple

import albumentations as A
import cv2
import matplotlib.pyplot as plt
import numpy as np
import segmentation_models_pytorch as smp
import torch
from albumentations.pytorch import ToTensorV2
from scipy.special import softmax
from sklearn.metrics import (auc, average_precision_score,
                               precision_recall_curve, roc_curve)
from tqdm import tqdm

# 本地应用/库的导入
import config
from utils import log_metrics_to_csv, get_preprocessing_transform, setup_directories

# --- 日志设置 ---
# 配置日志记录器
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[logging.StreamHandler()] # 输出日志到控制台
)
# 为 matplotlib 设置中文字体，请根据您的系统环境选择合适的字体
plt.rcParams['font.sans-serif'] = ['Noto Sans CJK JP', 'SimHei', 'Microsoft YaHei']  # 例如：SimHei, Microsoft YaHei
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# --- 核心工具函数 ---

# 将单通道的类别索引掩码转换为 RGB 彩色掩码
def colorize_mask(mask: np.ndarray, class_rgb_values: List[Tuple[int, int, int]]) -> np.ndarray:
    """将单通道的类别索引掩码转换为 RGB 彩色掩码。"""
    color_mask = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
    # 一个通用的调色板，可以轻松扩展
    palette = [
        (0, 0, 0), (220, 20, 60), (0, 255, 0), (0, 0, 255),
        (255, 255, 0), (255, 0, 255), (0, 255, 255), (244, 164, 96),
    ]
    for class_idx, color in enumerate(palette):
        if class_idx < len(class_rgb_values):
            # 假设掩码中的像素值 [0, 1, 2...] 对应类别的索引
            pixels_to_color = (mask == class_idx) 
            color_mask[pixels_to_color] = color
    return color_mask

# --- 分析与可视化函数 ---

def generate_and_save_curves(y_true: np.ndarray, y_probs: np.ndarray, class_names: List[str], output_dir: Path) -> None:
    """
    根据整个测试集的聚合预测结果，为每个类别生成并保存 ROC 和 PR 曲线。
    """
    logging.info("正在生成 ROC 和 PR 曲线...")
    num_classes = len(class_names)
    plt.style.use('seaborn-v0_8-whitegrid')

    # --- 准备画布 ---
    fig_roc, ax_roc = plt.subplots(figsize=(10, 8))
    ax_roc.plot([0, 1], [0, 1], 'k--', label='随机猜测')
    fig_pr, ax_pr = plt.subplots(figsize=(10, 8))

    # --- 为每个类别（跳过背景）生成曲线 ---
    for i in range(1, num_classes):
        class_name = class_names[i]
        y_true_binary = (y_true == i).astype(int)
        y_class_probs = y_probs[:, i]

        # ROC 曲线
        fpr, tpr, _ = roc_curve(y_true_binary, y_class_probs)
        roc_auc = auc(fpr, tpr)
        ax_roc.plot(fpr, tpr, label=f'类别 {class_name} (AUC = {roc_auc:.4f})')

        # PR 曲线
        precision, recall, _ = precision_recall_curve(y_true_binary, y_class_probs)
        avg_precision = average_precision_score(y_true_binary, y_class_probs)
        ax_pr.plot(recall, precision, label=f'类别 {class_name} (AP = {avg_precision:.4f})')

    # --- 美化并保存 ROC 曲线图 ---
    ax_roc.set(xlabel='False positive rate (FPR)', ylabel='True positive rate (TPR)', title='Receiver operating characteristic (ROC) curve')
    ax_roc.legend(loc='lower right')
    ax_roc.set_aspect('equal', adjustable='box')
    fig_roc.tight_layout()
    roc_save_path = output_dir / "ROC_Curves.png"
    fig_roc.savefig(roc_save_path)
    plt.close(fig_roc)
    logging.info(f"-> ROC 曲线图已保存至: {roc_save_path}")

    # --- 美化并保存 PR 曲线图 ---
    ax_pr.set(xlabel='Recall', ylabel='Precision', title='Precision-Recall (PR) Curve')
    ax_pr.legend(loc='lower left')
    ax_pr.set_aspect('equal', adjustable='box')
    fig_pr.tight_layout()
    pr_save_path = output_dir / "Precision_Recall_Curves.png"
    fig_pr.savefig(pr_save_path)
    plt.close(fig_pr)
    logging.info(f"-> PR 曲线图已保存至: {pr_save_path}")

def save_visual_comparison(
    image_name: str, original_image: np.ndarray, pred_mask: np.ndarray,
    gt_mask: np.ndarray = None, stats: Dict = None, stats_text: str = "",
    save_dir: str = ""
) -> None:
    """保存原始图像、预测掩码以及可选的真实掩码的对比图。"""
    pred_color_mask = colorize_mask(pred_mask, config.CLASS_RGB_VALUES)
    
    if gt_mask is not None and stats is not None:
        num_plots = 3
        figsize = (22, 8)
        gt_color_mask = colorize_mask(gt_mask, config.CLASS_RGB_VALUES)
        save_name = f"{Path(image_name).stem}-IOU_{stats['iou']:.4f}-ACC_{stats['acc']:.4f}.png"
    else:
        num_plots = 2
        figsize = (14, 6)
        save_name = f"{Path(image_name).stem}.png"

    fig, axes = plt.subplots(1, num_plots, figsize=figsize)
    fig.suptitle(f"Predict: {image_name}", fontsize=16)
    
    axes[0].imshow(original_image)
    axes[0].set_title("ORI_IMG")
    axes[0].axis('off')

    pred_title = "Predicted Mask"
    if stats:
        pred_title += f"\nIoU: {stats['iou']:.4f} | Acc: {stats['acc']:.4f}"
    axes[1].imshow(pred_color_mask)
    axes[1].set_title(pred_title)
    axes[1].axis('off')

    if num_plots == 3:
        axes[2].imshow(gt_color_mask)
        axes[2].set_title("Ground Truth")
        axes[2].axis('off')
        plt.figtext(0.5, 0.01, stats_text, ha="center", fontsize=10, 
                    bbox={"facecolor":"white", "alpha":0.7, "pad":5}, family='monospace')
        fig.subplots_adjust(bottom=0.3)
    
    plt.tight_layout()
    plt.savefig(save_dir / save_name)
    plt.close(fig)

# --- 主要逻辑函数 ---

def load_model(model_architecture: str, model_save_path: Path) -> torch.nn.Module:
    """
    加载已训练的分割模型。
    此函数经过优化，可以智能处理由 torch.nn.DataParallel 保存的
    (带有 'module.' 前缀) 和常规保存的模型权重。
    """
    logging.info(f"正在从以下路径加载模型: {model_save_path}")
    num_classes = len(config.CLASSES)
    
    # 1. 【修改】根据 config 动态创建基础模型架构
    logging.info(f"正在重建模型架构: '{model_architecture}'")
    try:
        model_params = config.ALL_MODEL_CONFIGS[model_architecture].copy() #
        model_class = getattr(smp, model_architecture)
    except AttributeError:
        logging.error(f"模型 '{model_architecture}' 在 segmentation_models_pytorch 库中不存在！")
        raise

    # 2. 【修改】准备参数，与训练时保持一致
    #    注意：加载权重时，我们不希望再次下载预训练权重，所以设为 None
    params = model_params
    params['in_channels'] = 3
    params['classes'] = num_classes
    params['encoder_weights'] = None  # 非常重要！避免在预测时重新下载ImageNet权重

    # ======================== 【新增代码段开始】 ======================== #
    # 自动检测 encoder_name 是否包含 'vit'
    encoder_name = model_params.get('encoder_name', '')
    if 'vit' in encoder_name.lower():
        params['dynamic_img_size'] = True
        logging.info(f"检测到 ViT 编码器 ('{encoder_name}')。自动设置 dynamic_img_size=True。")
    # ======================== 【新增代码段结束】 ======================== #

    # 3. 【修改】使用 **kwargs 创建模型实例
    model = model_class(**params).to(config.DEVICE)

    # 4. 从文件加载 state_dict
    state_dict = torch.load(model_save_path, map_location=torch.device(config.DEVICE))

    # 3. 检查是否存在 'module.' 前缀 (判断是否为 DataParallel 保存的权重)
    #    我们通过检查字典中任何一个键是否以 'module.' 开头来判断
    is_dataparallel = any(key.startswith('module.') for key in state_dict.keys())
    
    if is_dataparallel:
        logging.info("检测到模型由 DataParallel 保存，将移除 'module.' 前缀。")
        # 创建一个新的、不带前缀的 state_dict
        new_state_dict = {key.replace('module.', ''): value for key, value in state_dict.items()}
        model.load_state_dict(new_state_dict)
    else:
        # 如果没有 'module.' 前缀，直接加载
        logging.info("直接加载标准模型权重。")
        model.load_state_dict(state_dict)
    
    model.eval()
    return model

def calculate_and_save_final_metrics(stats_list: List[Dict], num_classes: int, save_dir: str) -> None:
    """聚合每张图像的统计数据，计算总体指标，并保存到 CSV 文件。"""
    if not stats_list:
        logging.warning("未收集到统计数据，跳过最终指标计算。")
        return

    logging.info("\n正在聚合结果以计算最终指标...")
    # 堆叠所有统计张量并按批次维度求和
    tp = torch.cat([s['tp'] for s in stats_list], dim=0).sum(dim=0)  # (N×B, C) -> (C,)
    fp = torch.cat([s['fp'] for s in stats_list], dim=0).sum(dim=0)  # (N×B, C) -> (C,)
    fn = torch.cat([s['fn'] for s in stats_list], dim=0).sum(dim=0)  # (N×B, C) -> (C,)
    tn = torch.cat([s['tn'] for s in stats_list], dim=0).sum(dim=0)  # (N×B, C) -> (C,)

    # ======================== 【新增代码段开始】 ======================== #
    # 根据 config.py 中的设置，筛选出用于计算宏观指标的类别
    ignore_indices = [config.CLASSES.index(cls) for cls in config.EVALUATION_CLASSES_TO_IGNORE if cls in config.CLASSES]
    
    keep_mask = torch.ones(num_classes, dtype=torch.bool, device=tp.device)
    if ignore_indices:
        keep_mask[ignore_indices] = False
        logging.info(f"最终评估时将忽略类别: {config.EVALUATION_CLASSES_TO_IGNORE}")

    tp_filtered = tp[keep_mask]
    fp_filtered = fp[keep_mask]
    fn_filtered = fn[keep_mask]
    tn_filtered = tn[keep_mask]
    # ======================== 【新增代码段结束】 ======================== #

    
    # ======================== 【修改代码段开始】 ======================== #
    # 使用过滤后的统计数据计算总体指标
    metrics = {
        "iou_score": smp.metrics.iou_score(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
        "f1_score": smp.metrics.f1_score(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
        "accuracy": smp.metrics.accuracy(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
        "recall": smp.metrics.recall(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
        "precision": smp.metrics.precision(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
    }
    # ======================== 【修改代码段结束】 ======================== #
    
    # 添加每个类别的原始统计数据（使用未经过滤的数据）
    for i in range(num_classes):
        metrics[f'tp_class_{i}'] = tp[i].item()
        metrics[f'tn_class_{i}'] = tn[i].item()
        metrics[f'fp_class_{i}'] = fp[i].item()
        metrics[f'fn_class_{i}'] = fn[i].item()

    csv_save_path = save_dir / "test_set_metrics.csv"
    log_metrics_to_csv(metrics, str(csv_save_path))
    logging.info(f"-> 整体测试集指标已保存至: {csv_save_path}")

# 2. 【新增】一个帮助函数，用于交互式选择模型目录
def select_trained_model_path(base_dir: Path, model_architecture: str) -> Path:
    """
    查找指定架构的所有训练运行目录，并让用户选择一个。
    """
    logging.info(f"正在 '{base_dir}' 中搜索模型 '{model_architecture}' 的训练记录...")
    
    # 查找所有以架构名开头的文件夹
    run_dirs = sorted([d for d in base_dir.iterdir() if d.is_dir() and d.name.startswith(model_architecture+'_')])
    
    if not run_dirs:
        logging.error(f"未找到任何 '{model_architecture}' 模型的训练记录。请先运行 train.py。")
        return None

    print("\n请选择要用于预测的模型:")
    for i, dir_path in enumerate(run_dirs):
        print(f"  [{i+1}] {dir_path.name}")
    
    while True:
        try:
            choice = input(f"请输入选项编号 (1-{len(run_dirs)}) 或按 Enter 退出: ")
            if not choice:
                return None
            choice_idx = int(choice) - 1
            if 0 <= choice_idx < len(run_dirs):
                selected_dir = run_dirs[choice_idx]
                logging.info(f"已选择模型: {selected_dir}")
                return selected_dir
            else:
                print("无效的选项，请重新输入。")
        except (ValueError, IndexError):
            print("无效的输入，请输入一个数字。")

# 【修改】函数签名和内部逻辑，不再接收一个列表，而是直接接收聚合后的统计数据
def calculate_and_save_final_metrics(tp, fp, fn, tn, num_classes: int, save_dir: Path) -> None:
    """根据聚合后的统计数据，计算总体指标，并保存到 CSV 文件。"""
    logging.info("\n正在计算最终的整体指标...")
    
    # 根据 config.py 中的设置，筛选出用于计算宏观指标的类别
    ignore_indices = [config.CLASSES.index(cls) for cls in config.EVALUATION_CLASSES_TO_IGNORE if cls in config.CLASSES]
    
    keep_mask = torch.ones(num_classes, dtype=torch.bool, device=tp.device)
    if ignore_indices:
        keep_mask[ignore_indices] = False
        logging.info(f"最终评估时将忽略类别: {config.EVALUATION_CLASSES_TO_IGNORE}")

    tp_filtered = tp[keep_mask]
    fp_filtered = fp[keep_mask]
    fn_filtered = fn[keep_mask]
    tn_filtered = tn[keep_mask]

    # 使用过滤后的统计数据计算总体指标
    metrics = {
        "iou_score": smp.metrics.iou_score(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
        "f1_score": smp.metrics.f1_score(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
        "accuracy": smp.metrics.accuracy(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
        "recall": smp.metrics.recall(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
        "precision": smp.metrics.precision(tp_filtered, fp_filtered, fn_filtered, tn_filtered, reduction='micro').item(),
    }
    
    # 添加每个类别的原始统计数据（使用未经过滤的数据）
    for i in range(num_classes):
        metrics[f'tp_class_{i}'] = tp[i].item()
        metrics[f'tn_class_{i}'] = tn[i].item()
        metrics[f'fp_class_{i}'] = fp[i].item()
        metrics[f'fn_class_{i}'] = fn[i].item()

    csv_save_path = save_dir / "test_set_metrics.csv"
    log_metrics_to_csv(metrics, str(csv_save_path))
    logging.info(f"-> 整体测试集指标已保存至: {csv_save_path}")

def main(model_architecture: str) -> None:
    """
    主函数，执行模型加载、图像预测、评估，并生成分析结果。
    """
    # --- 1. 【修改】选择模型并动态设置路径 ---
    selected_run_dir = select_trained_model_path(Path(config.PREDICT_BEST_MODEL_DIR), model_architecture)

    if not selected_run_dir:
        logging.info("未选择任何模型，程序退出。")
        sys.exit()

    # 【重要】动态覆写 config 中的路径，使其指向所选模型的目录
    # 这样，后续所有代码都会自动在正确的文件夹中加载模型和保存结果
    best_model_save_path = selected_run_dir / "best_model.pth"
    raw_mask_dir = selected_run_dir / "predicted_raw_masks"
    analysis_results_dir = selected_run_dir / "prediction_analysis"
    
    if not best_model_save_path.exists():
        logging.error(f"错误: 在 '{selected_run_dir}' 中未找到 'best_model.pth' 文件！")
        sys.exit()

    # setup_directories 函数现在用于创建预测结果的子目录
    utils.setup_directories(raw_mask_dir)
    utils.setup_directories(analysis_results_dir)

    model = load_model(model_architecture, best_model_save_path) # 1. 加载模型

    # ======================== 【新增/修改的代码段开始】 ======================== #
    # --- 2. 探测模型以确定正确的图像尺寸 ---
    logging.info("正在探测模型以确定预测时所需的输入尺寸...") #
    target_height, target_width = config.IMAGE_HEIGHT, config.IMAGE_WIDTH #

    try:
        # 检查编码器是否需要固定输入尺寸
        if model.encoder.is_fixed_input_size:
            required_size = model.encoder.input_size
            target_height = required_size[1]
            target_width = required_size[2]
            logging.warning(
                f"模型 '{model_architecture}' 的编码器需要固定的输入尺寸 "
                f"({target_height}x{target_width})。"
                f"将忽略 config.py 中的尺寸设置进行预测。"
            )
        else:
             logging.info(f"模型 '{model_architecture}' 支持动态输入尺寸。将使用 config.py 中定义的尺寸 ({target_height}x{target_width})。")
    except AttributeError:
        logging.info("无法自动检测模型尺寸要求。将使用 config.py 中定义的尺寸。")

    # --- 3. 获取用于预测的图像预处理流程 ---
    # 使用探测到的或配置中指定的目标尺寸
    transform = get_preprocessing_transform(target_height, target_width) #
    # ======================== 【新增/修改的代码段结束】 ======================== #

    # --- 4. 获取图像列表并判断模式 ---
    image_filenames = sorted(list(config.TEST_IMAGE_DIR.iterdir()))
    evaluate_mode = config.TEST_MASK_DIR.is_dir()

    if evaluate_mode:
        logging.info("正在以 [评估模式] 运行。将使用真实掩码进行评估。")
    else:
        logging.info("正在以 [仅预测模式] 运行。未找到真实掩码。")

    # --- 3. 初始化用于增量计算的变量 ---
    # 【修改】不再使用 list 来累积，而是直接累加 TP, FP, FN, TN
    num_classes = len(config.CLASSES)
    total_tp = torch.zeros(num_classes, dtype=torch.long, device=config.DEVICE)
    total_fp = torch.zeros(num_classes, dtype=torch.long, device=config.DEVICE)
    total_fn = torch.zeros(num_classes, dtype=torch.long, device=config.DEVICE)
    total_tn = torch.zeros(num_classes, dtype=torch.long, device=config.DEVICE)
    
    # 【新增】为 ROC/PR 曲线创建临时文件，用磁盘空间换内存
    # 使用 tempfile 模块来安全地创建临时文件
    temp_gt_file = tempfile.NamedTemporaryFile(delete=False, suffix='.npy')
    temp_probs_file = tempfile.NamedTemporaryFile(delete=False, suffix='.npy')
    logging.info(f"为ROC/PR曲线创建临时文件: {temp_gt_file.name}, {temp_probs_file.name}")

    # --- 4. 主推理循环 - 预测图片 ---
    for image_path in tqdm(image_filenames, desc="正在处理测试图像"):
        # --- 4.1. 读取图片并预处理 ---
        original_image = cv2.cvtColor(cv2.imread(str(image_path)), cv2.COLOR_BGR2RGB)
        original_h, original_w = original_image.shape[:2]
        image_tensor = transform(image=original_image)['image'].unsqueeze(0).to(config.DEVICE)

        # --- 4.2. 图片推理并保存 ---
        with torch.no_grad():
            pred_logits_tensor = model(image_tensor)
        pred_mask_tensor = torch.argmax(pred_logits_tensor, dim=1)
        pred_mask_numpy = pred_mask_tensor.squeeze().cpu().numpy().astype(np.uint8)
        pred_mask_resized_raw = cv2.resize(pred_mask_numpy, (original_w, original_h), interpolation=cv2.INTER_NEAREST)
        
        # 保存原始（单通道）的预测掩码
        cv2.imwrite(str(raw_mask_dir / f"{image_path.stem}.png"), pred_mask_resized_raw)
        
        # --- 4.3. 评估与可视化 ---
        gt_mask_numpy, per_image_stats, stats_text_display, gt_mask_raw = None, None, "", None
        if evaluate_mode:
            mask_path = config.TEST_MASK_DIR / image_path.name
            if mask_path.exists():
                gt_mask_raw = cv2.imread(str(mask_path), cv2.IMREAD_GRAYSCALE)
                gt_mask_numpy = cv2.resize(
                    gt_mask_raw, (target_width, target_height), 
                    interpolation=cv2.INTER_NEAREST
                )
                gt_mask_tensor = torch.from_numpy(gt_mask_numpy).long().to(config.DEVICE).unsqueeze(0)

                # ======================== 【修复代码段开始】 ======================== #

                # --- 1. [修复] 计算单张图像的统计数据 ---
                # 这个代码块是必需的，它为下面的 per_image_stats 和 stats_text_display 定义了 tp, fp, fn, tn
                tp, fp, fn, tn = smp.metrics.get_stats(
                    pred_mask_tensor, gt_mask_tensor, 
                    mode=config.SEG_MODE, num_classes=num_classes
                )
                
                # --- 2. 累加到总数，用于最终的整体评估 ---
                total_tp += tp.squeeze(0).to(config.DEVICE)
                total_fp += fp.squeeze(0).to(config.DEVICE)
                total_fn += fn.squeeze(0).to(config.DEVICE)
                total_tn += tn.squeeze(0).to(config.DEVICE)
                
                # --- 3. 将曲线数据写入临时文件 (这部分逻辑正确，无需修改) ---
                probs_tensor = torch.softmax(pred_logits_tensor, dim=1)
                gt_flat = gt_mask_tensor.cpu().numpy().flatten()
                probs_flat = probs_tensor.permute(0, 2, 3, 1).reshape(-1, num_classes).cpu().numpy()
                with open(temp_gt_file.name, 'ab') as f_gt, open(temp_probs_file.name, 'ab') as f_probs:
                    np.save(f_gt, gt_flat)
                    np.save(f_probs, probs_flat)

                # --- 4. [修复] 现在可以安全地计算单图指标用于可视化 ---
                # 因为上面的代码块已经定义了 tp, fp, fn, tn，所以这里不再报错
                per_image_stats = {
                    'iou': smp.metrics.iou_score(tp, fp, fn, tn, reduction='micro').item(),
                    'acc': smp.metrics.accuracy(tp, fp, fn, tn, reduction='micro').item()
                }

                # --- 5. [修复] 修正生成统计文本时的 Bug ---
                # `get_stats` 返回的 tp 张量形状是 (1, num_classes), 
                # 因此我们应该遍历第1维度 (类别)，并使用 [0, i] 进行索引
                stats_text_display = "Each Class TP / FP / FN / TN:\n" + "-"*35 + "\n"
                for i in range(1, tp.shape[1]): # 修正: 遍历类别维度 tp.shape[1]
                    if i < len(config.CLASSES):
                        class_name = config.CLASSES[i]
                        stats_text_display += (
                            f"  {class_name:<10}: "
                            # 修正: 使用正确的索引 tp[0, i]
                            f"{tp[0, i]:>5} / {fp[0, i]:>5} / {fn[0, i]:>5} / {tn[0, i]:>7}\n"
                        )
        
        save_visual_comparison(
            image_name=image_path.name,
            original_image=original_image,
            pred_mask=pred_mask_resized_raw,
            gt_mask=gt_mask_raw,
            stats=per_image_stats,
            stats_text=stats_text_display,
            save_dir = analysis_results_dir
        )

    logging.info("\n预测流程已完成！")
    logging.info(f"原始预测掩码已保存至: {raw_mask_dir}")
    logging.info(f"分析图像已保存至: {analysis_results_dir}")

    # --- 5. 最终分析（循环结束后） ---
    if evaluate_mode:
        # calculate_and_save_final_metrics(all_stats_for_metrics, len(config.CLASSES), save_dir = analysis_results_dir)
        calculate_and_save_final_metrics(total_tp.cpu(), total_fp.cpu(), total_fn.cpu(), total_tn.cpu(), num_classes, save_dir=analysis_results_dir)

        # if all_gt_for_curves:
        #     logging.info("\n正在聚合结果以生成 ROC 和 PR 曲线...")
        #     full_gt = torch.cat(all_gt_for_curves).numpy().flatten()
        #     full_logits = torch.cat(all_logits_for_curves)
            
        #     # 调整维度并计算概率
        #     num_classes = len(config.CLASSES)
        #     full_probs = softmax(
        #         full_logits.permute(0, 2, 3, 1).reshape(-1, num_classes).numpy(), 
        #         axis=1
        #     )
            
        #     generate_and_save_curves(
        #         y_true=full_gt, y_probs=full_probs,
        #         class_names=config.CLASSES, output_dir=analysis_results_dir
        #     )

        # 【新增】从临时文件中分块读取数据来生成曲线
        try:
            logging.info("\n正在从临时文件加载数据以生成 ROC 和 PR 曲线...")
            temp_gt_file.close()
            temp_probs_file.close()

            all_gt_parts, all_probs_parts = [], []
            with open(temp_gt_file.name, 'rb') as f_gt, open(temp_probs_file.name, 'rb') as f_probs:
                while True:
                    try:
                        all_gt_parts.append(np.load(f_gt))
                        all_probs_parts.append(np.load(f_probs))
                    # 【修复】将 EOFError 添加到要捕获的异常列表中
                    except (IOError, ValueError, EOFError):
                        break # 文件读取完毕
            
            if all_gt_parts:
                full_gt = np.concatenate(all_gt_parts)
                full_probs = np.concatenate(all_probs_parts)
                
                generate_and_save_curves(
                    y_true=full_gt, y_probs=full_probs,
                    class_names=config.CLASSES, output_dir=analysis_results_dir
                )
            else:
                logging.warning("临时文件中没有数据，跳过 ROC/PR 曲线生成。")
                
        finally:
            # 【新增】确保在最后删除临时文件
            import os
            os.remove(temp_gt_file.name)
            os.remove(temp_probs_file.name)
            logging.info("已清理临时文件。")

if __name__ == "__main__":
        # 创建一个参数解析器
    parser = argparse.ArgumentParser(description="训练图像分割模型。")

    # 添加 --architecture 参数
    # - `choices` 会自动从你的 config 文件中获取所有可用的模型名称
    # - `default` 设置了在不提供参数时的默认模型
    parser.add_argument(
        "-a", "--architecture",
        type=str,
        # default='Unet',
        choices=list(config.ALL_MODEL_CONFIGS.keys()),
        required=True,
        help="选择要训练的模型架构。"
    )

    # 解析命令行传入的参数
    args = parser.parse_args()

    # 将解析出的架构名称传入 main 函数并执行
    main(model_architecture=args.architecture)