DL4DistancePrediction4/Scripts/PredictPairRelation4Server.sh
#!/bin/bash

if [ -z "${ModelingHome}" ]; then
        echo "ERROR: Please set the environmental variable ModelingHome to the install folder of the RaptorX-3DModeling package"
        exit 1
fi

if [[ -z "${DL4DistancePredHome}" ]]; then
	echo "ERROR: please set the environmental variable DL4DistancePredHome to the install folder of DL4DistancePrediction4, e.g., RaptorX-3DModeling/DL4DistancePrediction4"
	exit 1
fi

#DeepModelFile=$DL4DistancePredHome/params/ModelFile4PairwisePred.txt
#ModelName=EC47C37C19CL99S35V2020MidModels

GPUMachineFile=$ModelingHome/params/GPUMachines.txt
DeepModelFile=""
ModelName=""
GPU=-1
ResultDir=""
alignmentType=0

printContactInCASPOnly=0

function Usage 
{
	echo $0 "[ -T alignmentType | -f DeepModelFile | -m ModelName | -d ResultDir | -g gpu | -h GPUMachineFile | -c ] proteinName RootDir"
        echo "	This script predicts distance/orientation for one protein from MSAs and optionally templates on a local or remote GPU"
	echo "	proteinName: the protein name"
	echo "	RootDir: a folder generated by BuildFeatures.sh, e.g. 1pazA_OUT/, that shall contain all needed feature files and folders including proteinName.seq, proteinName_contact and proteinName_thread"
	echo " "
	echo "	-T: indicate which query-template alignments to be used: 0 for none (default), 1 for alignments generated by HHpred and 2 for alignments generated by RaptorX threading"
        echo "		When this option is 0, predict distance/orientation from only sequence info such as MSAs"
        echo "		When not 0, predict distance/orientation matrices from both MSAs and query-template alignments"
	echo "		The alignment files shall be placed in RootDir/HHP and RootDir/DeepThreader, respectively"
	echo "	-f: a file containing some deep learning model files, default empty"
        echo "		Note that all the file paths shall be valid on the remote machine in order to work"
        echo "		When this option is not used, default DeepModelFile will be used"
        echo "	-m: a model name (defined in DeepModelFile) representing a set of deep learning models, default empty"
        echo "		Note that this name shall be valid on the remote machine in order to work"
        echo "		When this option is not used, the default model name will be used depending on alignment type"
	echo ""
	echo "	-d: the folder for result saving, default RootDir/DistancePred/. If RootDir is not writable, you shall provide another folder for result saving"
	echo "	-c: print only contact probability in CASP format, but not distance probability, default both"
	echo "	-g: -1 (default), 0-3. if -1, will automatically select a GPU"
	echo "	-h: a file containing remote machines or accounts with GPUs, to which you shall be able to ssh without password, default $GPUMachineFile"
	echo "		When this file does not exist, local GPUs will be used"
}

while getopts ":f:m:d:g:h:T:c" opt; do
        case ${opt} in
                f )
                  DeepModelFile=$OPTARG
                  ;;
                m )
                  ModelName=$OPTARG
                  ;;
                d )
                  ResultDir=$OPTARG
                  ;;
                g )
                  GPU=$OPTARG
                  ;;
		h )
		  GPUMachineFile=$OPTARG
		  ;;
		T )
		  alignmentType=$OPTARG
		  ;;
		c )
		  printContactInCASPOnly=1
		  ;;
                \? )
                  echo "Invalid Option: -$OPTARG" 1>&2
                  exit 1
                  ;;
                : )
                  echo "Invalid Option: -$OPTARG requires an argument" 1>&2
                  exit 1
                  ;;
        esac
done
shift $((OPTIND -1))

if [ $# -ne 2 ]; then
	Usage
	exit 1
fi

## we do not check if DeepModelFile exists or not since when RemoteAccount is used, it is not easy to do this.
## DeepModelFile will be checked later in PredictPairRelationLocal.sh anyway.
## an empty string is judged as a valid file in bash, so no need to tell if DeepModelFile is empty or not
##if [ ! -f $DeepModelFile ]; then
##	echo "ERROR: invalid file for the path info of deep learning models : $DeepModelFile"
##	exit 1
##fi

proteinName=$1

rootDir=`readlink -f $2`
if [ ! -d $rootDir ]; then
	echo "ERROR: invalid input folder for feature files $rootDir"
	exit 1
fi


if [ ! -f $GPUMachineFile ]; then
	UseRemoteMachine=false
else
	UseRemoteMachine=true
	localMachine=`hostname `
	machines=`grep -w on $GPUMachineFile | cut -f1 -d' ' `
	for machine in $machines
	do
		if [ $localMachine == "$machine" ]; then
        		UseRemoteMachine=false
                	break;
        	fi
	done

	if [[ -z "${machines}" ]]; then
                UseRemoteMachine=false
        fi
fi

if [ -z "${ResultDir}" ]; then
	ResultDir=$rootDir/DistancePred/
fi
if [ ! -d $ResultDir ]; then
	mkdir -p $ResultDir
fi

options="-g $GPU -d $ResultDir "
if [[ ! -z "${DeepModelFile}" ]]; then
        options=$options" -f $DeepModelFile "
fi

if [[ ! -z "${ModelName}" ]]; then
        options=$options" -m $ModelName "
fi

if [ $alignmentType -ne 0 ]; then
	options=$options" -T $alignmentType "
fi

if $UseRemoteMachine; then
	remoteAccount=`grep -w on $GPUMachineFile | shuf | head -1 |cut -f1 -d' '`
	#echo Running $DL4DistancePredHome/Scripts/PredictPairRelationRemote.sh $options -r $remoteAccount ${proteinName} $rootDir
        command="$DL4DistancePredHome/Scripts/PredictPairRelationRemote.sh $options -r $remoteAccount ${proteinName} $rootDir "
else
        command="$DL4DistancePredHome/Scripts/PredictPairRelationLocal.sh $options ${proteinName} $rootDir"
fi

echo Running $command
$command
if [ $? -ne 0 ]; then
	echo "ERROR: failed to run $command"
	exit 1
fi

## When seq-template alignments are used, we exit here since we do not need to generate the following information
if [ $alignmentType -ne 0 ]; then
	exit 0
fi

currDir=`pwd`
cd $ResultDir

$DL4DistancePredHome/Scripts/DeriveDistInfo4Threading.sh ${proteinName}.predictedDistMatrix.pkl
if [ $? -ne 0 ]; then
        echo "ERROR: failed to derive distance information for threading for $proteinName !"
        exit 1
fi

## print Cb-Cb contact prediction in CASP and matrix format
options=""
if [ $printContactInCASPOnly -eq 1 ]; then
	options=" -c "
fi

python $DL4DistancePredHome/PrintContactPrediction.py $options ${proteinName}.predictedDistMatrix.pkl
if [ $? -ne 0 ]; then
        echo "ERROR: failed to print predicted contact matrix for $proteinName !"
        exit 1
fi

## visualize the predicted contact map
python $DL4DistancePredHome/Utils/PlotContactMapByMatrix.py ${proteinName}.CM.txt
if [ $? -ne 0 ]; then
        echo "ERROR: failed to visualize the predicted contact matrix for ${proteinName} ! "
        exit 1
fi

## visualize the distance matrix
python $DL4DistancePredHome/Utils/EstimateAtomDistBounds.py ${proteinName}.predictedDistMatrix.pkl 
if [ $? -ne 0 ]; then
        echo "ERROR: failed to estimate inter-atom distance bound for ${proteinName} ! "
        exit 1
fi

python $DL4DistancePredHome/Utils/PlotDistanceMatrix.py ${proteinName}.bound.txt
if [ $? -ne 0 ]; then
        echo "ERROR: failed to visualize the predicted distance matrix for ${proteinName} ! "
        exit 1
fi

cd $currDir
