#!/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 FAST::Bio::PrimarySeq;
use FAST::Bio::Tools::MySeqStats;
use strict;

use vars qw($VERSION $DESC $NAME $COMMAND $DATE);
$VERSION = $FAST::VERSION; 
$DESC    = "";
$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;

## OPTION VARIABLES
my $normalize            = undef; # -n
my $table                = undef; # -t
my $join                 = $def_join_string; #-j
my $format               = $def_format;

my $man                  = undef;  # --man
my $help                 = undef;  # -h
my $moltype              = undef;  # -m, in case bioperl can't tell
my $log                  = undef;        # -l
my $logname              = $def_logname; # -L
my $comment              = undef;        # -C
my $strict               = undef;
my $iupac                = undef;
my $alphabet             = undef;
my $width                = undef;
my $precision            = 3;
my $fastq                = undef;

## TO GO INTO DICTIONARY
my %strict_symbols = ();
my %iupac_symbols = ();
$strict_symbols{'dna'}     = [split //,'AGCT'];
$iupac_symbols{'dna'}      = [split //,'AGCTRYSWKMBHDVXN'];
$strict_symbols{'rna'}     = [split //,'AGCU'];
$iupac_symbols{'rna'}      = [split //,'AGCTRYSWKMBHDVXN'];
$strict_symbols{'protein'} = [split //,'ACDEFGHIKLMNPQRSTVWY*'];
$iupac_symbols{'protein'}  = [split //,'ACDEFGHIKLMNPQRSTVWYBZX*'];

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;
						},
	   'log|l'                       => \$log,
	   'logname|L=s'                 => \$logname,
           'format=s'                    => \$format,
	   'comment|C=s'                 => \$comment,
	   'normalize|n'                 => \$normalize,
	   'table|t'                     => \$table,
	   'strict|s'                    => \$strict,
	   'iupac|i'                     => \$iupac,
	   'join|j=s'                    => \$join,
	   'alphabet|a=s'                => \$alphabet,
	   'width|w=i'                   => \$width,
	   'precision|p=i'               => \$precision,
           'fastq|q'                     => sub{$format = 'fastq';},
	  )
  or pod2usage(2);
	
$join = "\t" if ($join eq '\t');

pod2usage(-verbose => 1) if $help;
pod2usage(-verbose => 2) if $man;
my $fromSTDIN = ((-t STDIN) ? false : true);
pod2usage("$NAME: Requires at least one argument FILE [FILE2…FILEN] unless input from STDIN.\n") if (!($fromSTDIN) && (@ARGV == 0));
pod2usage("$NAME: Requires exactly zero arguments if input is from STDIN.\n") if ($fromSTDIN && (@ARGV != 0));

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

my $alpha_keys;
if ($alphabet) {
  $alpha_keys = [ split //,$alphabet ];
}

if ($normalize && !$width) {
  $width = $precision + 3;
}
elsif (!$width) {
  $width = 4;
}


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);
  }
}

my $sum = {};
my $grandtotal;
my $firstseq = 1;
my ($warned,$job_alphabet);
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 $seqobj = $IN->next_seq()) {
      if ($firstseq) {
	$job_alphabet = $moltype || $seqobj->alphabet;
	undef $firstseq;
	if ($table && ($iupac || $strict || $alphabet)){
	  my $keys = &keys();
	  print join "", map {sprintf "%${width}s",$_} @$keys;
	  print "   total" unless ($normalize);
	  print "\n";
	}
      }
      my ($freq,$total,$sum) = &process ($seqobj,$sum);
      $grandtotal += $total;
      &print ($freq,$total,$seqobj)
    }
    undef $IN;
  }
}
$table and do {
  if ($normalize){
    my $normalizer;
    my $keys = &keys($sum);
    foreach my $key (@$keys) {
      $normalizer += $$sum{$key};
    }
    if ($normalizer > 0) {
      foreach my $key (@$keys) {
	$$sum{$key} /= $normalizer;
      }
    }
    &print($sum,undef,undef);
  }
  else {
    &print($sum,$grandtotal,undef);
  }
};


#-----------------------
sub process {
#-----------------------
  my ($seqobj, $sum) = @_;
  
  if ($moltype){
    $seqobj->alphabet($moltype);
  }
  unless ($warned) {
    $job_alphabet and do {
      if ($seqobj->alphabet ne $job_alphabet) {
        warn "$NAME: input sequences of different alphabets.\n";
      }
      $warned = 1;
    };
  }  
  
  my $seqstats = FAST::Bio::Tools::MySeqStats->new($seqobj);
  my $freq = $seqstats->count_monomers();
  my $keys = &keys($freq);
  my $normalizer = 0;
  foreach my $key (@$keys) {
    $$sum{$key} += ( exists $$freq{$key} ? $$freq{$key} : 0);
    $normalizer += ( exists $$freq{$key} ? $$freq{$key} : 0);
  }
  if ($normalize && $normalizer > 0) {
    foreach my $key (@$keys) {
      if (exists $$freq{$key} ) {
	$$freq{$key} /= $normalizer;
      }
    }
  }
  return ($freq,$normalizer,$sum);
}

#-----------------------
sub keys {
#-----------------------
  my $freq = shift;
  my $keys;
  if ($alphabet) {
    $keys = $alpha_keys;
  }
  elsif ($strict) {
    $keys = $strict_symbols{$job_alphabet};
  }
  elsif ($iupac) {
    $keys = $iupac_symbols{$job_alphabet};
  }
  else {
    $keys = [ sort keys %$freq ];
  }
  return $keys;
}


#-----------------------
sub print {
#-----------------------
  my ($freq,$total,$obj) = @_;
  my $keys = &keys($freq);
  my $label;
  if ($obj) {
    $label = $obj->id() 
  }
  else {
    $label = "# ALL DATA";
  }

  if ($table){
    if ($iupac || $strict || $alphabet) {
      if ($normalize) {
	print join "",map {sprintf "%${width}.${precision}f",(exists $$freq{$_} ? $$freq{$_} : 0)} @$keys;
      }
      else {
	print join "",map {sprintf "%${width}d",  (exists $$freq{$_} ? $$freq{$_} : 0)} @$keys;
	printf "%8d",$total;
      }
    }
    else {
      if ($normalize) {
	print join " ",map {sprintf "%s:%-${width}.${precision}f",$_,(exists $$freq{$_} ? $$freq{$_} : 0)} @$keys;
      }
      else {
	print join " ",map {sprintf "%s:%-${width}d",  $_,(exists $$freq{$_} ? $$freq{$_} : 0)} @$keys;
	printf " total:%-${width}d",$total;
      }
    }
    printf "  %-${width}s\n",$label;
  }
   else{  
     my $olddesc = $obj->desc();
     my $compstring;
     if ($normalize) {
       $compstring = join $join,map { (sprintf "comp_%s:%.${precision}f",$_,(exists $$freq{$_} ? $$freq{$_} : 0)) } @$keys;
     }
     else{
       $compstring = join $join,map { (sprintf "comp_%s:%d",$_,(exists $$freq{$_} ? $$freq{$_} : 0))   } @$keys;
     }
     $obj->desc(join $join,$olddesc,$compstring);
     print $OUT $obj;
   }
}
__END__

=head1 NAME

B<fascomp> -- analysis of monomer frequencies 

=head1 SYNOPSIS

B<fascomp>   [OPTION]... [MULTIFASTA-FILE...]

=head1 DESCRIPTION

B<fascomp> takes multifasta format sequence or alignment data as
input, and counts the absolute or relative frequencies of monomers in
each sequence individually and for all data on input. By default,
absolute frequencies are computed and appended to the
description. Optionally, B<fascomp --table> will output a frequency
table to STDOUT as tagged values, both for individual sequences and
for all data on input. In table-mode if the character alphabet of the
job is determined by either setting B<--moltype> or from the type of
the first sequence on input, and either B<--strict>, B<--iupac> or
B<--alphabet> is set, header labels will be output as well. 

Options specific to B<fascomp>:
  B<-n>, B<--normalize>               compute relative frequencies
  B<-t>, B<--table>                   table-mode, output a table to STDOUT 
  B<-j>, B<--join>=<string>           use <string> to join tagged values in descriptions
  B<-s>, B<--strict>                  output moltype-dependent character header for table-mode
  B<-i>, B<--iupac>                   output character header including ambiguities for table-mode
  B<-a>, B<--alphabet>=<string>       tally only characters if they are in the set <string>, as in "ACGT-"
  B<-p>, B<--precision>=<int>         print relative frequencies with <int> digits after the decimal point 
  B<-w>, B<--width>=<int>             print frequencies in fields of width <int>

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<fascomp> 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<-n>
      B<--normalize>

Compute relative frequencies.

=item B<-t>
      B<--table>

Output a table to STDOUT. 

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

Use <string> to join tagged values in descriptions. Use "\t" to
indicate a tab-character.

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

Output moltype-dependent character header for table-mode.

=item B<-i>
      B<--iupac>

Output character header including ambiguities for table-mode.

=item B<-a>
      B<--alphabet>=<string>

Tally only characters if they are in the set <string>, as in "ACGT-".

=item B<-p> 
      B<--precision>=<int>

Print relative frequencies with <int> digits after the decimal point.

=item B<-w> 
      B<--width>=<int>

Print frequencies in fields of width <int>

=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

Compute and annotate description with normalized base frequencies: 

=over 8

B<fascomp> -n t/data/P450.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.> and Bioperl I<Stajich et al.>.

=cut
