#!/bin/bash

#
# gblocks2nexus.sh -- Shell wrapper for the Gblocks multiple sequence
#   alignment cleaning program
#
# (C) 2008 by
#     Markus Goeker (support [AT] goeker [DOT] org)
#
# This program is distributed under the terms of the Gnu Public License V2.
# For further information, see http://www.gnu.org/licenses/gpl.html
#
# If you use this script in a publication, please cite the web page
# at http://www.goeker.org/mg/aop/
#

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

# Chat Ramey's recommendations for secure shell scripts
IFS=$' \t\n'
unset -f unalias
\unalias -a
unset -f command
SYSPATH="$(command -p getconf PATH 2>/dev/null)"
if [ -z "$SYSPATH" ]; then
  SYSPATH="/usr/bin:/bin"
fi
PATH="$SYSPATH:$PATH"

set -eu
declare -r GBLOCKS=/usr/local/bin/Gblocks

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

# Convert Gblocks mask file output to NEXUS format
function gblocks2nexus
{
  awk '
    /^>/ {
      if ($0 ~ /#Mask$/)
        sequences[count] = clean(seq)
      else if ($0 != ">P1;Gblocks")
        labels[++count] = substr($0, 5)
      seq = ""
      next
    }
    {
      seq = (seq $0)
    }
    END {
      # Get exclusions status
      seq = clean(seq)
      np = split(seq, pos, "")
      for (i = 1; i <= np; i++) {
        current = pos[i]
        if (current == ".") {
          if (last == "#" || i == 1)
            exclude = (exclude " " i "-")
          else if (i == np)
            exclude = (exclude i)
        } else if (last == ".") 
          exclude = (exclude i - 1)
        last = current
      }
      # Print results
      print "#NEXUS"
      print ""
      print "begin data;"
      printf "   dimensions ntax=%g nchar=%g;\n", count, length(sequences[1])
      print "   format datatype=dna interleave=no missing=? gap=-;"
      print "   matrix"
      for (i = 1; i <= count; i++)
        maxlen = max(maxlen, length(labels[i]))
      format = ("%-" maxlen + 2 "s %s\n")
      for (i = 1; i <= count; i++)
        printf format, ("'"'"'" labels[i] "'"'"'"), sequences[i]
      printf format, "[", (seq "]")
      print "   ;"
      print "end;"
      if (exclude) {
        print "begin assumptions;"
        printf "   exset * gblocks =%s;\n", exclude
        print "end;"
      }
      print ""
    }
    function clean(str) {
      gsub(/[^A-Za-z@?#.-]+/, "", str)
      return str
    }
    function max(x, y) {
      return x > y ? x : y
    }
  ' "$@"
}

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

helpmsg=
gapmode=half
seqtype=dna

while getopts "g:hs:" opt; do
  case $opt in
    g ) gapmode=$OPTARG;;
    h ) helpmsg=yes;;
    s ) seqtype=$OPTARG;;
    * ) exit 1;;
  esac
done
shift $(($OPTIND - 1))

case $gapmode in
  [AaFf]* ) gapmode=f;;
  [Hh]*   ) gapmode=h;;
  [Nn]*   ) gapmode=n;;
  * )
    echo "Invalid gap mode: \"$gapmode\"" >&2
    exit 1
  ;;
esac

case $seqtype in
  [Cc]*   ) seqtype=c;;
  [DdNn]* ) seqtype=d;;
  [PpAa]* ) seqtype=p;;
  * )
    echo "Invalid sequence type: \"$seqtype\"" >&2
    exit 1
  ;;
esac

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

if [ $helpmsg ] || [ $# -eq 0 ]; then
  cat >&2 <<-__EOF

	${0##*/} creates Gblocks-cleaned NEXUS files

	Usage: ${0##*/} [options] fastafile(s)

	Options:
	  -g Gap mode (full/half/none).
	  -h Print this message.
	  -s Sequence type (DNA/Protein/Codons).

__EOF
  exit 1
fi

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

run=$RANDOM$RANDOM
run=${run:1:4}

for i; do
  cfasfile="$i-$run"
  maskfile="${cfasfile}Mask"
  "$GBLOCKS" "$i" -t="$seqtype" -b5="$gapmode" -k=y -e=-$run \
    -p=n > /dev/null || :
  if [ -s "$maskfile" ]; then
    gblocks2nexus "$maskfile" > ${i%.*}_gblocks.nex
    mv "$cfasfile" ${i%.*}_gblocks.fas
    rm -f "$maskfile"
  else
    echo "WARNING: converting \"$i\" did not result." >&2
    rm -f "$maskfile" "$cfasfile"
  fi
done