#!/bin/bash

set -e
shopt -s extglob

echo "======================================================="
echo "   Vina Split & OpenBabel Batch Conversion Script      "
echo "======================================================="

# 1. 显示文件列表并让用户选择（支持多选）
pdbqt_files=(*.pdbqt)
if [ ${#pdbqt_files[@]} -eq 0 ] || [ ! -e "${pdbqt_files[0]}" ]; then
    echo "❌ 错误: 当前目录下未找到任何 .pdbqt 文件。"
    exit 1
fi

echo "--> 可用的 PDBQT 文件列表:"
for i in "${!pdbqt_files[@]}"; do
    printf "  [%2d] %s\n" "$((i+1))" "${pdbqt_files[$i]}"
done

echo "-------------------------------------------------------"
echo "👉 请输入你想处理的文件编号（例如：输入 '1'，或多选输入 '1 3 5'，输入 'all' 选择全部）:"
read -p "选择: " user_choice

# 解析用户输入
selected_files=()
if [[ "$user_choice" == "all" ]]; then
    selected_files=("${pdbqt_files[@]}")
else
    for choice in $user_choice; do
        if [[ "$choice" =~ ^[0-9]+$ ]] && [ "$choice" -ge 1 ] && [ "$choice" -le "${#pdbqt_files[@]}" ]; then
            selected_files+=("${pdbqt_files[$((choice-1))]}")
        else
            echo "⚠️ 警告: 忽略无效的选择 '$choice'"
        fi
    done
fi

# 检查是否有有效选择
if [ ${#selected_files[@]} -eq 0 ]; then
    echo "❌ 错误: 未选择任何有效文件，程序退出。"
    exit 1
fi

echo "📋 已选择以下文件进行批处理:"
for f in "${selected_files[@]}"; do echo "  - $f"; done

# 2. 用户输入 pH 值
echo "-------------------------------------------------------"
read -p "--> 请输入 OpenBabel 转换的 pH 值 [默认: 7.0]: " ph_val
ph_val=${ph_val:-7.0}
echo "已设置 pH 值为: $ph_val"

# =======================================================
# 阶段 1: Vina Split (只需加载一次 Vina 环境)
# =======================================================
echo "-------------------------------------------------------"
echo "🔄 [Stage 1] 加载 Vina 环境并执行裂解..."
module restore autodock-vina

# 用于存储筛选出的 target_pose 文件路径，供阶段 2 使用
target_poses=()

for pdbqt_file in "${selected_files[@]}"; do
    base_name=$(basename "$pdbqt_file" .pdbqt)
    echo "Processing split for: $pdbqt_file ..."
    
    vina_split --input "$pdbqt_file"
    
    # 动态识别 Pose 1
    target_pose=""
    if [ -f "${base_name}_ligand_1.pdbqt" ]; then
        target_pose="${base_name}_ligand_1.pdbqt"
    elif [ -f "${base_name}_ligand_01.pdbqt" ]; then
        target_pose="${base_name}_ligand_01.pdbqt"
    else
        echo "❌ 错误: 未找到裂解后的第一个构象文件 (${base_name}_ligand_1.pdbqt 或 _01.pdbqt)。"
        exit 1
    fi
    
    # 清理其他多余的构象
    for file in "${base_name}"_ligand_*.pdbqt; do
        if [ -f "$file" ] && [ "$file" != "$target_pose" ]; then
            rm "$file"
        fi
    done
    
    # 将筛选好的 pose 存入数组
    target_poses+=("$target_pose")
done

# =======================================================
# 阶段 2: OpenBabel 转换 (只需加载一次 MD 环境)
# =======================================================
echo "-------------------------------------------------------"
echo "🔄 [Stage 2] 切换至 MD 环境并进行 OpenBabel 转换..."
module restore MDenv

for target_pose in "${target_poses[@]}"; do
    base_name="${target_pose%.pdbqt}"
    echo "Converting $target_pose to mol2 (pH: $ph_val) ..."
    
    obabel -ipdbqt "$target_pose" -omol2 -O "${base_name}.mol2" -p "$ph_val"
    
    # 清理中间文件
    rm "$target_pose"
done

echo "======================================================="
echo "🎉 所有文件处理完成！已生成对应的 *.mol2 构象文件。"
echo "======================================================="