#!/bin/bash
# Detached driver for one PROTAC candidate: chloroform 3 x 50 ns, analysis,
# then water 3 x 50 ns, analysis. The chloroform AM1-BCC charge cache is staged
# into the water run so sqm is only paid once.
#
#   run_candidate.sh <NAME> <GPU> <SMILES> [WAIT_LOGFILE] [WAIT_MARKER]
#
# With WAIT_LOGFILE and WAIT_MARKER the job queues instead of starting: it polls
# that log until the marker appears (the predecessor on the same GPU printed
# CANDIDATE_DONE), so two jobs never share a card. Waiting greps a file rather
# than process names, which would otherwise match this script's own cmdline.
# WAIT_LOGFILE resolves against this script's directory, not the caller's cwd,
# and after MAXWAIT the job stops waiting and starts anyway, which can put two
# jobs on one card.
set -u

NAME="$1"; GPU="$2"; SMI="$3"; WLOG="${4:-}"; WMARK="${5:-}"
MD="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CONDA_ENV="${CONDA_ENV:-protacmd}"
MAXWAIT=$((14*3600))     # never queue longer than this

if [ -n "${CONDA_SH:-}" ]; then
  source "$CONDA_SH"
elif [ -n "${CONDA_EXE:-}" ]; then
  source "$(dirname "$(dirname "$CONDA_EXE")")/etc/profile.d/conda.sh"
fi
conda activate "$CONDA_ENV"
cd "$MD"

log() { echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*"; }

log "CANDIDATE_START name=$NAME gpu=$GPU"

if [ -n "$WLOG" ] && [ -n "$WMARK" ]; then
  log "QUEUED on GPU $GPU: waiting for '$WMARK' in $WLOG ..."
  waited=0
  until grep -q "$WMARK" "$MD/$WLOG" 2>/dev/null; do
    sleep 60; waited=$((waited+60))
    if [ "$waited" -ge "$MAXWAIT" ]; then
      log "waited ${waited}s without seeing '$WMARK', starting anyway"; break
    fi
    [ $((waited % 1800)) -eq 0 ] && log "  ... still waiting (${waited}s)"
  done
  log "predecessor finished after ${waited}s, starting $NAME on GPU $GPU"
fi

nvidia-smi --query-gpu=index,utilization.gpu,memory.used --format=csv,noheader | sed 's/^/[gpu] /'

run_solvent() {
  local sol="$1"
  log "=== $NAME / $sol : 3 x 50 ns START ==="
  python protac_md.py --solvent "$sol" --replicas 3 --prod-ns 50 \
    --smiles "$SMI" --name "$NAME" --gpu "$GPU"
  local rc=$?
  log "=== $NAME / $sol : MD finished rc=$rc ==="
  if [ $rc -eq 0 ]; then
    log "--- $NAME / $sol : analysis START ---"
    python analyze_conformations.py --dir "$NAME/$sol" --name "$NAME"
    log "--- $NAME / $sol : analysis finished rc=$? ---"
  else
    log "$NAME / $sol MD failed (rc=$rc), skipping analysis"
  fi
}

run_solvent chloroform
CH="$MD/$NAME/chloroform/ff_cache.json"
if [ -s "$CH" ]; then
  mkdir -p "$MD/$NAME/water"; cp -f "$CH" "$MD/$NAME/water/ff_cache.json"
  log "staged charge cache -> $NAME/water/ff_cache.json"
fi
run_solvent water

log "CANDIDATE_DONE name=$NAME"