#!/bin/bash
set -e

usage() {
    echo; echo "Usage: $0 --fasta=/Path/to/infile.fasta --databases=/Path/to/databases --working=/Path/to/working_directory [--threads=4] [--log=/Path/to/run.log]"
    echo "  --fasta     Path to the input NT FASTA file"
    echo "  --databases Path to the database directory"
    echo "  --working   Path to the output working directory"
    echo "  --threads   Number of threads to use (Default=1)"
    echo "  --log       Log file that the status can be tee'd to (Default=Don't save the log)"
    echo "  -h, --help  Print this help message out"; echo;
}

## ========================
## Globals
DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
MODULES="${DIR}/../modules"
SCRIPTS="${DIR}/../scripts"
THREADS=1
WORKFLOW="taxonomic_identification"
LOG="" ## Initialize empty in case no log is needed

## ========================
## Process the input parameters

if [ $# -gt 5 ] || [ $# -lt 3 ]
then
    usage
    exit 1
fi

while true
do
    case $1 in
    --help|-h)
	    usage
	    exit;;
    --fasta=?*)
	    FASTA=${1#*=};;
    --fasta|fasta=)
	    echo "$0: missing argument for '$1' option"
	    usage
	    exit 1;;
    --working=?*)
	    WORKING=${1#*=};;
    --working|working=)
	    echo "$0: missing argument for '$1' option"
	    usage
	    exit 1;;
    --databases=?*)
	    DATABASES=${1#*=};;
    --databases|databases=)
	    echo "$0: missing argument for '$1' option"
            usage
            exit 1;;
    --threads=?*)
            THREADS=${1#*=};;
    --threads|threads=)
            echo "$0: missing argument for '$1' option"
            usage
            exit 1;;
    --log=?*)
	    LOG=${1#*=};;
    --log|log=)
	    echo "$0: missing argument for '$1' option"
	    usage
	    exit 1;;
    --)
	    shift
	    break;;
    -?*)
	    echo "$0: invalid option: $1"
	    usage
	    exit 1;;
    *)
	    break
    esac
    shift
done

## ========================
## Create working directories

BASE=$(basename "${FASTA}")
mkdir -p "${WORKING}"
mkdir -p "${WORKING}/${WORKFLOW}/kraken"
mkdir -p "${WORKING}/${WORKFLOW}/blastn"
mkdir -p "${WORKING}/${WORKFLOW}/blastx"
mkdir -p "${WORKING}/${WORKFLOW}/tipp"
mkdir -p "${WORKING}/${WORKFLOW}/taxonomic_assignment"

## Step 0a: Convert ambiguous bases to non-ambiguous bases (if we need to)
if [ ! -f "${WORKING}/ambig_2_unambig/${BASE}.unambig.fasta" ]; then
    mkdir -p ${WORKING}/initialize/ambig_2_unambig
    ${SCRIPTS}/ambiguous_2_unambiguous.pl -f ${FASTA} -o ${WORKING}/initialize/ambig_2_unambig/${BASE}.unambig.fasta -n 4
    status=$?
fi

## Step 0b: Perform 6-frame translation of the input query sequence (if we need to)
if [ ! -f "${WORKING}/six_frame_translation/${BASE}.translated.fasta" ]; then
    mkdir -p ${WORKING}/initialize/six_frame_translation
    ${SCRIPTS}/six_frame_translation.pl -f ${WORKING}/initialize/ambig_2_unambig/${BASE}.unambig.fasta \
	-o ${WORKING}/initialize/six_frame_translation/${BASE}.translated.fasta
    status=$?
fi

## ========================
## Step 1: Kraken against Kraken MiniDB. Will update this to all complete RefSeq genomes when that
##         dataset is compiled
echo -n "# Running Kraken against the MiniKrakenDB ......" | tee -a ${LOG} ## If user provides no log file, this tee just won't happen
"${MODULES}/kraken.sh" --fasta="${WORKING}/initialize/ambig_2_unambig/${BASE}.unambig.fasta" \
    --database="${DATABASES}/kraken/minikraken_20141208" \
    --out="${WORKING}/${WORKFLOW}/kraken/${BASE}.kraken.txt" \
    --threads="${THREADS}"
status=$?
echo -n " DONE "  | tee -a ${LOG}; date '+%H:%M:%S %Y-%m-%d' |tee -a ${LOG}

## ========================
## Step 2: BLASTn against the NCBI nt database
echo -n "# Running BLASTn against NCBI nt database ......" | tee -a ${LOG} ## If user provides no log file, this tee just won't happen
"${MODULES}/blastn.sh" --fasta="${WORKING}/initialize/ambig_2_unambig/${BASE}.unambig.fasta" \
    --database="${DATABASES}/blast/nt" \
    --out="${WORKING}/${WORKFLOW}/blastn/${BASE}.nt" \
    --threads="${THREADS}" \
    --evalue="1e-3"
status=$?
echo -n " DONE "  | tee -a ${LOG}; date '+%H:%M:%S %Y-%m-%d' |tee -a ${LOG}

## ========================
## Step 3: BLASTx of the nucleotide FASTA against NCBI nr database
echo -n "# Running BLASTx against NCBI nr database ......" | tee -a ${LOG} ## If user provides no log file, this tee just won't happen
"${MODULES}/blastx-fast.sh" --fasta="${WORKING}/initialize/ambig_2_unambig/${BASE}.unambig.fasta" \
    --database="${DATABASES}/blast/UNIREF100" \
    --out="${WORKING}/${WORKFLOW}/blastx/${BASE}.ur100" \
    --threads="${THREADS}" \
    --evalue="1e-3"
status=$?
echo -n " DONE "  | tee -a ${LOG}; date '+%H:%M:%S %Y-%m-%d' |tee -a ${LOG}

## ========================
## Step 4: TIPP against pyrg
echo -n "# Running TIPP against pyrg database ..........." | tee -a ${LOG} ## If user provides no log file, this tee just won't happen
"${MODULES}/tipp.sh" --fasta="${WORKING}/initialize/ambig_2_unambig/${BASE}.unambig.fasta" \
    --out="${WORKING}/${WORKFLOW}/tipp/tipp_results" \
    --threads="${THREADS}"
status=$?
echo -n " DONE "  | tee -a ${LOG}; date '+%H:%M:%S %Y-%m-%d' |tee -a ${LOG}

## ========================
## Step 5: Assign taxonomy using BLASTn and BLASTx outputs
echo -n "# Generating taxonomy report ..................." | tee -a ${LOG} ## If user provides no log file, this tee just won't happen
"${SCRIPTS}/btab_2_tax_report.pl" --blastx="${WORKING}/${WORKFLOW}/blastx/${BASE}.ur100.btab" \
    --blastn="${WORKING}/${WORKFLOW}/blastn/${BASE}.nt.btab" \
    --out="${WORKING}/${WORKFLOW}/taxonomic_assignment/results.txt" \
    --cutoff="3"
status=$?
echo -n " DONE "  | tee -a ${LOG}; date '+%H:%M:%S %Y-%m-%d' |tee -a ${LOG}

exit ${status}
