#!/bin/bash
# --- 1. Conda 环境设置 ---
CONDA_BASE_PATH="/home/wkmgc/miniconda3" # <--- 在这里修改为您自己的路径
CONDA_ENV_NAME="${SEG_CONDA_ENV:-seg_smp}" # 可用 SEG_CONDA_ENV=SMP bash predict.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

# --- 2. 模型与 GPU 配置 ---
GPUS_GROUP_1="2"
GPUS_GROUP_2="3"
GPUS_GROUP_3="4"
GPUS_GROUP_4="5"

# 注意：这里的模型架构列表应与 train.sh 保持一致，以确保能找到对应的训练好的模型
GROUP_1_ARCHS=("PSPNet" "Unet" "UnetPlusPlus" )
GROUP_2_ARCHS=("Linknet" "MAnet" "DeepLabV3" )
GROUP_3_ARCHS=("UPerNet" "Segformer" "DPT"  )
GROUP_4_ARCHS=("FPN" "DeepLabV3Plus" "PAN")

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


# --- 3. 依次启动所有预测任务 ---
# 脚本将按顺序逐一执行每个模型的预测，等待上一个完成后再开始下一个。

echo ">>> 准备启动第一组预测任务 (后台运行)..."
for arch in "${GROUP_1_ARCHS[@]}"; do
    echo "   -> 正在后台启动模型: ${arch} on GPUs: ${GPUS_GROUP_1}"
    # 【修改点】: 使用 'echo "1" |' 来自动回答 1_predict.py 中的交互式提问，选择第一个找到的模型。
    echo "1" | CUDA_VISIBLE_DEVICES=${GPUS_GROUP_1} python 1_predict.py -a "${arch}" > "${LOG_DIR}/${arch}.log" 2>&1 &
    echo "      - 模型 ${arch} 预测已在后台启动。日志文件: ${LOG_DIR}/${arch}.log"
    echo "      - 等待 60 秒..."
    sleep 60
done
echo ">>> 第一组所有模型均已启动。"
echo "----------------------------------------------------"

echo ">>> 准备启动第二组预测任务 (后台运行)..."
for arch in "${GROUP_2_ARCHS[@]}"; do
    echo "   -> 正在后台启动模型: ${arch} on GPUs: ${GPUS_GROUP_2}"
    # 【修改点】: 自动选择第一个模型
    echo "1" | CUDA_VISIBLE_DEVICES=${GPUS_GROUP_2} python 1_predict.py -a "${arch}" > "${LOG_DIR}/${arch}.log" 2>&1 &
    echo "      - 模型 ${arch} 预测已在后台启动。日志文件: ${LOG_DIR}/${arch}.log"
    echo "      - 等待 50 秒..."
    sleep 60
done
echo ">>> 第二组所有模型均已启动。"
echo "----------------------------------------------------"

echo ">>> 准备启动第三组预测任务 (后台运行)..."
for arch in "${GROUP_3_ARCHS[@]}"; do
    echo "   -> 正在后台启动模型: ${arch} on GPUs: ${GPUS_GROUP_3}"
    # 【修改点】: 自动选择第一个模型
    echo "1" | CUDA_VISIBLE_DEVICES=${GPUS_GROUP_3} python 1_predict.py -a "${arch}" > "${LOG_DIR}/${arch}.log" 2>&1 &
    echo "      - 模型 ${arch} 预测已在后台启动。日志文件: ${LOG_DIR}/${arch}.log"
    echo "      - 等待 50 秒..."
    sleep 60
done
echo ">>> 第三组所有模型均已启动。"
echo "----------------------------------------------------"

echo ">>> 准备启动第四组预测任务 (后台运行)..."
for arch in "${GROUP_4_ARCHS[@]}"; do
    echo "   -> 正在后台启动模型: ${arch} on GPUs: ${GPUS_GROUP_4}"
    # 【修改点】: 自动选择第一个模型
    echo "1" | CUDA_VISIBLE_DEVICES=${GPUS_GROUP_4} python 1_predict.py -a "${arch}" > "${LOG_DIR}/${arch}.log" 2>&1 &
    echo "      - 模型 ${arch} 预测已在后台启动。日志文件: ${LOG_DIR}/${arch}.log"
    echo "      - 等待 50 秒..."
    sleep 60
done
echo ">>> 第四组所有模型均已启动。"
echo "----------------------------------------------------"

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

# 退出前取消激活环境
conda deactivate