import argparse, shutil
import logging
from ultralytics import YOLO
from datetime import datetime
import yolo_config as config
import gc, torch

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

def train_model(model_key: str):
    """
    根据给定的模型密钥训练YOLO分割模型。

    Args:
        model_key (str): 在 yolo_config.MODEL_CONFIGS 中定义的模型密钥。
    """
    # 生成时间戳以区分不同训练运行
    timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    project_folder = f"{model_key}_{timestamp}"
    
    if model_key not in config.MODEL_CONFIGS:
        logging.error(f"错误：模型 '{model_key}' 不在 yolo_config.py 中定义。")
        logging.info(f"可用模型: {list(config.MODEL_CONFIGS.keys())}")
        return

    model_config = config.MODEL_CONFIGS[model_key]
    model_file = model_config['weights'] 
    logging.info(f"开始训练模型: {model_key} ({model_file})")

    try:
        # --- 1. 加载模型 ---
        model = YOLO(model_file) # 如果没有模型会自动下载
        logging.info("模型加载成功。")

        # --- 2. 模型训练 ---
        logging.info(f"数据集配置文件: {config.DATASET_YAML_PATH}")
        logging.info(f"训练参数: Epochs={config.EPOCHS}, Batch Size={config.BATCH_SIZE}, Img Size={config.IMAGE_SIZE}")

        model.train(
            data=str(config.DATASET_YAML_PATH),
            epochs=config.EPOCHS,
            imgsz=config.IMAGE_SIZE,
            batch=config.BATCH_SIZE,
            optimizer=config.OPTIMIZER,
            lr0=config.LEARNING_RATE,
            device=config.DEVICE,
            project=str(config.OUTPUTS_DIR), # 指定输出的根目录
            name=project_folder, # 指定本次训练的项目名
            workers=config.WORKERS,
            exist_ok=True, # 如果项目已存在，则覆盖
            patience=config.PATIENCE, # 提前停止训练的轮数
            save_period=config.SAVE_PERIOD # 每隔多少轮保存一次模型
        )

        logging.info("模型训练完成。")
        logging.info(f"训练结果保存在: {config.OUTPUTS_DIR / project_folder}")
        return project_folder

    except Exception as e:
        logging.error(f"训练过程中发生错误: {e}")
        return None

# --- 3. 复制结果到硬盘并删除原文件夹 (新增逻辑) ---
def move_results_to_hardisk(project_folder: str):
    source_dir = config.OUTPUTS_DIR / project_folder
    # 确保目标硬盘目录存在
    outputs_folder_name = config.OUTPUTS_DIR.name
    destination_dir = 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}")

if __name__ == "__main__":
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description="训练 YOLO 分割模型。")
    parser.add_argument(
        "--model",
        type=str,
        required=True,
        choices=list(config.MODEL_CONFIGS.keys()),
        help=f"选择要训练的模型。可选: {list(config.MODEL_CONFIGS.keys())}"
    )
    args = parser.parse_args()

    # 1. 开始训练
    project_folder = train_model(args.model)
    # 2. 训练完成后，移动结果到硬盘
    if project_folder:
        move_results_to_hardisk(project_folder)
    else:
        logging.error("由于训练失败，未执行结果移动操作。")
    # 3. 释放资源
    torch.cuda.empty_cache()  # 释放 CUDA 显存
    gc.collect()              # 释放系统内存
    logging.info("已清理 CUDA 显存与系统内存。")