#!/bin/bash

# ==============================================================================
# 0. 前置文件检查 (Pre-run check)
# ==============================================================================
# 检查当前目录下是否存在 inputs/ion.mdp 或 input/ion.mdp
if [ -f "inputs/ion.mdp" ]; then
    ION_MDP="inputs/ion.mdp"
elif [ -f "input/ion.mdp" ]; then
    ION_MDP="input/ion.mdp"
else
    echo -e "\033[31m[错误/Error] 'inputs/ion.mdp' NOT found in the current folder!\033[0m"
    echo "Please create the 'inputs' directory and put 'ion.mdp' inside it first."
    echo "Command suggestion: mkdir -p inputs && cp /path/to/template/ion.mdp inputs/"
    exit 1
fi

echo "Using found MDP file: $ION_MDP"
echo "--------------------------------------------------"

# ==============================================================================
# 1. 拓扑与坐标处理 (Topology & Coordinate preparation)
# ==============================================================================
# get the itp and gro from ligand outputs.
cp ./ligand.acpype/ligand_GMX.itp .
cp ./ligand.acpype/ligand_GMX.gro .

# Edit topol file first
echo "ligand              1" >> topol.top
sed -i '/#include "amber14sb.ff\/forcefield.itp"/a #include "ligand_GMX.itp"' topol.top

# Now edit gro file
head -n -1 protein.gro > complex.gro
sed '1,2d;$d' ligand_GMX.gro >> complex.gro
tail -n 1 protein.gro >> complex.gro

PROT_ATOMS=$(sed -n '2p' protein.gro | awk '{print $1}')
LIG_ATOMS=$(sed -n '2p' ligand_GMX.gro | awk '{print $1}')
TOTAL_ATOMS=$((PROT_ATOMS + LIG_ATOMS))
echo "Total atoms will be: $TOTAL_ATOMS"
sed -i "2s/.*/$TOTAL_ATOMS/" complex.gro

echo " ================== enjoy your complex.gro! ==============="

# ==============================================================================
# 2. 溶剂化与加离子 (Solvation & Ionization)
# ==============================================================================
gmx editconf -f complex.gro -o complex_box.gro -c -d 1.0 -bt dodecahedron
gmx solvate -cp complex_box.gro -cs spc216.gro -o complex_solv.gro -p topol.top

# 直接使用前面检查好的路径，不再进行交互式询问 (Directly use checked path)
gmx grompp -f "$ION_MDP" -c complex_solv.gro -p topol.top -o ions.tpr

printf "SOL\n" | gmx genion -s ions.tpr -o complex_neutral.gro -p topol.top -pname NA -nname CL -neutral
printf "0 & ! a H*\nq\n" | gmx make_ndx -f ligand_GMX.gro -o index_ligand.ndx
printf "3\n" | gmx genrestr -f ligand_GMX.gro -n index_ligand.ndx -o posre_ligand.itp -fc 1000 1000 1000
sed -i '/#include "ligand_GMX.itp"/a \\n; Ligand position restraints\n#ifdef POSRES_LIG\n#include "posre_ligand.itp"\n#endif' topol.top

echo "============ prepare to run your md now! ============"