#!/bin/bash
#SBATCH -J xxxxxxxx
#SBATCH --partition=gpu-prodq
#SBATCH --account=gpu_users
#SBATCH --gres=gpu:1
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --cpus-per-task=20
#SBATCH --qos=gpu
#SBATCH -o %x-%j.out
#SBATCH -e %x-%j.err

# Load the GROMACS module with CUDA support for GPU-accelerated MD simulations
module load GROMACS/2021.3-foss-2021a-CUDA-11.3.1

# Define global variables used throughout the script
PREFIX=""           # Prefix applied to all output file names
TOP="topol.top"          # Main GROMACS topology file
INIT_STRUCT="receptor.pdb"  # Initial protein structure in PDB format
LIG="LIG.pdb"            # Ligand structure in PDB format

RESULTS_DIR="${PREFIX}_results"

# ==============================================================================
# SECTION 1: Generate .mdp parameter files for each simulation stage
# ==============================================================================

# --- 1.1 ions.mdp ---
# Steepest-descent energy minimization parameters used only to prepare the
# system for ion addition (no PME, simple cutoff electrostatics)
cat <<EOF > ${PREFIX}_ions.mdp
; ions.mdp - used as input into grompp to generate ions.tpr
integrator  = steep         ; Steepest descent energy minimization algorithm
emtol       = 1000.0        ; Stop when maximum force < 1000.0 kJ/mol/nm
emstep      = 0.01          ; Step size for minimization
nsteps      = 50000         ; Maximum number of minimization steps

nstlist         = 1         ; Update neighbor list every step
cutoff-scheme   = Verlet    ; Buffered neighbor searching (recommended)
ns_type         = grid      ; Grid-based neighbor searching
coulombtype     = cutoff    ; Simple cutoff for electrostatics (no PME needed here)
rcoulomb        = 1.0       ; Short-range electrostatic cutoff (nm)
rvdw            = 1.0       ; Short-range van der Waals cutoff (nm)
pbc             = xyz       ; Periodic boundary conditions in all three dimensions
EOF

# --- 1.2 minim.mdp ---
# Energy minimization with PME electrostatics, used after ion addition
# to relax the full solvated system before dynamics
cat <<EOF > ${PREFIX}_minim.mdp
; minim.mdp - used as input into grompp to generate em.tpr
integrator  = steep         ; Steepest descent minimization
emtol       = 1000.0        ; Convergence criterion: max force < 1000.0 kJ/mol/nm
emstep      = 0.01          ; Minimization step size
nsteps      = 50000         ; Maximum minimization steps

nstlist         = 1         ; Neighbor list update frequency
cutoff-scheme   = Verlet    ; Buffered Verlet neighbor searching
ns_type         = grid      ; Grid-based neighbor search
coulombtype     = PME       ; Particle Mesh Ewald for long-range electrostatics
rcoulomb        = 1.0       ; Electrostatic cutoff (nm)
rvdw            = 1.0       ; van der Waals cutoff (nm)
pbc             = xyz       ; Periodic boundary conditions in all dimensions
EOF

# --- 1.3 nvt.mdp ---
# NVT equilibration (constant Number of particles, Volume, Temperature)
# The protein and ligand are position-restrained to allow solvent relaxation
cat <<EOF > ${PREFIX}_nvt.mdp
title                   = NVT equilibration 
define                  = -DPOSRES  ; Activate position restraints on protein/ligand
integrator              = md        ; Leap-frog MD integrator
nsteps                  = 50000     ; 50000 steps × 2 fs = 100 ps total
dt                      = 0.002     ; Time step: 2 fs

; Save coordinates, velocities, energies and log every 500 steps (1 ps)
nstxout                 = 500
nstvout                 = 500
nstenergy               = 500
nstlog                  = 500

continuation            = no        ; First dynamics run (generate velocities)
constraint_algorithm    = lincs     ; LINCS algorithm for bond constraints
constraints             = h-bonds   ; Constrain all bonds involving hydrogen
lincs_iter              = 1         ; LINCS accuracy iterations
lincs_order             = 4         ; LINCS expansion order

cutoff-scheme           = Verlet
ns_type                 = grid
nstlist                 = 10        ; Update neighbor list every 20 fs
rcoulomb                = 1.0       ; Electrostatic cutoff (nm)
rvdw                    = 1.0       ; van der Waals cutoff (nm)
DispCorr                = EnerPres  ; Long-range dispersion correction for energy and pressure

coulombtype             = PME       ; Particle Mesh Ewald electrostatics
pme_order               = 4         ; Cubic interpolation for PME
fourierspacing          = 0.16      ; FFT grid spacing (nm)

; Temperature coupling using velocity-rescaling thermostat
tcoupl                  = V-rescale
tc-grps                 = Protein_LIG Water_and_ions   ; Two separate coupling groups
tau_t                   = 0.1     0.1                  ; Coupling time constant (ps)
ref_t                   = 300     300                  ; Target temperature: 300 K

pcoupl                  = no        ; No pressure coupling during NVT
pbc                     = xyz

; Generate initial velocities from Maxwell-Boltzmann distribution at 300 K
gen_vel                 = yes
gen_temp                = 300
gen_seed                = -1        ; Random seed
EOF

# --- 1.4 npt.mdp ---
# NPT equilibration (constant Number of particles, Pressure, Temperature)
# Allows the simulation box to relax to the correct density at 1 bar
cat <<EOF > ${PREFIX}_npt.mdp
title                   = NPT equilibration 
define                  = -DPOSRES  ; Keep position restraints active
integrator              = md
nsteps                  = 50000     ; 100 ps
dt                      = 0.002

nstxout                 = 500
nstvout                 = 500
nstenergy               = 500
nstlog                  = 500

continuation            = yes       ; Continuation from NVT run
constraint_algorithm    = lincs
constraints             = h-bonds
lincs_iter              = 1
lincs_order             = 4

cutoff-scheme           = Verlet
ns_type                 = grid
nstlist                 = 10
rcoulomb                = 1.0
rvdw                    = 1.0
DispCorr                = EnerPres

coulombtype             = PME
pme_order               = 4
fourierspacing          = 0.16

tcoupl                  = V-rescale
tc-grps                 = Protein_LIG Water_and_ions
tau_t                   = 0.1     0.1
ref_t                   = 300     300

; Enable Parrinello-Rahman barostat for pressure coupling
pcoupl                  = Parrinello-Rahman
pcoupltype              = isotropic             ; Isotropic box scaling
tau_p                   = 2.0                   ; Pressure coupling time constant (ps)
ref_p                   = 1.0                   ; Target pressure: 1 bar
compressibility         = 4.5e-5                ; Isothermal compressibility of water (bar⁻¹)
refcoord_scaling        = com                   ; Scale reference coordinates with center of mass

pbc                     = xyz
gen_vel                 = no        ; Velocities inherited from NVT checkpoint
EOF

# --- 1.5 md.mdp ---
# Production MD run: 50 ns at 300 K and 1 bar with no position restraints
cat <<EOF > ${PREFIX}_md.mdp
title                   = MD 
integrator              = md
nsteps                  = 200000000 ; 200,000,000 × 2 fs = 400 ns production run
dt                      = 0.002

; Suppress heavy .trr coordinate/velocity/force output to save disk space
nstxout                 = 0
nstvout                 = 0
nstfout                 = 0
nstenergy               = 5000      ; Save energies every 10 ps
nstlog                  = 5000      ; Update log every 10 ps

; Save compressed trajectory every 10 ps (whole system)
nstxout-compressed      = 5000
compressed-x-grps       = System

continuation            = yes       ; Continuation from NPT checkpoint
constraint_algorithm    = lincs
constraints             = h-bonds
lincs_iter              = 1
lincs_order             = 4

cutoff-scheme           = Verlet
ns_type                 = grid
nstlist                 = 10
rcoulomb                = 1.0
rvdw                    = 1.0

coulombtype             = PME
pme_order               = 4
fourierspacing          = 0.16

tcoupl                  = V-rescale
tc-grps                 = Protein_LIG Water_and_ions
tau_t                   = 0.1     0.1
ref_t                   = 300     300

pcoupl                  = Parrinello-Rahman
pcoupltype              = isotropic
tau_p                   = 2.0
ref_p                   = 1.0
compressibility         = 4.5e-5

pbc                     = xyz
DispCorr                = EnerPres
gen_vel                 = no
EOF


# ==============================================================================
# SECTION 2: System preparation
# ==============================================================================

# Convert the protein PDB file to GROMACS format (.gro) and generate topology.
# Force field 8 = OPLS-AA/L, water model = SPC/E, -ignh ignores input hydrogens
gmx pdb2gmx -f "$INIT_STRUCT" \
            -o "${PREFIX}_prot_processed.gro" \
            -water spce \
            -ignh <<< $'8\n0\n0' 

# Convert the ligand PDB file to GROMACS .gro format
gmx editconf -f "$LIG" \
             -o "${PREFIX}_lig_processed.gro" 

# Merge protein and ligand .gro files into a single system .gro file.
# Step 1: Count total atoms from both files (line 2 of each .gro holds atom count)
total=$(($(sed -n '2p' *prot_processed.gro | tr -dc '0-9') + $(sed -n '2p' *lig_processed.gro | tr -dc '0-9')))

# Step 2: Update atom count in the protein .gro file header
sed -i "2s/.*/$total/" *prot_processed.gro

# Step 3: Remove the last line (box vectors) from the protein .gro file
sed -i '$d' *prot_processed.gro

# Step 4: Append ligand atom coordinates (skip header and box line)
cat *lig_processed.gro | tail -n +3 | head -n -1 >> *prot_processed.gro

# Step 5: Append the ligand box vector line to close the merged .gro file
tail -n 1 *lig_processed.gro >> *prot_processed.gro

# Define a cubic simulation box with at least 1.0 nm between protein and box edge,
# and center the system in the box
gmx editconf -f *prot_processed.gro \
             -o "${PREFIX}_prot_newbox.gro" \
             -c \
             -d 1.0 \
             -bt cubic 

# ------------------------------------------------------------------------------
# Update topol.top to include ligand topology and position restraint files
# ------------------------------------------------------------------------------

file="topol.top"

# Insert the ligand force field include (.itp) right after the main forcefield include
sed -i '/^#include ".*forcefield.itp"/a #include "LIG.itp"' "$file"

# Insert ligand position restraint block before the water topology section
sed -i '/; Include water topology/i ; Ligand position restraints\n#ifdef POSRES_LIG\n#include "posre_lig.itp"\n#endif' "$file"

# Add LIG entry to the [molecules] section if not already present
if ! grep -q "^LIG" "$file"; then
    if grep -q "^\[ molecules \]" "$file"; then
        if grep -q "^; Compound.*#mols" "$file"; then
            # Insert after the column header comment line
            sed -i '/^; Compound.*#mols/a LIG                 1' "$file"
        else
            # Insert before SOL or ION if present, otherwise right after [ molecules ]
            if grep -q "^SOL" "$file"; then
                sed -i '/^SOL/i LIG                 1' "$file"
            elif grep -q "^ION" "$file"; then
                sed -i '/^ION/i LIG                 1' "$file"
            else
                sed -i '/^\[ molecules \]/a LIG                 1' "$file"
            fi
        fi
    else
        # Create the entire [molecules] section if it does not exist
        echo -e "\n[ molecules ]\n; Compound        #mols\nJZ4                 1" >> "$file"
    fi
fi

# Solvate the protein-ligand system with SPC/E water molecules
# and update the topology with the number of water molecules added
gmx solvate -cp *prot_newbox.gro \
            -cs spc216.gro \
            -o "${PREFIX}_solv.gro" \
            -p "$TOP" 

# Prepare the .tpr input file for ion addition using the minimal ions.mdp parameters
gmx grompp -f *ions.mdp \
           -c *solv.gro \
           -p "$TOP" \
           -o "${PREFIX}_ions.tpr" -maxwarn 5 

# Add Na+ and Cl- ions to neutralize the net charge of the system.
# Group 15 corresponds to the SOL (water) group in the default index
gmx genion -s *ions.tpr \
           -o "${PREFIX}_solv_ions.gro" \
           -p "$TOP" \
           -pname NA \
           -nname CL \
           -neutral <<< "15" 

# Prepare the energy minimization .tpr file using the solvated + ionized system
gmx grompp -f *minim.mdp \
           -c *solv_ions.gro \
           -p "$TOP" \
           -o "${PREFIX}_em.tpr" -maxwarn 5 

# Run energy minimization to remove steric clashes introduced during solvation/ion placement
# -ntmpi 1 forces single MPI thread (useful for GPU runs)
gmx mdrun -v -deffnm "${PREFIX}_em" -ntmpi 1

# Create an index file for the ligand heavy atoms (non-hydrogen) 
# to be used for generating position restraints
gmx make_ndx -f "$LIG" \
             -o "${PREFIX}_index_lig.ndx" <<< $'0 & ! a H*\nq' 

# Generate position restraint file for the ligand with force constant 1000 kJ/mol/nm²
# Group 3 = heavy atoms of the ligand
gmx genrestr -f "$LIG" \
             -n *index_lig.ndx \
             -o posre_lig.itp \
             -fc 1000 1000 1000 <<< $'3\nq' 

# Create the final index file for the full solvated system.
# Merge Protein (group 4) and Ligand (group 2) into a new group called Protein_LIG
# This combined group is required for temperature coupling in NVT/NPT/MD
gmx make_ndx -f *em.gro -o "${PREFIX}_index.ndx" << EOF
4 | 2
name 24 Protein_LIG
q
EOF

# Prepare .tpr for NVT equilibration using the energy-minimized structure
# -r specifies the reference coordinates for position restraints
gmx grompp -f *nvt.mdp \
           -c *em.gro \
           -r *em.gro \
           -p topol.top \
           -n *index.ndx \
           -o "${PREFIX}_nvt.tpr"

# Run NVT equilibration to bring the system to target temperature (300 K)
gmx mdrun -deffnm nvt -v -s *nvt.tpr

# Prepare .tpr for NPT equilibration using NVT output structure and checkpoint
# -t passes the checkpoint file to continue velocities from NVT
gmx grompp -f *npt.mdp \
           -c *nvt.gro \
           -t *nvt.cpt \
           -r *nvt.gro \
           -p topol.top \
           -n *index.ndx \
           -o "${PREFIX}_npt.tpr"

# Run NPT equilibration to stabilize pressure and box density
gmx mdrun -deffnm npt -v -s *npt.tpr

echo "MD preparation completed successfully! Launching production simulation..."

# Hand off to the production MD script
./MD_simul.sh