#!/bin/bash
#
# 并行预测脚本 - 将大型数据集分成 N 份并行处理
#
# 用法:
#   bash scripts/run_parallel_predict.sh [输入文件] [输出文件] [进程数] [GPU设备]
#
# 示例:
#   bash scripts/run_parallel_predict.sh Data/fragment/Frags-Enamine-18M.csv output.csv 4 cuda:0
#
# 说明:
#   - 进程数建议 = GPU显存(GB) / 2.5
#   - 32GB 显存建议使用 4-12 个进程
#   - 所有进程将使用同一个 GPU（串行使用）
#
# 作者: AI Assistant
# 日期: 2025-10-17

set -e

# ============================================================================
# 参数设置
# ============================================================================

# 默认参数
INPUT_FILE="${1:-Data/fragment/Frags-Enamine-18M.csv}"
OUTPUT_FILE="${2:-Data/fragment/Frags-Enamine-18M_predicted.csv}"
N_PROCESSES="${3:-4}"
GPU_DEVICE="${4:-cuda:0}"

# SMILES 和 ID 列名
SMILES_COLUMN="smiles"
ID_COLUMN="chem_id"

# ============================================================================
# 打印配置信息
# ============================================================================

echo "============================================================"
echo "🚀 并行预测 Enamine 抗菌活性"
echo "============================================================"
echo ""
echo "配置:"
echo "  输入文件: $INPUT_FILE"
echo "  输出文件: $OUTPUT_FILE"
echo "  并行进程数: $N_PROCESSES"
echo "  GPU 设备: $GPU_DEVICE"
echo "  SMILES 列: $SMILES_COLUMN"
echo "  ID 列: $ID_COLUMN"
echo ""
echo "说明:"
echo "  - 使用批量预测工具: utils/batch_predictor.py"
echo "  - 每个模型实例约占用 2.5GB GPU 显存"
echo "  - 建议进程数 = GPU显存(GB) / 2.5"
echo "  - 32GB 显存 → 建议 4-12 个进程"
echo ""
echo "============================================================"

# ============================================================================
# 检查文件
# ============================================================================

if [ ! -f "$INPUT_FILE" ]; then
    echo "❌ 错误: 输入文件不存在: $INPUT_FILE"
    exit 1
fi

if [ ! -f "utils/batch_predictor.py" ]; then
    echo "❌ 错误: 预测脚本不存在: utils/batch_predictor.py"
    exit 1
fi

# ============================================================================
# 运行预测
# ============================================================================

echo ""
echo "🔄 开始预测..."
echo ""

# 使用 nohup 在后台运行（可选）
if [ "$5" == "--background" ]; then
    LOG_FILE="${OUTPUT_FILE%.csv}.log"
    echo "后台运行模式，日志保存到: $LOG_FILE"
    
    nohup pixi run python utils/batch_predictor.py \
        --input "$INPUT_FILE" \
        --output "$OUTPUT_FILE" \
        --smiles-column "$SMILES_COLUMN" \
        --id-column "$ID_COLUMN" \
        --device "$GPU_DEVICE" \
        --n-processes "$N_PROCESSES" \
        --batch-size 1000 \
        --verbose \
        > "$LOG_FILE" 2>&1 &
    
    echo "✓ 进程已在后台启动"
    echo "  进程 ID: $!"
    echo "  查看日志: tail -f $LOG_FILE"
else
    # 前台运行
    pixi run python utils/batch_predictor.py \
        --input "$INPUT_FILE" \
        --output "$OUTPUT_FILE" \
        --smiles-column "$SMILES_COLUMN" \
        --id-column "$ID_COLUMN" \
        --device "$GPU_DEVICE" \
        --n-processes "$N_PROCESSES" \
        --batch-size 1000 \
        --verbose
fi

echo ""
echo "============================================================"
echo "✅ 完成"
echo "============================================================"