#!/bin/bash
# Jason Key 9e4d1a4 2021-07-29 11:04:26 -0400
# Copyright © 2021 SBGrid Consortium.  All rights reserved.
#
# Wrapper script for alphafold
#
#
# To set single GPU
# export CUDA_VISIBLE_DEVICES=0
#
# Useful variables used by this script :
#
# ALPHAFOLD_DB       : set path to database directories when set
# ALPHAFOLD_PTM      : Use pTM models when set
# ALPHAFOLD_PRESET   : use reduced_dbs or CASP14 databases$Id$
# ALPHAFOLD_TEMPLATE : date string for limiting template search
#
# Modified by Normand Cyr 2021-08-12 to reflect the installation on our system
#
USAGE1="$(basename $0): A wrapper script for running Alphafold jobs in the SBGrid Software installation"
USAGE2="Usage: $(basename $0) [fasta file] [path to output directory]"
USAGE3="Default path to output is /tmp/alphafold/. Data stored in the default directory are deleted after 7 days."

if [ $# -eq 0 ]; then
  echo $USAGE1
  echo $USAGE2
  echo ""
  echo $USAGE3
  exit 1
fi

input_fasta=$1
if [ ! -f ${input_fasta} ]; then
  echo "${input_fasta} is not a file. Exiting... "
  exit 1
fi

# Output directory
output_dir="/tmp/alphafold"
if [ $2 ]; then
  output_dir=$2
fi
mkdir -p "${output_dir}"

## ALPHAFOLD databases
data_dir="/home/sbio/afdb"
if [ ! -z ${ALPHAFOLD_DB} ] && [ -d ${ALPHAFOLD_DB} ]; then
  data_dir=${ALPHAFOLD_DB}
fi

if [ ! -d ${data_dir} ]; then
  echo "${data_dir} is not a directory. Exiting... "
  exit 1
else
  echo "Using databases in ${data_dir}"
fi

# Set pTM models
if [ ! -z ${ALPHAFOLD_PTM} ]; then
  model_names="model_1_ptm,model_2_ptm,model_3_ptm,model_4_ptm,model_5_ptm"
  echo -e "Using pTM models \n"
else
  model_names="model_1,model_2,model_3,model_4,model_5"
fi

# Filter by date
if [ ! -z ${ALPHAFOLD_TEMPLATE} ]; then
  template_date=${ALPHAFOLD_TEMPLATE}
  echo -e "Using max-template-date $template_date \n"
else # use today's date
  template_date=$(date +"%Y-%m-%d")
fi

# database presets
if [ ! -z ${ALPHAFOLD_PRESET} ] && [ ${ALPHAFOLD_PRESET} == 'reduced_dbs' ]; then
  preset_args="--preset=reduced_dbs --small_bfd_database_path=${data_dir}/small_bfd"
  echo -e "Using --preset=reduced_dbs \n"
elif [ ! -z ${ALPHAFOLD_PRESET} ] && [ ${ALPHAFOLD_PRESET} == 'casp14' ]; then
  preset_args="-preset=casp14 \
--bfd_database_path=${data_dir}/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
--uniclust30_database_path=${data_dir}/uniclust30/uniclust30_2018_08/uniclust30_2018_08 "
  echo -e "Using --preset=casp14 \n"
# default is full_dbs
else
  preset_args="-preset=full_dbs \
--bfd_database_path=${data_dir}/bfd/bfd_metaclust_clu_complete_id30_c90_final_seq.sorted_opt \
--uniclust30_database_path=${data_dir}/uniclust30/uniclust30_2018_08/uniclust30_2018_08 "
fi

##
base_args="--data_dir=${data_dir} --output_dir=${output_dir} --fasta_paths=${input_fasta} \
--model_names=${model_names} --max_template_date=${template_date} "

database_args="--uniref90_database_path=${data_dir}/uniref90/uniref90.fasta \
--mgnify_database_path=${data_dir}/mgnify/mgy_clusters.fa \
--pdb70_database_path=${data_dir}/pdb70/pdb70 \
--template_mmcif_dir=${data_dir}/pdb_mmcif/mmcif_files \
--obsolete_pdbs_path=$data_dir/pdb_mmcif/obsolete.dat "

args="$base_args $preset_args $database_args"
alphafold_py="/programs/x86_64-linux/alphafold/2.0.0/bin.capsules/run_alphafold.py"

if [ ! -f $alphafold_py ]; then
  echo "Error - $alphafold_py not found"
  exit 1
else
  echo "${alphafold_py} ${args}"
  ${alphafold_py} ${args}
fi
