#!/bin/bash

# =================================================================
#               YOLO 模型批量并行训练脚本
# =================================================================

# --- 1. Conda 环境设置 ---
CONDA_BASE_PATH="/home/wkmgc/miniconda3" # <--- 在这里修改为您自己的 Conda 路径
CONDA_ENV_NAME="${SEG_CONDA_ENV:-seg_smp}" # 可用 SEG_CONDA_ENV=SMP bash yolo_train.sh 临时覆盖

# 初始化并激活 Conda 环境
if [ -f "${CONDA_BASE_PATH}/etc/profile.d/conda.sh" ]; then
    source "${CONDA_BASE_PATH}/etc/profile.d/conda.sh"
    conda activate "${CONDA_ENV_NAME}"
    if [ $? -ne 0 ]; then
        echo "错误: 激活 Conda 环境 '${CONDA_ENV_NAME}' 失败！"
        exit 1
    fi
    echo "Conda 环境 '${CONDA_ENV_NAME}' 已成功激活。"
else
    echo "错误: 找不到 conda.sh 脚本。请检查您的 CONDA_BASE_PATH 设置是否正确。"
    exit 1
fi

echo "正在为当前终端会话设置 HTTP/HTTPS 代理..."
export https_proxy=http://127.0.0.1:1080
export http_proxy=http://127.0.0.1:1080
echo "代理已设置为 http://127.0.0.1:1080"
echo "如果模型权重不存在，将通过此代理进行下载。"

# --- 2. 模型与 GPU 配置 ---
# 根据您的 GPU 硬件和训练计划进行分组
# yolo_train.py 会自动使用分配给它的所有可见 GPU
GPUS_GROUP_0="0"
GPUS_GROUP_1="1"
GPUS_GROUP_2="2"
GPUS_GROUP_3="3"
# TODO
GPUS_GROUP_4="0"
GPUS_GROUP_5="1"
GPUS_GROUP_6="2"
GPUS_GROUP_7="3"

# GPUS_GROUP_4="4"
# GPUS_GROUP_5="5"
# GPUS_GROUP_6="6"
# GPUS_GROUP_7="7"

# 从 yolo_config.py 中选择要训练的模型，并分配到不同的组
# 建议将计算量/模型大小相似的模型放在同一组
GROUP_0_MODELS=("YOLO11l-seg" "YOLOv8m-seg")
GROUP_1_MODELS=("YOLO11x-seg")
GROUP_2_MODELS=("YOLO11n-seg" "YOLO11s-seg" "YOLO11m-seg") 
GROUP_3_MODELS=("YOLOv9e-seg") 
GROUP_4_MODELS=("YOLOv9c-seg")  
GROUP_5_MODELS=("YOLOv8x-seg")
GROUP_6_MODELS=("YOLOv8l-seg" "YOLO12-seg") # YOLOv9e-seg # 模型太大
GROUP_7_MODELS=("YOLOv8n-seg" "YOLOv8s-seg" ) # 训练完毕

# 1. 从 config.py 中读取 OUTPUTS_DIR 的值
OUTPUTS_DIR=$(python -c "from yolo_config import OUTPUTS_DIR; print(OUTPUTS_DIR)")
# 检查是否成功获取了 OUTPUTS_DIR
if [ -z "$OUTPUTS_DIR" ]; then
    echo "OUTPUTS_DIR: $OUTPUTS_DIR"
    echo "Error 1: Could not read OUTPUTS_DIR from yolo_config.py. Exiting."
    echo "Error 2: Or the directory specified by OUTPUTS_DIR does not exist. Please create it first."
    # exit 1
fi
# 2. 定义带有时间戳的日志目录名
LOG_DIR_NAME="yolo_train_logs_parallel_$(date +%Y-%m-%d_%H-%M-%S)"
# 3. 拼接成最终的完整路径
LOG_DIR="$OUTPUTS_DIR/$LOG_DIR_NAME"
mkdir -p "${LOG_DIR}"
echo "所有模型的日志将保存在 ./${LOG_DIR}/ 目录中。"
echo "----------------------------------------------------"


# --- 3. 训练执行函数 ---
# 定义一个函数来启动一组训练，以避免代码重复
start_training_group() {
    # 使用 nameref (引用) 来传递数组
    local -n models=$1
    local gpus=$2
    local group_name=$3

    echo ">>> 准备启动 ${group_name} 的训练任务 (后台运行)..."
    # 遍历指定组中的所有模型
    for model_key in "${models[@]}"; do
        echo "   -> 正在后台启动模型: ${model_key} on GPUs: ${gpus}"
        # 使用 '&' 将命令放入后台运行
        # 通过 --model 参数将模型名称传递给 yolo_train.py
        CUDA_VISIBLE_DEVICES=${gpus} python yolo_train.py --model "${model_key}" > "${LOG_DIR}/${model_key}.log" 2>&1 &
        echo "      - 模型 ${model_key} 已在后台启动。日志文件: ${LOG_DIR}/${model_key}.log"
        echo "      - 等待 30 秒，确保 GPU 资源稳定分配..."
        sleep 30
    done
    echo ">>> ${group_name} 的所有模型均已启动。"
    echo "----------------------------------------------------"
}

# --- 4. 依次启动所有训练任务 ---
# 脚本将快速地按顺序启动每一组任务到后台
start_training_group GROUP_0_MODELS "${GPUS_GROUP_0}" "第零组" # TODO
start_training_group GROUP_1_MODELS "${GPUS_GROUP_1}" "第一组"
start_training_group GROUP_2_MODELS "${GPUS_GROUP_2}" "第二组"
start_training_group GROUP_3_MODELS "${GPUS_GROUP_3}" "第三组"


# --- 5. 等待所有后台任务完成 ---
echo ""
echo "--- 所有模型均已在后台启动。现在等待所有训练任务完成... ---"
# 'wait' 命令会暂停脚本，直到所有由此脚本启动的后台子进程全部执行完毕
wait
echo "--- 所有后台训练任务已全部完成！ ---"

# 适用于4卡版本
start_training_group GROUP_4_MODELS "${GPUS_GROUP_4}" "第四组" # TODO
start_training_group GROUP_5_MODELS "${GPUS_GROUP_5}" "第五组"
start_training_group GROUP_6_MODELS "${GPUS_GROUP_6}" "第六组"
start_training_group GROUP_7_MODELS "${GPUS_GROUP_7}" "第七组"
# --- 5. 等待所有后台任务完成 ---
echo ""
echo "--- 所有模型均已在后台启动。现在等待所有训练任务完成... ---"
# 'wait' 命令会暂停脚本，直到所有由此脚本启动的后台子进程全部执行完毕
wait
echo "--- 所有后台训练任务已全部完成！ ---"

# --- 6. 退出脚本 ---
# 训练完成，取消激活 Conda 环境
echo "训练流程结束。"
# --- 取消网络代理设置 ---
echo "正在取消当前终端会话的 HTTP/HTTPS 代理设置..."
unset https_proxy
unset http_proxy
echo "代理已取消。"
conda deactivate
echo "取消激活 Conda 环境。"