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

use vars qw($VERSION $DESC $NAME $COMMAND $DATE);
$VERSION = $FAST::VERSION; 
$DESC    = "Translates gapped and ungapped sequences and alignments.";
$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;  # "fasta";
my $def_logname = $FAST::DEF_LOGNAME; # "FAST.log.txt";
my $def_join_string = $FAST::DEF_JOIN_STRING;

my $def_stop_char = '*'; # -s
my $def_unk_char  = 'X'; # -u
my $def_frame     = 0;   # -F
my $def_table     = 1;   # -t // TABLE_ID (CODE)

## OPTION VARIABLES
my $man                  = undef;  # --man
my $help                 = undef;  # -h
my $tables               = undef;  # 

my $format               = $def_format;  # -f
my $log                  = undef;        # -l
my $logname              = $def_logname; # -L
my $comment              = undef;        # -C

my $annotate             = undef;        # -a
my $join                 = $def_join_string; # -j

my $brand                = undef;        # -b

my $gapped               = undef;        # -g, gapped one-letter aa IUPAC transalation, aligns to codons
my $keep                 = undef;        # -k, keep input sequences on output
my $xl_as_cds            = undef;        # --cds, force all translations to start with M
my $codon2aa             = undef;        # -o, turn codon aln to aa aln
my $moltype              = undef;        # -m, force moltype

my $stop                 = $def_stop_char; # -s
my $unknown              = $def_unk_char;  # -u
my $frame                = $def_frame;     # -F
my $table                = $def_table;     # -t

## in increasing priority:
my $all_three            = undef;        # -3
my $all_six              = undef;        # -6

GetOptions('help|h'         		 => \$help, 
	   'man'            		 => \$man,
	   'tables|codes'                => \$tables,
	   'format=s'                    => \$format,
	   'log|l'                       => \$log,
	   'logname|L=s'                 => \$logname,
	   'comment|C=s'                 => \$comment,
	   'annotate|a'                  => \$annotate,
	   'join|j=s'                    => \$join,
	   'moltype|m=s'                 => sub{  my (undef,$val) = @_; 
						  die "$NAME: --moltype or -m option argument must be \"dna\", \"rna\"" 
						    unless $val =~ /dna|rna/i; 
						  $moltype = $val;
						},
           'stop|s=s'                    => sub{  my (undef,$val) = @_; 
		   				  die "$NAME: --stop or -s option takes a single character argument in POSIX class [:print:]." 
						    unless $val =~ /\A[[:print:]]\Z/;
						  $stop = $val;
						},
           'unknown|u=s'                 => sub{  my (undef,$val) = @_; 
		   				  die "$NAME: --unknown or -u option takes a single character argument in POSIX class [:print:]." 
						    unless $val =~ /\A[[:print:]]\Z/;
						  $unknown = $val;
						},	   
           'frame|f=i'                   => sub{  my (undef,$val) = @_; 
		   				  die "$NAME: --frame or -f option takes an integer argument between 0 and 2 inclusively." 
						    unless $val >= 0 and $val <= 2;
						  $frame = $val;
						},
	   'code|c|table|t=i'            => sub{  my (undef,$val) = @_; 
						  die "$NAME: --table or -t option takes an integer argument, an NCBI genetic code table index in the set {1..6,9..16,21..23}." 
						    unless $val >= 1 and $val <= 6 or $val >= 9 and $val <= 16 or $val >= 21 and $val <= 23;
						  $table = $val;
						},
	   'all-three|3'                 => \$all_three,
	   'all-six|6'                   => \$all_six,	   
	   'gapped|g'                    => \$gapped,
	   'keep|k'                      => \$keep,
	   'CDS'                         => \$xl_as_cds,
	   'codon2aa|o'                  => \$codon2aa,
           'fastq|q'                     => sub { $format = 'fastq'; },
	  )
  or pod2usage(2);
		  
pod2usage(-verbose => 1) if $help;
pod2usage(-verbose => 2) if $man;

if ($tables) {
  my $tables = FAST::Bio::Tools::CodonTable->tables;
  my ($id, $name);
  foreach $id (sort {$a <=> $b} sort keys %$tables) {
    $name = $$tables{$id};
    print "$id = $name\n";
  }
  die;
}

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));
pod2usage("$NAME: Options --annotate/-a and --keep/-k are incompatible.\n") if ($annotate && $keep);

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

$keep and do {
  $gapped = true;
};

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

my $OUT = FAST::Bio::SeqIO->newFh('-format' => 'fasta');
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()) {
      $keep and do {
	print $OUT $seq;
      };
      if ($all_six) {
	&xl_3_frames($seq,$seq,false);
	my $rc = $seq->revcom;
	my $id = join '-',$seq->display_id,'rc';
	$rc->display_id($id);
	$keep and do {
	  print $OUT $rc;
	};
	&xl_3_frames($seq,$rc,true);
      }
      elsif ($all_three) {
	&xl_3_frames($seq,$seq,false);
      }
      else {
	&xl ($seq,$seq,$frame,false);
      }
      print $OUT $seq if ($annotate);
    }
    undef $IN;
  }
}


#-----------------------
sub xl_3_frames {
#-----------------------
  my $seq = shift;
  my $xlseq = shift;
  my $rc  = shift;
  for ($frame=0;$frame<3;$frame++) {
    &xl ($seq,$xlseq,$frame,$rc);
  }
}

#-----------------------
sub xl {
#-----------------------
  my $seq = shift;
  my $xlseq = shift;
  my $frame = shift;
  my $rc = shift;

  my $rcprint = ($rc ? "rc-" : "");

  my $aa;
  if ($gapped) {
    my $gapseq = FAST::Bio::GapSeq->copy($xlseq);
    $gapseq->alphabet($seq->alphabet());
    $aa = $gapseq->translate($stop,$unknown,$frame,$table,$xl_as_cds); 
  }
  elsif ($codon2aa) {
    my @aaseq = ();
    my $seqseq = $xlseq->seq;
    my @seqparts = split /(-+)/,$seqseq;
    foreach my $part (@seqparts) {
      if ($part =~ /-/) {
	$part =~ s/---/-/g;
	push @aaseq,$part;
      }
      else {
	my $seq = FAST::Bio::Seq->new(-seq => $part, -alphabet => $xlseq->alphabet());
	my $as = $seq->translate($stop,$unknown,$frame,$table,undef); ## handle cds 
	push @aaseq, $as->seq;
      }
    }
    my $aaseq = join '',('-' x $frame),@aaseq;
    $aa = FAST::Bio::Seq->new(-seq => $aaseq, -alphabet => $seq->alphabet());
  }
  else {
    if ($xlseq->seq() =~ /-/) {
      die "$NAME: expects gap-free input unless -g or -o options are used.\n";
    }
    $aa = $xlseq->translate($stop,$unknown,$frame,$table,$xl_as_cds); 
    my $aaFix = join '',('-' x $frame),$aa->seq;
    $aa->seq($aaFix);
  }
  if ($annotate) {
    my $olddesc = $seq->desc();
    $seq->desc(join $join,$olddesc,(join "",$rcprint,"xl$frame:",$aa->seq())); 
  }
  else {
    my $id = join '-',$seq->display_id,(join "",$rcprint,"xl$frame");   
    $aa->display_id("$id");
    print $OUT $aa; 
  }
}


__END__

=head1 NAME

B<fasxl> -- translate sequences by a genetic code.

=head1 SYNOPSIS

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

=head1 DESCRIPTION

B<fasxl> takes multifasta format DNA or RNA codon sequences or
alignments as input, and generates biological translations of those
sequences as output. Gapped sequences on input are allowed with the
-g, -k or -o options; the gap character "-" is required for correct
interpretation.

Options specific to B<fasxl>:
  B<-a>, B<--annotate>          output translations as tagged values in descriptions 
  B<-j>, B<--join>=<string>     use <string> to join data 
  B<-g>, B<--gapped>            gapped translations in alignment with input
  B<-k>, B<--keep>              keep input sequences in output (forces -g).
  B<-o>, B<--codon2aa>          turn codon alignment into a protein alignment 
                           	(best used for codon-based alignments, where gaps are mod 3 
                            	length and seqs are to be translated in frame 0)
  B<--CDS>                   treat as CDS (for bacterial code, for example, treat init codons as M)(deprecated) 
  B<-f>, B<--frame>=<int>       frame for translation [0 (default),1, or 2].
  B<-c>, B<-t>, B<--code>=<int> NCBI genetic code tableID for translating sequences [1]
  B<-3>                      translate each sequence in all three forward frames
  B<-6>                      translate each sequence in all six frames
  B<-s>, B<--stop>=<char>       character representing stop codons ["*"]
  B<-u>, B<--unknown>=<char>    character representing unknown amino acids ["X"]
  B<--codes>, B<--tables>       print NCBI table ids of genetic codes for -t option

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<-m>, -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<fasxl> 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<-a>,
      B<--annotate>

This option outputs translations as tagged values in descriptions.
By default translations are outputted in the sequence.

=item B<-j>,
      B<--join>=<string>

Use a <string> to join data in the description. Default is a space
character. Use "\t" to indicate a tab-character.

=item B<-g>,
      B<--gapped>

This option enables translation of gapped sequences.

=item B<-k>,
      B<--keep>

Keep inputted sequences in the output. This option also enables the
translation of gapped sequences.

=item B<-o>,
      B<--codon2aa>

turn codon alignment into a protein alignment 
(best used for codon-based alignments, where gaps are mod 3 
length and seqs are to be translated in frame 0)

=item B<-c>,
      B<--cds>

This option is deprecated.
Treat as CDS (for bacterial code, for example, treat init codons as M).

=item B<-f>,
      B<--frame>=<int>

Specify the frame for translation [0,1,2]. By default frame 0 is used.

=item B<-c [int]>,
      B<--code=[int]>
      B<-t [int]>,
      B<--table=[int]>

Use NCBI genetic code tableID <int> for translating sequences.

=item B<--tables>,
      B<--codes>

Output a list of NCBI genetic code tableIDs and exit.

=item B<-3>

Translate each sequence in all three forward frames.

=item B<-6>

Translate each sequence in all six frames.

=item B<-s>,
      B<--stop>=<char>

Specify a string to representing stop codons. By default "*" is used.

=item B<-u>,
      B<--unknown>=<char>

Specify a string representing unknown amino acids. By default "X" is used.

=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<-q>
      B<--fastq>

Use fastq format as input and output.

=back

=head1 EXAMPLES

Translate each nucleotide seqeunces contained in file data.fas in all three forward reading frames.
 
=over 8

B<fasxl> -3 data.fas

=back

Translate each nucleotide seqeunces contained in file data.fas in all six reading frames and only select sequences that start with Met and end with a stop codon.

=over 8

B<fasxl> -6 data.fas | B<fasgrep> -s "^M.*\*$" 

=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
