#!/usr/bin/perl -w
use POSIX;
use Getopt::Long qw(:config bundling require_order auto_version);
use Pod::Usage;
use FAST;
use FAST::Bio::SeqIO;
use strict;

use vars qw($VERSION $DESC $NAME $COMMAND $DATE);
$VERSION = $FAST::VERSION; 
$DESC    = "Sort and identify sequences based on NCBI taxonomy.\n";
$NAME    = $0;
$NAME    =~ s/^.*\///;
$COMMAND = join " ",$NAME,@ARGV;
$DATE = POSIX::strftime("%c",localtime());

use constant { true => 1, false => 0 };

## DEFAULT OPTION VALUES
my $def_format  = $FAST::DEF_FORMAT;  #7/1/13 "fasta";
my $def_logname = $FAST::DEF_LOGNAME; #7/1/13 "FAST.log.txt";
my $def_join_string = $FAST::DEF_JOIN_STRING;
my $def_tag  = "taxaddress:";

## OPTION VARIABLES
my $man                  = undef;  # --man
my $help                 = undef;  # -h
my $moltype              = undef;  # -m, in case bioperl can't tell
my $format               = $def_format;  # --format
my $log                  = undef;        # -l
my $logname              = $def_logname; # -L
my $comment              = undef;        # -C
my $tax_id_mode          = undef; #-T
my $split_on_regex       = undef; #-S
my $field                = undef; #-f
my $identifier           = undef; #-i
my $index                = undef; #--index
my $tag                  = $def_tag;  # -m
my $annotate             = undef;
my $join                 = $def_join_string; # -j
my $fastq                = undef;

GetOptions('help|h'         		 => \$help, 
	   'man'            		 => \$man,
	   'moltype|m=s'                 => sub{  my (undef,$val) = @_; 
						  die "$NAME: --moltype or -m option argument must be \"dna\", \"rna\" or \"protein\"" 
						    unless $val =~ /dna|rna|protein/i; 
						  $moltype = $val;
						},
	   'format=s'                    => \$format,
	   'log|l'                       => \$log,
	   'logname|L=s'                 => \$logname,
	   'comment|C=s'                 => \$comment,
	   'field|f=i'                   => sub{  my (undef,$val) = @_; 
						  die "$NAME: --field or -f option expects non-zero integer argument\n" 
						    unless $val != 0; 
						  $field = $val;
						},
	   'split-on-regex|S=s'          => \$split_on_regex,
	   'tax-id-mode|T'               => \$tax_id_mode,
	   'identifier|i'                => \$identifier,
	   'index'                       => \$index,
	   'annotate|a'                  => \$annotate,
	   'join|j=s'                    => \$join,
           'q|fastq'                     => sub {$format = 'fastq'},
	  ) 
  or exit(1);

$join = "\t" if ($join eq '\t');

pod2usage(-verbose => 1) if $help;
pod2usage(-verbose => 2) if $man;
my $fromSTDIN = ((-t STDIN) ? false : true);

unless ($tax_id_mode) {
  pod2usage("$NAME: expects path to NCBI \"nodes.dmp\" file, path to NCBI \"names.dmp\" file, and at least one input filename or glob.\n") if (!($fromSTDIN) && (@ARGV < 3));
  pod2usage("$NAME: expects path to NCBI \"nodes.dmp\" file and path to NCBI \"names.dmp\" file when input is on STDIN.\n") if ($fromSTDIN && @ARGV != 2);
}
else {
  pod2usage("$NAME: in tax-id-mode, expects path to NCBI \"nodes.dmp\" file and at least one input filename or glob.\n") if (!($fromSTDIN) && (@ARGV < 2));
  pod2usage("$NAME: in tax-id-mode, expects path to NCBI \"nodes.dmp\" file when input is on STDIN.\n") if ($fromSTDIN && @ARGV != 1);
}

&FAST::log($logname, $DATE, $COMMAND, $comment, $fromSTDIN) if ($log); 

my $nodesfile;
my $namesfile;
my %id;
my %name;
my %parent;
my %children;
my %address;
my %childaddress;
my %childrenseen;
my %seq;

unless ($tax_id_mode){
  ($nodesfile,$namesfile) = splice @ARGV,0,2;
}
else {
  ($nodesfile) = shift @ARGV;
}

if ($namesfile) {
  die "$NAME: NCBI-taxonomy-names-file $namesfile cannot be found\n" unless (-e $namesfile);
  open (NAMES,$namesfile) or die "Can't open NCBI-taxonomy-names-file $namesfile\n";
  while (<NAMES>) {
    my ($id,$name,@stuff) = split /\t\|\t/,$_;
    $id{$name} = $id;
    if ($index) {
      unless (exists $name{$id}) { ## take the first name from names.dmp
	$name{$id} = $name;
      }
    }
  }
  close NAMES;
}

die "$NAME: NCBI-taxonomy-nodes-file $nodesfile cannot be found\n" unless (-e $nodesfile);
open (NODES,$nodesfile) or die "Can't open NCBI-taxonomy-nodes-file $nodesfile\n";
while (<NODES>) {
  my ($tax,$parent,@stuff) = split /\t\|\t/,$_;
  $parent{$tax} = $parent;
  unless ($parent == $tax) { ## both equal to one at top of nodes.dmp
    $children{$parent}{$tax} = 1;
  }
}
close NODES;

$address{ 1 } = "00";
if ($index) {
  print $tag,join " ",$address{ 1 },($tax_id_mode ? 1 : $name{1}),"\n";
}
&address_DFS(1);
sub address_DFS(){
  my $tax = shift;
  $childaddress{$tax} = 0;
  foreach my $child (sort {$a <=> $b} keys %{ $children{$tax} } ) { 
    $address{ $child } = join "",(join ".",$address{ $tax },(sprintf "%02X",$childaddress{ $tax }++));
    ## overflow is no problem. If there are more than 256 children, childaddress will print with 3 or more figures.
    if ($index) {
      print $tag,join " ",$address{ $child },($tax_id_mode ? $child : $name{$child}),"\n";
    }
    if (exists $children{$child}) {  
      &address_DFS($child);
    }
  }
}

die if $index;

my $field_index;
if (defined $field and $field > 0) {
  $field_index = $field - 1;
}
else { # $field < 0
  $field_index = $field;
}

my $re;
if ($split_on_regex){
  $re = qr/$split_on_regex/;
}


my $OUT = FAST::Bio::SeqIO->newFh('-format' => $format);
my $IN;
unless (@ARGV) {
    if ($moltype) {
	$IN = FAST::Bio::SeqIO->new(-fh => *STDIN{IO}, '-format' => $format, '-alphabet' => $moltype);
    }
    else {
	$IN = FAST::Bio::SeqIO->new(-fh => *STDIN{IO}, '-format' => $format);
    }
}

while ($IN or @ARGV) {
  if (@ARGV) {
    my $file = shift (@ARGV);
    unless (-e $file) {
      warn "$NAME: Could not find file $file. Skipping.\n";
      next;
    }
    elsif ($moltype) {
      $IN = FAST::Bio::SeqIO->new(-file => $file, '-format' => $format, '-alphabet' => $moltype);
    }
    else {
      $IN = FAST::Bio::SeqIO->new(-file => $file, '-format' => $format);
    }
  }
  if ($IN) { 
    while (my $seq = $IN->next_seq()) {
      my $id;
      my $data = $seq->desc();
      $data = $seq->id() if ($identifier);
      if ($field) {
	my @data = ();
	if ($split_on_regex) {
	  @data = split $re,$data;
	}
	else {
	  @data = split ' ',$data;
	}
	$data = $data[$field_index];
      }
      next unless $data;
  
      if ($tax_id_mode) {
	push @{ $seq{$data} },$seq;
	$id = $data;
      }
      else {
	next unless (exists $id{$data});
	push @{ $seq{$id{$data}} },$seq;
	$id = $id{$data};
      }
      while($id != 1) {
	my $parent = $parent{$id};
	$childrenseen{$parent}{$id} = 1; 
	$id = $parent;
      }
    }
    undef $IN;
  }
}

&DFS_seq(1);

sub DFS_seq {
  my $node = shift;
  if (exists $seq{$node} ) {
    foreach my $seq (@{ $seq{$node} }){
      if ($annotate) {
	$seq->desc(join $join,$seq->desc,(join "",$tag,$address{$node}));
      }
      print $OUT $seq;      

    }
  }
  if (exists $childrenseen{$node} ) {
    foreach my $child (map { $_->[0] } sort address_sort map {[$_,$address{$_}]} keys %{ $childrenseen{$node} }) {
      &DFS_seq($child);
    }
  }
}

sub address_sort {
  return 0 if ($a->[1] eq $b->[1]);
  my @a = split /\./, $a->[1];
  my @b = split /\./, $b->[1];
  while (@a) {
    my $aa = shift @a or return -1;
    my $bb = shift @b or return  1;
    return -1 if ($aa < $bb);
    return  1 if ($aa > $bb);
  }
  warn "${NAME}::address_sort: invalid condition. Please contact the author.\n";
}

__END__

=head1 NAME

B<fastaxsort> - sort sequence records based on NCBI taxonomy

=head1 SYNOPSIS

B<fastaxsort>  [OPTION]...  [NODES-FILE] [NAMES-FILE] [MULTIFASTA-FILE]...

B<fastaxsort>  --tax-id-mode [OPTION]... [NODES-FILE] [MULTIFASTA-FILE]...

=head1 DESCRIPTION

B<fastaxsort> takes NCBI Taxonomy data and multifasta format sequence
or alignment data as input and, if directed to valid NCBI Taxonomic
names or IDs in the sequence records, outputs those records
sorted taxonomically, or more specifically, sorted by pre-order
depth-first traversal through the NCBI Taxonomy tree.

NCBI Taxonomy data must be downloaded separately from <NCBI
Taxonomy|http://www.ncbi.nlm.nih.gov/taxonomy>, particularly one of
the files marked "taxdump" from for example
<ftp://ftp.ncbi.nih.gov/pub/taxonomy>. Only the files "nodes.dmp" and
"names.dmp" in the downloaded data are used.

At least some part of sequence records must contain NCBI taxonomic
names or IDs. By default, the entire description is expected to
exactly match exactly one NCBI taxonomic name (or ID in
--tax-id-mode). B<fastaxsort> can optionally sort sequence records
taxonomically by their identifiers, or by indexed fields within
descriptions or identifiers, where fields are generated by splitting
with a delimiter; by default, one or more white-space
characters. Alternative delimiters may be specified by a user-defined
regex. Positive integers index fields from the beginning; the first
field has index one. Negative integers index fields from the end.

Options specific to B<fastaxsort>:
  B<-T>, B<--tax-id-mode>              sort records using NCBI taxonomic IDs in data
  B<-i>, B<--identifier>               sort records using sequence identifiers
  B<-f>, B<--field>=<int>              sort records using fields
  B<-S>, B<--split-on-regex>=<regex>   split descriptions or identifiers using regex 
  B<-a>, B<--annotate>                 annotate records with a dot-hex taxonomic address
  B<-j>, B<--join>=<string>            use <string> to join taxonomic addresses to descriptions 
  B<--index>                        output an index mapping dot-hex addresses to NCBI Taxonomy

Options general to FAST:
  B<-h>, B<--help>                     print a brief help message
  B<--man>             	        print full documentation
  B<--version>                      print version
  B<-l>, B<--log>                         qcreate/append to logfile	
  B<-L>, B<--logname>=<string>            use logfile name <string>
  B<-C>, B<--comment>=<string>            save comment <string> to log
  B<--format>=<format>                 use alternative format for input  
  B<--moltype>=<[dna|rna|protein]>      specify input sequence type
  B<-q>, B<--fastq>                        use fastq format as input and output

=head1 INPUT AND OUTPUT

B<fastaxsort> is part of FAST, the FAST Analysis of Sequences Toolbox, based
on Bioperl. Most core FAST utilities expect input and return output in
multifasta format. Input can occur in one or more files or on
STDIN. Output occurs to STDOUT. The FAST utility B<fasconvert> can
reformat other formats to and from multifasta.

=head1 OPTIONS

=over 8

=item B<-T>
      B<--tax-id-mode>

NCBI Taxonomic data in sequence records are IDs, not names.

=item B<-i>
      B<--indentifier>

Taxa are sorted using sequence identifiers (default uses whole descriptions)

=item B<-f>
      B<--field>

Sort sequence records by values at a specific field in sequence
descriptions or identifiers.  With this option, the description or
identifier is split into fields using strings of white space as field
delimiters (the default Perl delimiter for splitting lines of data
into fields, which are invalid characters in sequence identfiers).

This option takes a mandatory integer option argument giving the index
for which field to sort by. One-based indexing is used, so the first
field after the sequence identifier has index 1. As standard in Perl,
negative indices count backwards from the last field; field "-1" is
the last field, "-2" is the second-to-last etc. Sequence records for
which the specified field does not exist will sort on a null key.

=item B<-S>
      B<--split-on-regex>

Use regex <regex> to split descriptions/identifiers for the -f option
instead of the perl default (which splits on one or more whitespace
characters). Special characters must be quoted to protect them from
the shell.

=item B<-a>
      B<--annotate>

Add FAST taxonomic addresses in dot-hex notation to sequence record descriptions

=item B<-j [string]>
      B<--join=[strong]>

Use [string] to append FAST taxonomic addresses to sequence record
descriptions, instead of default " ". Use "\t" to indicate a
tab-character.

=item B<--index>

Instead of printing the sorted sequence records, print a key that maps
B<fastaxsort> taxonomically generated taxonomic addresses in
dot-hexadecimal notation to NCBI taxonomic names or IDs.


=item B<-h>,
      B<--help>

Print a brief help message and exit.

=item B<--man>

Print the manual page and exit.

=item B<--version>

Print version information and exit.

=item B<-l>,
      B<--log>

Creates, or appends to, a generic FAST logfile in the current working
directory. The logfile records date/time of execution, full command
with options and arguments, and an optional comment.

=item B<-L [string]>,
      B<--logname=[string]>

Use [string] as the name of the logfile. Default is "FAST.log.txt".

=item B<-C [string]>,
      B<--comment=[string]>

Include comment [string] in logfile. No comment is saved by default.

=item B<--format=[format]> 		  

Use alternative format for input. See man page for "fasconvert" for
allowed formats. This is for convenience; the FAST tools are designed
to exchange data in Fasta format, and "fasta" is the default format
for this tool.

=item B<-m [dna|rna|protein]>,
      B<--moltype=[dna|rna|protein]> 		  

Specify the type of sequence on input (should not be needed in most
cases, but sometimes Bioperl cannot guess and complains when
processing data).

=item B<-q>
      B<--fastq>

Use fastq format as input and output.

=back

=head1 EXAMPLES

Print all sequences with "-DNA" in the ID:

=over 8

Sort sequences where the taxonomic identifier is found in the third field of the description:

B<fastaxsort> -f 3 -S " \| " nodes.dmp names.dmp tRNAdb-CE.sample2000.fas

=back


=head1 SEE ALSO

=over 8

=item C<man perlre>

=item C<perldoc perlre>

Documentation on perl regular expressions.

=item C<man FAST>

=item C<perldoc FAST>

Introduction and cookbook for FAST

=item L<The FAST Home Page|http://compbio.ucmerced.edu/ardell/FAST>"

=back 

=head1 CITING

If you use FAST, please cite I<Lawrence et al. (2015). FAST: FAST Analysis of
Sequences Toolbox. Bioinformatics> and Bioperl I<Stajich et al.>.

=cut
