BuildFeatures/BuildMSAs.sh
#!/bin/bash

if [[ -z "${DistFeatureHome}" ]]; then
        echo "ERROR: Please set environmental variable DistFeatureHome to the installation folder of BuildFeatures "
        exit 1
fi

if [ -z "${HHDIR}" ]; then
	echo "ERROR: please set environmental variable HHDIR to the install folder of HHblits"
	exit 1
fi

if [ ! -d $HHDIR ]; then
        echo "ERROR: invalid folder $HHDIR "
        exit 1
fi

if [ -z "${HHDB}" ]; then
	echo "ERROR: please set environmental variable HHDB to the sequence database to be searched by HHblits"
	exit 1
fi

HHPkg=$DistFeatureHome/HHblitsWrapper/
DB4Thread=$HHDB
DB4HHMSA=$HHDB
DB4Jack=${JackDB}

## the following flags are used for MSA method selection
Flag4ThreadMSA=1
Flag4HHblitsV2=2
Flag4JackHmmer=4
Flag4HHblitsV3=8
Flag4MetaData=16

ResultDir=`pwd`

MSAmode=25
numCPUs=5

function Usage 
{
	echo $0 "[ -d ResultDir | -m MSAmethod | -t SeqDB4Threading | -h SeqDB4HHMSA | -j SeqDB4Jackhmmer | -c numCPUs ] seqFile"
	echo "	This script generates MSAs and features for threading and contact/distance prediction using HHblits and Jackhmmer"
	echo "	seqFile: the query protein sequence file in FASTA format"
	echo "	-d: the folder for result saving, default current work directory; a subfolder proteinName_OUT/ will be created in this folder"
	echo "	-m: an integer indicating MSA generation methods, formed by combining 1, 2, 4, 8 and 16, default $MSAmode (1+8+16)"
	echo "	   1: run HHblits to generate MSA for local structure property prediction and for threading"
	echo "	   2: run HHblits 2.0 to generate MSA for contact (obsolete)"
	echo "	   4: run Jackhmmer to generate MSA for contact and distance prediction"
	echo "	   8: run HHblits to generate MSA for contact and distance prediction"
	echo "	   16: search MetaGenome database for each MSA generated by the above methods"
	echo "	   By default, building MSAs for threading and contact/distance prediction using HHblits and metagenome data"
	echo " "
	echo "	-t: sequence database used by HHblits for sequence profile generation, default $DB4Thread"
	echo "	-h: sequence database used by HHblits for MSA generation for contact/distance prediction, default $DB4HHMSA"
	echo "	-j: sequence database used by Jackhmmer for MSA generation for contact/distance prediction, default $DB4Jack"
	echo "	-c: the number of CPUs to be used by HHblits and Jackhmmer, default $numCPUs"
}

while getopts ":m:d:t:h:j:c:" opt; do
        case ${opt} in
                t )
                  DB4Thread=`readlink -f $OPTARG`
                  ;;
                m )
                  MSAmode=$OPTARG
                  ;;
                d )
                  ResultDir=$OPTARG
                  ;;
                h )
                  DB4HHMSA=`readlink -f $OPTARG`
                  ;;
                j )
                  DB4Jack=`readlink -f $OPTARG`
                  ;;
		c )
		  numCPUs=$OPTARG
		  ;;
                \? )
                  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 1 ]; then
	Usage
	exit 1
fi

input_fasta=$1
if [ ! -f ${input_fasta} ]; then
	echo "ERROR: invalid query sequence file ${input_fasta}"
	exit 1
fi

fulnam=`basename $input_fasta`
relnam=${fulnam%.*}

out_folder=`readlink -f $ResultDir`
outdir=$out_folder/${relnam}_OUT
if [ ! -d $outdir ]; then
	mkdir -p $outdir
fi

re='^[0-9]+$'
if ! [[ $MSAmode =~ $re ]]; then
	echo "ERROR: MSA_mode is not a positive integer"
	exit 1
fi
if [ $MSAmode -lt 1 -o $MSAmode -gt 31 ]; then
	echo "ERROR: MSAmode shall be an integer between 1 and 31"
	exit 1
fi

RunThread=$(( ($MSAmode & $Flag4ThreadMSA) == $Flag4ThreadMSA ))
#RunHHblitsV2=$(( ($MSAmode & $Flag4HHblitsV2) == $Flag4HHblitsV2 ))
RunHHblitsV2=False
RunHHblitsV3=$(( ($MSAmode & $Flag4HHblitsV3) == $Flag4HHblitsV3 ))
RunJackhmmer=$(( ($MSAmode & $Flag4JackHmmer) == $Flag4JackHmmer ))
UseMetaData=$(( ($MSAmode & $Flag4MetaData) == $Flag4MetaData ))

$HHPkg/util/Verify_FASTA $input_fasta $outdir/$relnam.seq

#-> 1. generate MSA and sequence profile by hhsuite for threading
if (( $RunThread )); then

	if [ ! -f ${DB4Thread}_hhm.ffindex ]; then
        	echo "ERROR: invalid or damaged sequence database to be searched by hhblits: $DB4Thread"
        	exit 1
	fi

	mkdir -p $outdir/${relnam}_thread

	currDir=`pwd`
	cd $outdir/${relnam}_thread

	if [ ! -f "$outdir/${relnam}_thread/$relnam.a3m" ] || [ ! -f "$outdir/${relnam}_thread/$relnam.tgt" ]; then
		echo "Building MSA for threading from $DB4Thread by hhblits ..."
		iteration=2
		e_value=0.001
		neffmax=6

		$HHDIR/bin/hhblits -i $outdir/$relnam.seq -cpu $numCPUs -d $DB4Thread -o $relnam.hhr -oa3m $relnam.a3m -n $iteration -e $e_value -neffmax $neffmax

		tmp_root=$relnam"_tmp"/
		mkdir -p $tmp_root
		$DistFeatureHome/util/A3M_To_TGT -i $outdir/$relnam.seq -I $relnam.a3m -o $relnam.tgt -t $tmp_root
		rm -rf $tmp_root
	fi
	cd $currDir
fi

meff_cdhit=$DistFeatureHome/Meff_CDHIT/meff_cdhit
if [ ! -f $meff_cdhit ]; then
	echo "ERROR: invalid utility $meff_cdhit"
	exit 1
fi

MSADir=$outdir/${relnam}_contact

#-> 2. generate MSA for contact and distance prediction by HHsuites 2.0
if (( $RunHHblitsV2 )); then
	echo "ERROR: HHsuites v2.0 is not used any more"
	exit 1

	if [ ! -f ${DB4HHMSA}_hhm.ffindex ]; then
        	echo "ERROR: damaged sequence database $DB4HHMSA for HHblitsV2"
        	exit 1
	fi

	mkdir -p $MSADir

	if [ ! -f "$MSADir/${relnam}_v2uce3/$relnam.a3m" ] || [ ! -f "$MSADir/${relnam}_v2uce5/$relnam.a3m" ]; then
		echo "Building MSA from $DB4HHMSA by hhsuite 2.0 ..."
		iteration=3

		e_value=0.001
		$HHPkg/BuildMSAByHH.sh -i $outdir/$relnam.seq -n $iteration -e $e_value -C -1 -o $MSADir/${relnam}_v2uce3 -d $DB4HHMSA -c $numCPUs & 

		e_value=1
		$HHPkg/BuildMSAByHH.sh -i $outdir/$relnam.seq -n $iteration -e $e_value -C -1 -o $MSADir/${relnam}_v2uce5 -d $DB4HHMSA -c $numCPUs &
	
		wait

		#-> calculate meff
		numLines=`wc -l $MSADir/${relnam}_v2uce3/$relnam.a3m | cut -f1 -d' ' `
		if [ $numLines -ge 200000 ]; then
			echo 11 > $MSADir/${relnam}_v2uce3/$relnam.meff
		else
			$meff_cdhit -i $MSADir/${relnam}_v2uce3/$relnam.a3m > $MSADir/${relnam}_v2uce3/$relnam.meff &
		fi

		numLines=`wc -l $MSADir/${relnam}_v2uce5/$relnam.a3m | cut -f1 -d' ' `
		if [ $numLines -ge 200000 ]; then
			echo 11 > $MSADir/${relnam}_v2uce5/$relnam.meff
		else
			$meff_cdhit -i $MSADir/${relnam}_v2uce5/$relnam.a3m > $MSADir/${relnam}_v2uce5/$relnam.meff &
		fi

		wait
	fi
fi

#-> 3. generate MSA for contact and distance prediction
if (( $RunHHblitsV3 )); then
	if [ ! -f ${DB4HHMSA}_hhm.ffindex ]; then
        	echo "ERROR: invalid or damaged sequence database to be searched by hhblits: $DB4HHMSA"
        	exit 1
	fi

	mkdir -p $MSADir

	if [ ! -f "$MSADir/${relnam}_uce3/$relnam.a3m" ] || [ ! -f "$MSADir/${relnam}_uce5/$relnam.a3m" ]; then
		echo "Building MSAs for $relnam from $DB by hhsuite..."
		iteration=3

		e_value=0.001
		$HHPkg/BuildMSA4DistPred.sh -n $iteration -e $e_value -C -1 -o $MSADir/${relnam}_uce3 -d $DB4HHMSA -c $numCPUs $outdir/$relnam.seq &

		e_value=0.00001
		$HHPkg/BuildMSA4DistPred.sh -n $iteration -e $e_value -C -1 -o $MSADir/${relnam}_uce5 -d $DB4HHMSA -c $numCPUs $outdir/$relnam.seq &

		wait

		#-> calculate meff
		numLines=`wc -l $MSADir/${relnam}_uce3/$relnam.a3m | cut -f1 -d' ' `
                if [ $numLines -ge 200000 ]; then
                        echo 11 > $MSADir/${relnam}_uce3/$relnam.meff
                else
			$meff_cdhit -i $MSADir/${relnam}_uce3/$relnam.a3m > $MSADir/${relnam}_uce3/$relnam.meff &
		fi

		numLines=`wc -l $MSADir/${relnam}_uce5/$relnam.a3m | cut -f1 -d' ' `
                if [ $numLines -ge 200000 ]; then
                        echo 11 > $MSADir/${relnam}_uce5/$relnam.meff
                else
			$meff_cdhit -i $MSADir/${relnam}_uce5/$relnam.a3m > $MSADir/${relnam}_uce5/$relnam.meff &
		fi

		wait
	fi
fi

#--| 4. generate MSA for contact/distance prediction by Jackhmmer with uniref90
if (( $RunJackhmmer )); then
	if [ ! -f $DB4Jack ]; then
		echo "ERROR: invalid file for the sequence database to be searched by jackhmmer: $DB4Jack"
		echo "	you may fix this by setting the environmental variable $JackDB to the database file"
		echo "	or provide this database file as the input argument of this script"
		exit 1
	fi
	JackPkg=$DistFeatureHome/EVAlign/

	mkdir -p $MSADir

	if [ ! -f "$MSADir/${relnam}_ure3/$relnam.a3m" ] || [ ! -f "$MSADir/${relnam}_ure5/$relnam.a3m" ]; then
		echo "Building MSAs for $relnam from $DB4Jack by Jackhmmer..."
		iteration=3

		e_value=0.001
		#$JackPkg/BuildMSAByJack.sh -i $outdir/$relnam.seq -n $iteration -t $e_value -b 0 -o $MSADir/${relnam}_ure3 -d $DB4Jack -c $numCPUs &
		$JackPkg/BuildMSAByJack.sh -n $iteration -t $e_value -b 0 -o $MSADir/${relnam}_ure3 -d $DB4Jack -c $numCPUs $outdir/$relnam.seq &

		e_value=0.00001
		#$JackPkg/BuildMSAByJack.sh -i $outdir/$relnam.seq -n $iteration -t $e_value -b 0 -o $MSADir/${relnam}_ure5 -d $DB4Jack -c $numCPUs &
		$JackPkg/BuildMSAByJack.sh -n $iteration -t $e_value -b 0 -o $MSADir/${relnam}_ure5 -d $DB4Jack -c $numCPUs $outdir/$relnam.seq &

		wait
	
		if [ ! -f $MSADir/${relnam}_ure3/$relnam.a3m ]; then
			echo "ERROR: failed to generate $MSADir/${relnam}_ure3/$relnam.a3m by jackhmmer"
		else

			#-> calculate meff
			numLines=`wc -l $MSADir/${relnam}_ure3/$relnam.a3m | cut -f1 -d' ' `
                	if [ $numLines -ge 200000 ]; then
                        	echo 11 > $MSADir/${relnam}_ure3/$relnam.meff
                	else
				$meff_cdhit -i $MSADir/${relnam}_ure3/$relnam.a3m > $MSADir/${relnam}_ure3/$relnam.meff &
			fi

		fi

		if [ ! -f $MSADir/${relnam}_ure5/$relnam.a3m ]; then
			echo "ERROR: failed to generate $MSADir/${relnam}_ure5/$relnam.a3m by jackhmmer"
		else
			numLines=`wc -l $MSADir/${relnam}_ure5/$relnam.a3m | cut -f1 -d' ' `
                	if [ $numLines -ge 200000 ]; then
                        	echo 11 > $MSADir/${relnam}_ure5/$relnam.meff
                	else
				$meff_cdhit -i $MSADir/${relnam}_ure5/$relnam.a3m > $MSADir/${relnam}_ure5/$relnam.meff &
			fi
		fi

		wait
	fi
fi

## remove the folders with empty a3m files
for method in uce3 uce5 ure3 ure5 v2uce3 v2uce5
do
	if [ ! -f "$MSADir/${relnam}_${method}/$relnam.a3m" ]; then
                        continue
        fi

	if [ ! -s "$MSADir/${relnam}_${method}/$relnam.a3m" ]; then
		echo "WARNING: empty a3m file $MSADir/${relnam}_${method}/$relnam.a3m "
		rm -rf $MSADir/${relnam}_${method}
	fi
done

if (( $UseMetaData )); then
	MetaDataSearch=$DistFeatureHome/MetaDataSearch/BuildMSAFromMetaData.sh
	if [ ! -x $MetaDataSearch ]; then
        	echo "ERROR: non-executable script for metagenome database search: $MetaDataSearch"
        	exit 1
	fi
	## generate additional MSAs from meta data. the resultant MSAs contain both the sequences from the original MSAs and the sequences searched from the metagenome database
	for method in uce3 uce5 ure3 ure5 v2uce3 v2uce5
	do
		if [ ! -f "$MSADir/${relnam}_${method}/$relnam.a3m" ]; then
			continue
		fi

		meff=0
		if [ -f "$MSADir/${relnam}_${method}/$relnam.meff" ]; then
			meff=`cat $MSADir/${relnam}_${method}/$relnam.meff | cut -f1 -d'.'`
		fi

		## if ln(meff) > 6 or meff>403, do not use metagenome data
		if [ $meff -ge 6 ]; then
			continue
		fi

		$MetaDataSearch -m 1 -o $MSADir/${relnam}_${method}_meta $outdir/$relnam.seq $MSADir/${relnam}_${method}/$relnam.a3m &	
		sleep 1
	done
	wait

	## copy sequence and calculate meff
	for method in uce3 uce5 ure3 ure5 v2uce3 v2uce5
	do
		if [ ! -f "$MSADir/${relnam}_${method}_meta/$relnam.a3m" ]; then
			continue
		fi

		cp $outdir/$relnam.seq $MSADir/${relnam}_${method}_meta/
		$meff_cdhit -i $MSADir/${relnam}_${method}_meta/$relnam.a3m > $MSADir/${relnam}_${method}_meta/$relnam.meff &	
		sleep 1
	done
	wait 
fi
