#!/usr/bin/env perl

# MANUAL FOR six_frame_translation.pl

=pod

=head1 NAME

six_frame_translation.pl -- performs 6-frame translation for each sequence

=head1 SYNOPSIS

 six_frame_translation.pl --fasta=/Path/to/infile.fasta --out=/Path/to/output.fasta
                     [--help] [--manual]

=head1 DESCRIPTION

 Performs a 6-frame translation for each sequence using the Standard Codon
 Usage Table (transl_table=1)
 
=head1 OPTIONS

=over 3

=item B<-f, --fasta>=FILENAME

Input file in FASTA format. (Required) 

=item B<-o, --out>=FILENAME

Output file in FASTA format. (Required) 

=item B<-h, --help>

Displays the usage message.  (Optional) 

=item B<-m, --manual>

Displays full manual.  (Optional) 

=back

=head1 DEPENDENCIES

Requires the following Perl libraries.



=head1 AUTHOR

Written by Daniel Nasko, 
Center for Bioinformatics and Computational Biology, University of Maryland.

=head1 REPORTING BUGS

Report bugs to dnasko@umiacs.umd.edu

=head1 COPYRIGHT

Copyright 2017 The S2FAST Team.
License GPLv3+: GNU GPL version 3 or later <http://gnu.org/licenses/gpl.html>.  
This is free software: you are free to change and redistribute it.  
There is NO WARRANTY, to the extent permitted by law.  

=cut


use strict;
use warnings;
use Getopt::Long;
use File::Basename;
use Pod::Usage;

#ARGUMENTS WITH NO DEFAULT
my($fasta,$outfile,$help,$manual);

GetOptions (	
				"f|fasta=s"	=>	\$fasta,
				"o|out=s"	=>	\$outfile,
				"h|help"	=>	\$help,
				"m|manual"	=>	\$manual) || pod2usage({-exitval => 2, -verbose => 1, -output => \*STDERR});

# VALIDATE ARGS
pod2usage(-verbose => 2)  if ($manual);
pod2usage( {-exitval => 0, -verbose => 2, -output => \*STDERR} )  if ($help);
pod2usage( -msg  => "\n\n ERROR!  Required argument --fasta not found.\n\n", -exitval => 2, -verbose => 1)  if (! $fasta );
pod2usage( -msg  => "\n\n ERROR!  Required argument --out not found.\n\n", -exitval => 2, -verbose => 1)  if (! $outfile);

## Globals
my %Codons = BuildCodonHash();
my ($header,$seq) = ("","");
my $warn=0;
my $fh;

## validate
check_fasta($fasta);

## open the outfile
open(my $ofh,'>',"$outfile") || die "\n Cannot write to the outfile: $outfile\n";

## open the infile if compressed / or if not compressed
if ($fasta =~ m/\.gz$/) { ## if a gzip compressed infile
    open($fh, '-|', 'gunzip', '-c', $fasta) || die "\n\n Cannot open the input file: $fasta\n\n";
}
else { ## If not gzip comgressed
    open($fh,'<',"$fasta") || die "\n\n Cannot open the input file: $fasta\n\n";
}
while(<$fh>) {
    $_ =~ s/\r[\n]*/\n/gm;
    chomp;
    if ($_ =~ m/^>/) {
	unless ($seq eq "") { ## Allows us to parse wrapped FASTA files
	    $seq = uc($seq);
	    my @Translations = Translate($seq);
	    for (my $i=0; $i<scalar(@Translations); $i++) {
		my $frame = $i+1;
	    	print $ofh ">" . $header . "_trans_" . $frame . "\n" . $Translations[$i] . "\n";
	    }
	}
	$header = $_;
	$header =~ s/^>//;
	$seq = "";
    }
    else {
	$seq = $seq . $_;
    }
}
close($fh);
## Need to loop through all of this one more time for the last sequence.
$seq = uc($seq);
my @Translations = Translate($seq);
for (my $i=0; $i<scalar(@Translations); $i++) {
    my $frame = $i+1;
    print $ofh ">" . $header . "_trans_" . $frame . "\n" . $Translations[$i] . "\n";
}
close($ofh);

## Print out a warning if there are non-standard bases.
if ($warn > 0) {
    print STDERR "\n Warning: $warn codons contained non-standard ATGC bases. These were translated to 'X's...\n\n";
}

exit 0;

sub Translate
{
    ## Input: A nucleotide (NT) sequence
    ## Output: An array of the 6 possible peptide sequences that one NT sequence could be (i.e. frames 1,2,3,-1,-2,-3)
    my $s = $_[0];
    my @Translated;
    for (my $i=0; $i<=2; $i++) {
	my $peptide = "";
	for (my $j=$i; $j<length($s); $j+=3) {
	    my $codon = substr $s, $j, 3;
	    if (length($codon) == 3) {
		if (exists $Codons{$codon}) {
		    $peptide = $peptide . $Codons{$codon};
		}
		else { $peptide = $peptide . "X"; $warn++; }
	    }
	}
	push(@Translated, $peptide);
    }
    $s = revcomp($s);
    for (my $i=0; $i<=2; $i++) {
        my $peptide = "";
        for (my $j=$i; $j<length($s); $j+=3) {
            my $codon = substr $s, $j, 3;
            if (length($codon) == 3) {
                if (exists $Codons{$codon}) {
                    $peptide = $peptide . $Codons{$codon};
                }
                else { $peptide = $peptide . "X"; $warn++; }
            }
        }
        push(@Translated, $peptide);
    }
    return @Translated;
}
sub revcomp
{
    ## Input: A nucleotide sequence
    ## Output: A reverse compliment of that nucleotide sequence
    my $s = $_[0];
    $s = scalar reverse $s;
    $s = uc($s);
    $s =~ tr/ATGC/TACG/;
    return $s;
}
sub BuildCodonHash
{
    ## Input: None
    ## Output: A dictionary with all codon mappings
    my %aacode = (
  TTT => "F", TTC => "F", TTA => "L", TTG => "L",
  TCT => "S", TCC => "S", TCA => "S", TCG => "S",
  TAT => "Y", TAC => "Y", TAA => "*", TAG => "*",
  TGT => "C", TGC => "C", TGA => "*", TGG => "W",
  CTT => "L", CTC => "L", CTA => "L", CTG => "L",
  CCT => "P", CCC => "P", CCA => "P", CCG => "P",
  CAT => "H", CAC => "H", CAA => "Q", CAG => "Q",
  CGT => "R", CGC => "R", CGA => "R", CGG => "R",
  ATT => "I", ATC => "I", ATA => "I", ATG => "M",
  ACT => "T", ACC => "T", ACA => "T", ACG => "T",
  AAT => "N", AAC => "N", AAA => "K", AAG => "K",
  AGT => "S", AGC => "S", AGA => "R", AGG => "R",
  GTT => "V", GTC => "V", GTA => "V", GTG => "V",
  GCT => "A", GCC => "A", GCA => "A", GCG => "A",
  GAT => "D", GAC => "D", GAA => "E", GAG => "E",
  GGT => "G", GGC => "G", GGA => "G", GGG => "G",
	);
    return %aacode;
}

sub check_fasta
{
    ## Input: the input FASTA file
    ## Output: Nothing, just exit if two cases are met.
    my $file = $_[0];
    if ( -z $file ) { die "\n Error: The input FASTA file appears to be empty\n"; } ## If the file is empty then just exit 0
    my $seq_count=0;
    open(my $fh, '<', $file) || die "\n Error: Cannot open the fasta file $file\n";
    while(<$fh>) {
	chomp;
	if ($_ =~ m/^>/) { $seq_count++; }
    }
    close($fh);
    if ($seq_count == 0) { die "\n Error: the input FASTA file does not appear to be in FASTA format as there are no header lines.\n\n"; }
    return 0;
}
