#!/usr/bin/env bash
set -euo pipefail

# Run this script inside an already prepared GPU environment/node.
# It does not request a compute node.

cd /public/home/scnb9biwet/jiangqq/GPSite-main

GPU_ID="${1:-0}"
INPUT_FASTA="${2:-./conf/example/demo.fa}"
OUT_DIR="${3:-./results/}"

PROTTRANS_DIR="/path/to/prot_t5_xl_uniref50"
ESMFOLD_HUB_DIR="/path/to/weight"
ESMFOLD_WEIGHT="${ESMFOLD_HUB_DIR}/checkpoints/esmfold_3B_v1.pt"
ESM2_WEIGHT="${ESMFOLD_HUB_DIR}/checkpoints/esm2_t36_3B_UR50D.pt"

if [ ! -d "${PROTTRANS_DIR}" ]; then
  echo "ERROR: ProtTrans directory not found: ${PROTTRANS_DIR}" >&2
  exit 1
fi

if [ ! -f "${ESMFOLD_WEIGHT}" ]; then
  echo "ERROR: ESMFold weight not found: ${ESMFOLD_WEIGHT}" >&2
  exit 1
fi

if [ ! -f "${ESM2_WEIGHT}" ]; then
  echo "ERROR: ESM2 weight not found: ${ESM2_WEIGHT}" >&2
  exit 1
fi

if [ ! -f "${INPUT_FASTA}" ]; then
  echo "ERROR: input FASTA not found: ${INPUT_FASTA}" >&2
  exit 1
fi

chmod +x ./scripts/feature_extraction/mkdssp

export PROTTRANS_PATH="${PROTTRANS_DIR}"
export ESMFOLD_HUB_DIR="${ESMFOLD_HUB_DIR}"

echo "Using GPU: ${GPU_ID}"
echo "Input FASTA: ${INPUT_FASTA}"
echo "Output directory: ${OUT_DIR}"
echo "ProtTrans: ${PROTTRANS_DIR}"
echo "ESMFold hub dir: ${ESMFOLD_HUB_DIR}"

python ./scripts/predict.py \
  -i "${INPUT_FASTA}" \
  -o "${OUT_DIR}" \
  --gpu "${GPU_ID}"