#!/bin/bash
#########################################################################
#	SLURM scrip for running DRAM annotator on Orion cluster
#		Dependencies: 
#					  DRAM conda environment (/mnt/auve/mycondaenvs/DRAM)
#	gtdbtk.bac120.summary.tsv For taxonomy annotation
#	results.tsv	CheckM results for quality annotation
#
#	This script copies all the MAGs to a local disk on any Orion's node 
#	and runs locally on the node. At the end it copies all the results to the PEP
#	or any other path in the cluster into a DRAM.Results.dir folder.
#
#	to run: 
# sbatch  dram.GTDB.CM.SLURM.sh inputdir trans_table
#eg: sbatch  dram.GTDB.CM.SLURM.sh MAGS 11
#
# Author: Arturo Vera
# April 2021
#########################################################################

###############SLURM SCRIPT###################################

## Job name:
#SBATCH --job-name=DRAM
#
## Wall time limit:
#SBATCH --time=36:00:00
#
## Other parameters:
#SBATCH --cpus-per-task 14
#SBATCH --mem=150G
#SBATCH --partition=hugemem
#SBATCH -x cn-2

###########################################################

#####Array list######

#LIST=list.txt
#input=$(head -n $SLURM_ARRAY_TASK_ID $LIST | tail -n 1)


## Set up job environment:

module --quiet purge  # Reset the modules to the system default
module load Miniconda3

##Activate conda environments

export PS1=\$

##DRAM is at /mnt/users/auve/mycondaenvs/DRAM

source activate /net/cn-1/mnt/SCRATCH/bin420-21/condaenvironments/metaG/DRAM

####Do some work:########

## For debuggin
echo "Hello" $USER
echo "my submit directory is:"
echo $SLURM_SUBMIT_DIR
echo "this is the job:"
echo $SLURM_JOB_ID
echo "I am running on:"
echo $SLURM_NODELIST
echo "I am running with:"
echo $SLURM_CPUS_ON_NODE "cpus"
echo "Today is:"
date

##Variables

input=$1 #Directory with genomes
trans_table=$2 #Translation table used by prodigal
gtdbtk=$3
checkm=$4

## Copying data to local node for faster computation

cd $TMPDIR

#Check if $USER exists in $TMPDIR

if [[ -d $USER ]]
	then
        	echo "$USER exists on $TMPDIR"
	else
        	mkdir $USER
fi


echo "copying files to" $TMPDIR/$USER/tmpDir_of.$SLURM_JOB_ID

cd $USER
mkdir tmpDir_of.$SLURM_JOB_ID
cd tmpDir_of.$SLURM_JOB_ID

#Copy the MAGs to the $TMPDIR

echo "copying MAGs to" $TMPDIR/$USER/tmpDir_of.$SLURM_JOB_ID

cp -r $SLURM_SUBMIT_DIR/$input .
cp -r $SLURM_SUBMIT_DIR/$gtdbtk .
cp -r $SLURM_SUBMIT_DIR/$checkm .

##################DRAM##############################

echo "DRAM started at"
date +%d\ %b\ %T

time DRAM.py annotate \
-i $input'/*.fa' \
--trans_table $trans_table \
--gtdb_taxonomy $gtdbtk  \
--checkm_quality $checkm \
-o dram.annotation.$input.dir \
--threads $SLURM_CPUS_ON_NODE

echo "Distilling..."

time DRAM.py distill \
-i dram.annotation.$input.dir/annotations.tsv \
-o dram.genome_summaries.$input.dir \
--trna_path dram.annotation.$input.dir/trnas.tsv \
--rrna_path dram.annotation.$input.dir/rrnas.tsv

echo "DRAM finished at"
date +%d\ %b\ %T

mkdir DRAM.Results.$input.dir
mv dram.annotation.$input.dir DRAM.Results.$input.dir
mv dram.genome_summaries.$input.dir DRAM.Results.$input.dir
###########Moving results to PEP partition or anywhere the main script was submitted############

echo "moving results to" $SLURM_SUBMIT_DIR/$input

cd $TMPDIR/$USER/tmpDir_of.$SLURM_JOB_ID

time cp -r *.dir $SLURM_SUBMIT_DIR/$input

echo "DRAM results are in: " $SLURM_SUBMIT_DIR/$input/DRAM.Results.dir.$input.dir

####removing tmp dir. Remember to do this for not filling the HDD in the node!!!!###

cd $TMPDIR/$USER/
rm -r tmpDir_of.$SLURM_JOB_ID

echo "I've done at"
date
