#!/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    = "Select sequences based on exact matching of part of a description to NCBI taxonomy IDs or names.\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";

## 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 $negate               = undef;        # -v
my $strict_negate        = undef;        # -s
my $identifier           = undef;        # -i
my $field                = undef;        # -f
my $split_on_regex       = undef;        # -S
my $tax_id_mode          = undef;        #-T
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,
	   'negate|v'                    => \$negate,
	   'strict-negate|s'             => \$strict_negate,
	   'identifier|i'                => \$identifier,
	   '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,
           'q|fastq'                     => sub {$format = 'fastq'},
	  ) 
  or exit(1);

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 \"nodes.dmp\", path to \"names.dmp\", a taxonomic query, and at least one input filename or glob.\n") if (!($fromSTDIN) && (@ARGV < 4));
  pod2usage("$NAME: expects path to \"nodes.dmp\", path to \"names.dmp\" and a taxonomic query when input is on STDIN.\n") if ($fromSTDIN && @ARGV != 3);
}
else {
  pod2usage("$NAME: in --tax-id-mode, expects path to \"nodes.dmp\", a taxonomic query, and at least one input filename or glob.\n") if (!($fromSTDIN) && (@ARGV < 3));
  pod2usage("$NAME: in --tax-id-mode, expects path to \"nodes.dmp\" and a taxonomic query when input is on STDIN.\n") if ($fromSTDIN && @ARGV != 2);
}
&FAST::log($logname, $DATE, $COMMAND, $comment, $fromSTDIN) if ($log); 

my $query; 
my $numquery; 
my $nodesfile;
my $namesfile;
my $line;
my %id;
my %parent;
my %children;
my %match;

unless ($tax_id_mode){
  ($nodesfile,$namesfile,$query) = splice @ARGV,0,3;
}
else {
  ($nodesfile,$query) = splice @ARGV,0,2;
  die "$NAME: in --tax-id-mode with two arguments, the second argument is expected to be an NCBI-taxon-ID and must be an integer. Try $NAME -h for help.\n" if ($query =~ /[^\d]/);
  $numquery = $query;
}

if ($namesfile) {
  die "NCBI-taxonomy-names-file $namesfile cannot be found\n" unless (-e $namesfile);
  $line = `grep \"$query\" $namesfile`;
  die "$NAME: The taxon name $query you entered cannot be found in NCBI-taxonomy-names-file $namesfile\n" unless ($line);
  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 ($name eq $query) {
      $numquery = $id;
    }
  }
  close NAMES;
  die "$NAME: The taxon name $query you entered cannot be found in NCBI-taxonomy-names-file $namesfile\n" unless ($numquery);
}

die "$NAME: NCBI-taxonomy-nodes-file $nodesfile cannot be found\n" unless (-e $nodesfile);
$line = `grep \"$numquery\" $nodesfile`;
die "$NAME: The taxon ID \"$numquery\" cannot be found in NCBI-taxonomy-nodes-file $nodesfile (or your NCBI taxonomy files do not correspond)\n" unless ($line);
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;
  push @{ $children{$parent} },$tax;
}
close NODES;

my @queue;
## here is where we compute the match criterion
push @queue,$numquery;
while (@queue){
  my $first = shift @queue;
  $match{$first} = (($negate || $strict_negate)? 0 : 1);
  push @queue, @{ $children{$first} } if (exists  $children{$first});
}

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


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 $data = $seq->desc();
      $data = $seq->id() if ($identifier);
      if ($field) {
        my @datae = ();
        if($split_on_regex){
	  @datae = split(qr/$split_on_regex/,$data);}
	else{
	  @datae = split $data;
	} #changed to this with more unique variable names, and added $regex_split option herei nstead of above.  	  
	$data = $datae[$index];
      }
      next unless $data;
      if ($tax_id_mode) {
	print $OUT $seq if (($strict_negate and exists $parent{$data} and not exists($match{$data})) or ($negate and not exists($match{$data})) or $match{$data});
      }
      else {
	if($negate and not exists $id{$data}){print $OUT $seq;next;}
	next unless (exists $id{$data});
        print $OUT $seq if (($strict_negate and exists $parent{$id{$data}} and not exists($match{$id{$data}})) or ($negate and not exists($match{$id{$data}})) or $match{$id{$data}});
      }      
    }
    undef $IN;
  }
}

__END__

=head1 NAME

B<fastax> -- select sequence records by NCBI taxonomic names or IDs

=head1 SYNOPSIS

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

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

=head1 DESCRIPTION

B<fastax> takes NCBI Taxonomy data and sequence or alignment data on
input and, if directed to valid NCBI Taxonomic ID labels in the
sequence records, outputs sequence records that belong to a taxonomic
clade as specified by a query argument. For example, a query of
"Alphaproteobacteria" will match data such as "Rickettsiales" and
"Caulobacter crescentus NA1000" but not "Escherichia coli O157:H7" or
"Metazoa"

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
L<ftp://ftp.ncbi.nih.gov/pub/taxonomy>. Only the files "nodes.dmp" and
"names.dmp" are used.

By default, the entire description must exactly match a valid NCBI
taxonomic identifier for a record to match. B<fastax> can optionally
match identifiers or indexed fields in descriptions split by default
on strings of white-space or a user-defined regex.

Options specific to B<fastax>:
  B<-T>, B<--tax-id-mode>              search records by NCBI taxonomic IDs instead of names
  B<-i>, B<--identifier>               search records by sequence identifiers 
  B<-f>, B<--field>=<int>              search records using indexed fields
  B<-S>, B<--split-on-regex>=<regex>   split descriptions or identifiers into fields using <regex>
  B<-v>, B<--negate>                   return records that do *not* match query
  B<-s>, B<--strict-negate>            return records that do *not* match query only if it does match a valid taxon

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>                      create/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<fastax> 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<-v>,
      B<--negate>

return all sequences that are not from the taxon

=item B<-s>,
      B<--strict-negate>

return all sequences that are not from the taxon, but only
if they are from a recognized taxon

=item B<-i>,
      B<--identifier>

taxa are searched over sequence identifiers (default is over descriptions)

=item B<-f [int]>,
      B<--field=[int]>

sequence descriptions or identifers are split into fields
and field <int> (1-based) is searched for taxonomic identifiers

=item B<-S [string]>,
      B<--regex-split=[string]>

in field-mode (B<-f>) split on perl-regex <regex> instead of the default separator,
which is strings of white-space. Do "man perlre" and "man perlfunc" for split.

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

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

=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 sequence that belong to Rhizobiales:

=over 8

B<fastax> -f 3 -S " \| " nodes.dmp names.dmp Rhizobiales t/data/tRNAdb-CE.sample2000.fas

=back

=head1 TODO

=over

=item * Create, store and access persistent NCBI taxonomy data for
easier and faster usage, possibly using Storable and File::Share

=back

=head1 SEE ALSO

Introduction and cookbook for FAST

=over

=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.> and Bioperl I<Stajich et al.>.

=cut
