#!/bin/bash -eu

# RTAX: Rapid and accurate taxonomic classification of short paired-end
#       sequence reads from the 16S ribosomal RNA gene.
#
# David A. W. Soergel 1*, Rob Knight 2, and Steven E. Brenner 1
#
# 1 Department of Plant and Microbial Biology,
#   University of California, Berkeley
# 2 Howard Hughes Medical Institute and Department of Chemistry
#   and Biochemistry, University of Colorado at Boulder
# * Corresponding author: soergel@cs.umass.edu
#
# http://www.davidsoergel.com/rtax
#
# Version 0.984  (August 7, 2013)
#
# For usage instructions: just run the script with no arguments
#
#
# Copyright (c) 2011 Regents of the University of California
#
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#     * Redistributions of source code must retain the above copyright notice,
#       this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of the University of California, Berkeley nor the
#       names of its contributors may be used to endorse or promote products
#       derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


function echoerr() { echo "$@" 1>&2; }

echoerr
echoerr 'RTAX: Rapid and accurate taxonomic classification of short paired-end'
echoerr '      sequence reads from the 16S ribosomal RNA gene.'
echoerr ''
echoerr 'David A. W. Soergel (1), Rob Knight (2), and Steven E. Brenner (1)'
echoerr ''
echoerr '1 Department of Plant and Microbial Biology,'
echoerr '  University of California, Berkeley '
echoerr '2 Howard Hughes Medical Institute and Department of Chemistry '
echoerr '  and Biochemistry, University of Colorado at Boulder'
echoerr '* Corresponding author (current address): soergel@cs.umass.edu'
echoerr ''
echoerr 'Version 0.984  (August 7, 2013)'
echoerr ''
echoerr 'http://dev.davidsoergel.com/rtax'
echoerr ''
echoerr 'Please cite:'
echoerr 'Soergel D.A.W., Dey N., Knight R., and Brenner S.E.  2012.'
echoerr 'Selection of primers for optimal taxonomic classification of'
echoerr 'environmental 16S rRNA gene sequences.  ISME J (6), 1440–1444'
echoerr ''



# we want to call subsidiary scripts in the "scripts" subdir, but we don't know where we're installed or what the working directory is.
# solution from http://hintsforums.macworld.com/archive/index.php/t-73839.html

IFS=$' \t\n'
#declare -x PATH=/bin:/usr/bin:/usr/local/bin
arg=$0; [[ -L $0 ]] && arg=$(stat -f '%Y' "$0")
pth=$(2>/dev/null cd "${arg%/*}" >&2; echo "`pwd -P`/${arg##*/}")
par=$(dirname "$pth")
scripts="$par/scripts"

# The above solution works with the widespread GNU Coreutils 5.97.  Thanks to Bernd Brandt for this alternate solution for Coreutils 8.5.

#IFS=$' \t\n'
#arg=$0; [[ -L $0 ]] && arg=$(stat -c '%N' "$0") 
#pth=${arg##*-> \`};
#par=$(dirname "$pth")
#scripts="$par/scripts"

function abspath {
	if [[ -d "$1" ]]
	then
		pushd "$1" >/dev/null
		pwd
		popd >/dev/null
	#elif [[ -e $1 ]]
	#then
    else
		pushd $(dirname $1) >/dev/null
		echo $(pwd)/$(basename $1)
		popd >/dev/null
	#else
	#	echo $1 does not exist! >&2
	#	return 127
	fi
}

# Grab the command-line arguments

usageHelp="Usage: ${0##*/} -r <refdb> -t <taxonomy> -a <queryA> -b <queryB> -d <delimiter> -i <regex> -o <classifications.out>"
refdbHelp="-r  reference database in FASTA format"
taxonomyHelp="-t  taxonomy file with sequence IDs matching the reference database"
queryAHelp="-a  FASTA file containing query sequences (single-ended or read 1)"
queryBHelp="-b  FASTA file containing query sequences (read b, with matching IDs)"
revcompAHelp="-x  Reverse-complement query A sequences (required if they are provided in the reverse sense)"
revcompBHelp="-y  Reverse-complement query B sequences (required if they are provided in the reverse sense)"
idListHelp="-l  text file containing sequence IDs to process, one per line"
delimiterHelp="-d  delimiter separating the two reads when provided in a single file"
idRegexHelp="-i  regular expression used to select part of the fasta header to use as the sequence id.  Default: \"(\\\\S+)\""
outputHelp="-o  output path"
tmpdirHelp="-m  temporary directory.  Will be removed on successful completion, but likely not if there is an error."
fallbackHelp="-f  for sequences where only one read is available, fall back to single-ended classification.  Default: drop these sequences."
noGenericFallbackHelp="-g  for sequences where one read is overly generic, do not fall back to single-ended classification.  Default: classify these sequences based on only the more specific read."
badOptionHelp="Option not recognised"
printHelpAndExit()
{
	echoerr "$usageHelp"
	echoerr "$refdbHelp"
	echoerr "$taxonomyHelp"
	echoerr "$queryAHelp"
	echoerr "$queryBHelp"
	echoerr "$revcompAHelp"
	echoerr "$revcompBHelp"
	echoerr "$idRegexHelp"
	echoerr "$idListHelp"
	echoerr "$delimiterHelp"
	echoerr "$tmpdirHelp"
	echoerr "$fallbackHelp"
	echoerr "$noGenericFallbackHelp"
	echoerr "$outputHelp"
	exit $1
}
printErrorHelpAndExit()
{
        echoerr "-------------------------------------------------------------"
        echoerr "ERROR : $@"
        echoerr "-------------------------------------------------------------"
        echoerr
        printHelpAndExit 1
}

refdb=""
taxonomy=""
queryA=""
queryB=""
revcompA=""
revcompB=""
delimiter=""
idList=""
idRegex=""
verbose=""
fallback=""
noGenericFallback=""
outpath=""
tempdir=""

while getopts "hr:t:a:b:i:l:d:m:o:vfgxy" optionName; do
case "$optionName" in
h) printHelpAndExit 0;;
r) refdb=`abspath "$OPTARG"`;;
t) taxonomy=`abspath "$OPTARG"`;;
a) queryA=`abspath "$OPTARG"`;;
b) queryB=`abspath "$OPTARG"`;;
x) revcompA="--revcompA";;
y) revcompB="--revcompB";;
d) delimiter="$OPTARG";;
i) idRegex="$OPTARG";;
l) idList=`abspath "$OPTARG"`;;
o) outpath=`abspath "$OPTARG"`;;
m) tempdir=`abspath "$OPTARG"`;;
v) verbose="0";;
f) fallback="--singleOK";;
g) noGenericFallback="--nosingleOKgeneric";;
[?]) printErrorHelpAndExit "$badOptionHelp";;
esac
done



if [[ -n "$verbose" ]]
then 
echoerr
echoerr Reference database : $refdb
echoerr Taxonomy file      : $taxonomy
echoerr ID List            : $idList
echoerr Query read A       : $queryA
echoerr RevComp read A     : $revcompA
echoerr Query read B       : $queryB
echoerr RevComp read B     : $revcompB
# echoerr Header Regex : $idRegex
echoerr Read delimiter     : $delimiter
echoerr Output path        : $outpath
echoerr Temporary dir      : $tempdir
echoerr Fallback           : $fallback
echoerr No Generic Fallback: $noGenericFallback
echoerr
fi



if [[ -z "$refdb" || -z "$taxonomy" || -z "$queryA" || -z "$outpath" ]]; then
    printErrorHelpAndExit "-r, -t, -a, and -o options are required"
fi

if [[ -n "$queryB" && -n "$delimiter" ]]; then
    printErrorHelpAndExit "Paired ends can be input using -b or -d, but not both"
fi

# find USEARCH, or hardcode if needed (could use an environment variable...)
# usearch=/path/to/usearch
set +e
usearch=`which usearch`
if [ $? -eq 0 ] 
then
    echoerr using usearch: $usearch
    echoerr
else
    echoerr "usearch not found in path.  Please correct, or just hardcode the path in the rtax script."
    exit 1
fi

set -e

# create a temporary working directory

if [ ! $tempdir ]; then
    tempdir=/tmp/rtax.$RANDOM
    #tempdir=./rtax.$RANDOM
fi
mkdir -p $tempdir
cd $tempdir

# perform the search

if [ $idList ]; then
	idList="--idList $idList"
fi


if [ $delimiter ]; then
    $scripts/splitDelimitedFasta.pl $delimiter $queryA
    queryAa=`basename $queryA`.a
    queryAb=`basename $queryA`.b
    $scripts/rtaxSearch.pl --database $refdb --queryA $queryAa --queryB $queryAb $idList $revcompA $revcompB --usearch $usearch --idRegex "$idRegex" $fallback $noGenericFallback > rawhits
else
    if [ $queryB ]; then
        $scripts/rtaxSearch.pl --database $refdb --queryA $queryA --queryB $queryB $idList $revcompA $revcompB --usearch $usearch --idRegex "$idRegex" $fallback $noGenericFallback > rawhits
    else
        fallback="--singleOK"
        $scripts/rtaxSearch.pl --database $refdb --queryA $queryA $idList $revcompA --usearch $usearch --idRegex "$idRegex" $fallback > rawhits
    fi
fi


# choose the "best" taxonomy assignment by walking down the tree, including info on how many ref sequences agree at each level
# and pipe directly into the cleanup scripts instead of using temp files

$scripts/rtaxVote.pl $taxonomy rawhits | $scripts/classificationQualityFilter.pl 1 0.0 0.5 | $scripts/classificationQualityStripper.pl > $outpath

# The output of this process includes the pcid-width column and a tab-delimited tax string

# filter the taxonomy strings to the finest rank where half of the hits agree (out of all the hits, i.e. including in the denominator those hits with no annotation)
#$scripts/classificationQualityFilter.pl 1 0.0 0.5 < rawtaxonomy > filteredtaxonomy

# clean up the resulting output and provide it to STDOUT
#$scripts/classificationQualityStripper.pl < filteredtaxonomy 

# remove the temp directory
# (comment out to debug or grab intermediate results)

rm -rf $tempdir