import logging, argparse, shutil
from datetime import datetime
from pathlib import Path
from typing import Dict, Tuple
import gc

import numpy as np
import segmentation_models_pytorch as smp
import torch
import torch.optim as optim
from torch.optim import lr_scheduler # 学习率调度器
from torch.amp import GradScaler
from torch.utils.data import DataLoader

# 本地应用/库的导入
import config, os
import utils
from dataset import SegmentationDataset
from loss import MultiClassLoss, UNetPlusPlusLoss

# --- 日志设置 ---
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    handlers=[
        logging.StreamHandler() # 输出日志到控制台
    ]
)

# --- 辅助函数：移动训练结果文件夹从源输出目录移动到配置的硬盘目录 ---
def move_results_to_hardisk(project_folder: str):
    """
    将指定的训练结果文件夹从源输出目录移动到配置的硬盘目录。
    移动操作包括复制整个文件夹树，然后在复制成功后删除原始文件夹。

    参数:
        project_folder (str): 要移动的项目文件夹的名称。
                              该名称通常由模型架构和时间戳构成。
    """
    source_dir = Path(config.OUTPUTS_DIR) / project_folder
    # 确保目标硬盘目录存在
    outputs_folder_name = Path(config.OUTPUTS_DIR).name
    destination_dir = Path(config.HARDISK_DIR) / outputs_folder_name / project_folder
    destination_dir.parent.mkdir(parents=True, exist_ok=True)
    
    logging.info(f"准备将结果从 {source_dir} 移动到 {destination_dir}...")
    
    try:
        # 步骤 1: 复制文件夹
        shutil.copytree(source_dir, destination_dir)
        logging.info(f"成功复制结果到: {destination_dir}")

        # 步骤 2: 复制成功后，删除原文件夹
        logging.info(f"正在删除原文件夹: {source_dir}")
        shutil.rmtree(source_dir)
        logging.info("成功删除原文件夹。")
        logging.info(f"任务完成，最终结果已保存至: {destination_dir}")

    except Exception as e:
        logging.error(f"移动文件夹时发生错误: {e}")
        logging.warning(f"原始训练结果仍保留在: {source_dir}")

# 1. 根据配置文件初始化模型、损失函数和优化器。
def initialize_components(model_architecture:str, num_classes: int, seg_mode: str) -> Tuple:
    """
    根据配置文件动态初始化模型、损失函数和优化器。
    """
    # --- 初始化模型 ---
    logging.info(f"正在初始化模型: '{model_architecture}'...")

    # 使用传入的 model_architecture 作为 key 来获取参数
    try:
        model_params = config.ALL_MODEL_CONFIGS[model_architecture]
        model_class = getattr(smp, model_architecture)
    except KeyError:
        logging.error(f"模型 '{model_architecture}' 的配置未在 config.py 的 ALL_MODEL_CONFIGS 中定义！")
        raise
    except AttributeError:
        logging.error(f"模型 '{model_architecture}' 在 segmentation_models_pytorch 库中不存在！")
        raise

    # 2. 准备参数字典
    #    首先复制 config 中的参数，然后添加固定的 `in_channels` 和 `classes`
    params = model_params.copy()
    params['in_channels'] = 3
    params['classes'] = num_classes

    # ======================== 【新增代码段开始】 ======================== #
    # 自动检测 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. 使用字典解包 (**) 将所有参数传递给模型构造函数
    #    这种方式非常灵活，无论你在 config 中定义了多少参数，都能正确传递
    model = model_class(**params).to(config.DEVICE)
    
    # --- 设置损失函数 (保持不变) ---
    logging.info(f"为 '{seg_mode}' 模式设置损失函数。")
    if seg_mode == 'multiclass':
        loss_fn = MultiClassLoss(mode=seg_mode)
    elif seg_mode == 'multilabel':
        loss_fn = UNetPlusPlusLoss(mode=seg_mode)
    else:
        raise ValueError(f"无效的 SEG_MODE: '{seg_mode}'。必须是 'multiclass' 或 'multilabel'。")

    # --- 优化器和梯度缩放器 (保持不变) ---
    optimizer = optim.Adam(model.parameters(), lr=config.LEARNING_RATE, weight_decay=config.WEIGHT_DECAY)
    # --- TODO 添加学习率调度器 TODO ---
    # T_max 是调度器周期的最大迭代次数，通常设置为总的 epoch 数量
    scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=config.EPOCHS, eta_min=1e-6)
    # -----------------------
    scaler = GradScaler('cuda')
    
    return model, loss_fn, optimizer, scaler

def main(model_architecture: str) -> None:
    """
    编排图像分割模型的主要训练流程。

    该函数执行以下步骤:
    1. 设置输出目录。
    2. 准备数据增强，并为训练集和验证集创建 DataLoader 实例。
    3. 初始化 U-Net++ 模型、损失函数、优化器和梯度缩放器。
    4. 运行主训练循环，其中包括:
       - 训练一个 epoch。
       - 在验证集上评估模型。
       - 将各项指标记录到 CSV 文件。
       - 基于验证损失实现早停机制。
       - 保存性能最佳的模型和定期的断点。
       - 绘制训练进度曲线图。
    返回:
        str: 本次训练运行的文件夹名称 (run_name)。
    """
    logging.info(f"使用设备: {config.DEVICE}")
    
    # --- 1.1. 使用 pathlib 进行现代化的路径管理 ---
    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    run_name = f"{model_architecture}_{timestamp}"
    output_dir = Path(config.OUTPUTS_DIR) / run_name # <-- 使用 / 操作符
    metrics_csv_path = output_dir / config.METRICS_CSV_NAME
    best_model_path = output_dir / config.BEST_MODEL_SAVE_NAME

    utils.setup_directories(output_dir)

    # --- 1.2. 探测模型以确定最佳图像尺寸 ---
    logging.info("正在探测模型以确定输入尺寸要求...") #
    num_classes = len(config.CLASSES) #
    
    # 创建一个临时模型实例，仅用于检查其属性
    # 注意：这里的调用现在更简单，不需要传递 dynamic_img_size
    probe_model, _, _, _ = initialize_components(
        model_architecture, num_classes, config.SEG_MODE #
    )
    
    target_height, target_width = config.IMAGE_HEIGHT, config.IMAGE_WIDTH #

    try:
        # 检查编码器是否需要固定输入尺寸
        if probe_model.encoder.is_fixed_input_size:
            required_size = probe_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 中定义的尺寸。")

    del probe_model # 删除临时模型，释放内存

    # --- 1.3. 数据加载 ---
    logging.info(f"正在设置数据增强和加载器，目标尺寸为 {target_height}x{target_width}...") #
    # 使用探测到的或配置中指定的目标尺寸
    train_transform = utils.get_training_augmentation(target_height, target_width) #
    val_transform = utils.get_validation_augmentation(target_height, target_width) #

    train_dataset = SegmentationDataset(
        image_dir=Path(config.TRAIN_IMAGE_DIR),
        mask_dir=Path(config.TRAIN_MASK_DIR),
        classes=config.CLASSES,
        class_rgb_values=config.CLASS_RGB_VALUES,
        augmentation=train_transform,
        seg_mode=config.SEG_MODE
    )
    val_dataset = SegmentationDataset(
        # 已修正路径，使用配置文件中的验证数据目录
        image_dir=Path(config.VAL_IMAGE_DIR),
        mask_dir=Path(config.VAL_MASK_DIR),
        classes=config.CLASSES,
        class_rgb_values=config.CLASS_RGB_VALUES,
        augmentation=val_transform,
        seg_mode=config.SEG_MODE
    )

    train_loader = DataLoader(train_dataset, batch_size=config.BATCH_SIZE, shuffle=True, pin_memory=config.PIN_MEMORY, num_workers=config.NUM_WORKERS)
    val_loader = DataLoader(val_dataset, batch_size=config.BATCH_SIZE, shuffle=False, pin_memory=config.PIN_MEMORY, num_workers=config.NUM_WORKERS)
    
    # --- 2. 模型、损失函数、优化器、梯度缩放器 ---
    num_classes = len(config.CLASSES)
    model, loss_fn, optimizer, scaler = initialize_components(model_architecture, num_classes, config.SEG_MODE)
    
    # --- 启用多 GPU 训练 ---
    if torch.cuda.device_count() >= 1:
        logging.info(f"正在使用 {torch.cuda.device_count()} 块 GPU 进行并行训练。")
        model = torch.nn.DataParallel(model)

    # --- 3. 训练循环 ---
    best_val_loss = float('inf')
    epochs_no_improve = 0
    
    for epoch in range(config.EPOCHS):
        logging.info(f"\n--- 第 {epoch+1}/{config.EPOCHS} 轮 ---")
        # --- 3.1. 训练集训练数据 ---
        train_loss = utils.train_fn(train_loader, model, optimizer, loss_fn, scaler, config.DEVICE, model_architecture)
        # --- 3.2. 验证集测试数据【获得指标】 ---
        val_metrics_dict = utils.evaluate_fn(val_loader, model, loss_fn, config.DEVICE, num_classes, model_architecture,  beta=config.FBETA_BETA)

        # --- 3.3. 日志记录 及 显示 ---
        log_data = {'epoch': epoch + 1, 'train_loss': train_loss}
        log_data.update(val_metrics_dict)
        utils.log_metrics_to_csv(log_data, metrics_csv_path)
        
        val_loss = val_metrics_dict.get('val_loss')
        iou_score = val_metrics_dict.get('iou_score', 0)
        fbeta_score = val_metrics_dict.get('fbeta_score', 0)

        logging.info(f"训练集Loss: {train_loss:.4f} | 验证集Loss: {val_loss:.4f} | IoU: {iou_score:.4f} | F-beta: {fbeta_score:.4f}")

        # --- 3.4. 早停与模型检查点 ---
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            epochs_no_improve = 0
            if isinstance(model, torch.nn.DataParallel):
                torch.save(model.module.state_dict(), best_model_path)
            else:
                torch.save(model.state_dict(), best_model_path)
            logging.info(f"验证损失改善至 {best_val_loss:.4f}。正在保存最佳模型至 '{best_model_path}'。")
        else:
            epochs_no_improve += 1
            logging.info(f"验证损失已连续 {epochs_no_improve} 轮未改善。")
        
        if isinstance(config.EARLY_STOPPING_PATIENCE, int) and epochs_no_improve >= config.EARLY_STOPPING_PATIENCE:
            logging.info(f"\n训练在 {epoch + 1} 轮后触发早停。")
            logging.warning(f"验证损失已连续 {config.EARLY_STOPPING_PATIENCE} 轮未改善。")
            break

        # --- 3.5. 检查点及曲线图保存 ---
        checkpoint_path = output_dir / f"epoch_{epoch+1}.pth"
        if isinstance(model, torch.nn.DataParallel):
            torch.save(model.module.state_dict(), checkpoint_path)
        else:
            torch.save(model.state_dict(), checkpoint_path)
        logging.info(f"第 {epoch+1} 轮的检查点已保存至 '{checkpoint_path}'。")

        # ---------- 清理旧检查点 ----------
        # 只在 epoch > 0 时才开始删，避免第一轮就误删
        if epoch > 0:
            # 要尝试删除的 epoch 编号 = 上一轮
            last_epoch = epoch
            # 如果上一轮不是“10 的整数倍”就删掉
            if last_epoch % config.CKPT_SAVE_INTERVAL != 0:
                old_ckpt = output_dir / f"epoch_{last_epoch}.pth"
                if old_ckpt.exists():
                    old_ckpt.unlink() # 文件删除

        utils.plot_training_progress(metrics_csv_path, output_dir)
        logging.info("所有训练曲线图已更新并保存。")

    return run_name

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

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

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

    # 将解析出的架构名称传入 main 函数并执行
    project_folder = None
    try:
        project_folder = main(model_architecture=args.architecture)
    finally:
        # 强制进行垃圾回收，释放 Python 对象占用的系统内存 (RAM)
        gc.collect()
        logging.info("已执行垃圾回收 (gc.collect())。")
        # 清理 PyTorch 在 CUDA 上缓存的显存
        torch.cuda.empty_cache()
        logging.info("训练结束，已清理 CUDA 缓存。")

    # 训练和清理完成后，移动结果到硬盘
    if project_folder:
        move_results_to_hardisk(project_folder)
    else:
        logging.error("由于训练失败或未生成项目文件夹，未执行结果移动操作。")