#!/usr/bin/env perl

# MANUAL FOR orf_predict.pl

=pod

=head1 NAME

orf_predict.pl -- predicts open reading frames from a nucleotide FASTA file

=head1 SYNOPSIS

 orf_predict.pl --fasta=/Path/to/infile.fasta --out=/Path/to/output_file
                     [--help] [--manual]

=head1 DESCRIPTION

 Performs a 6-frame translation for each sequence using the Standard Codon
 Usage Table (transl_table=1) and then identifies complete or incomplete
 open reading frames >21 NT.
 
=head1 OPTIONS

=over 3

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

Input file in FASTA format. (Required) 

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

Output file base name. (Required)
Two FASTA files of ORFs are written, one NT and one AA.

=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 $min_length = 21; ## the smallest an open reading frame can be is 21 base pairs
my $max_stop_freq = 0.01; ## wont look at translations with >= 1% stop codons
my $ffh;

## validate
check_fasta($fasta);

## open the outfile
open(my $aafh,'>',"$outfile.pep") || die "\n Cannot write to the outfile: $outfile.pep\n";
open(my $ntfh,'>',"$outfile.nuc") || die "\n Cannot write to the outfile: $outfile.nuc\n";

## open the infile if compressed / or if not compressed
if ($fasta =~ m/\.gz$/) { ## if a gzip compressed infile
    open($ffh, '-|', 'gunzip', '-c', $fasta) || die "\n\n Cannot open the input file: $fasta\n\n";
}
else { ## If not gzip comgressed
    open($ffh,'<', $fasta) || die "\n\n Cannot open the input file: $fasta\n\n";
}
while(<$ffh>) {
    $_ =~ s/\r[\n]*/\n/gm;
    chomp;
    if ($_ =~ m/^>/) { ## Header line
	unless ($seq eq "") { ## Unless we're on the *first* header line
	    $seq = uc($seq);
	    my @Orfs = find_orfs($seq); ## Find all of the potential open reading frames in the sequence
	    print_orfs($header, @Orfs); ## Print all of the open reading frames found for this sequence
	}
	$header = $_;
	$header =~ s/^>//;
	$header =~ s/ .*//;
	$seq = "";
    }
    else {
	$seq = $seq . $_; ## extend the sequence
    }
}
close($ffh);

## Need to loop through all of this one more time for the last sequence.
$seq = uc($seq);
my @Orfs = find_orfs($seq);
print_orfs($header, @Orfs);

close($ntfh);
close($aafh);

## 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 print_orfs
{
    ## Input: The FASTA header and then an array of ORF sequences for that header.
    ## Output: Nothing, just print each ORF sequence to the output file.
    my $header = shift(@_);
    my @Orfs = @_;
    for (my $i=0; $i<scalar(@Orfs); $i++) {
	my $orf_num = $i+1;
	my $type = orf_type($Orfs[$i]);
	print $ntfh ">" . $header . "_orf_" . $orf_num . " type=$type" . "\n" . $Orfs[$i] . "\n";
	print $aafh ">" . $header . "_orf_" . $orf_num . " type=$type" . "\n" . translate($Orfs[$i]) . "\n";
    }
}

sub find_orfs
{
    ## Input: A nucleotide sequence
    ## Output: An array of all open reading frames found in that sequence
    ##         that meet the minimum size and stop codon frequency checks
    my $s = $_[0];
    my @Orfs;
    my @Frames = six_frames($seq);
    for (my $i=0; $i<scalar(@Frames); $i++) {
	my $freq = stop_freq($Frames[$i]);
	if ($freq < $max_stop_freq) {
	    my @Tmp = break_orfs($Frames[$i]);
	    push(@Orfs, @Tmp);
	}
    }
    return @Orfs;
}

sub break_orfs
{
    ## Input: A nucleotide sequence
    ## Output: An array of open reading frames broken on stop codons
    ##         and checked for length
    my $s = $_[0];
    my %Stops = ('TAA' => 1, 'TGA' => 1, 'TAG' => 1); ## these are the three stop codon types
    my @Orfs;
    my $orf = "";
    for (my $j=0; $j<length($s); $j+=3) {
        my $codon = substr $s, $j, 3;
        if (length($codon) == 3) {
	    if (exists $Stops{$codon}) {
		$orf = $orf . $codon;
		if (length($orf) >= $min_length) { push(@Orfs, $orf); }
		$orf = "";
	    }
	    else {
		$orf = $orf . $codon;
	    }
	}
    }
    if (length($orf) >= $min_length) { push(@Orfs, $orf); }
    return @Orfs;
}

sub six_frames
{
    ## Input: A nucleotide (NT) sequence
    ## Output: An array of the 6 possible NT sequences that it could be coding in (i.e. frames 1,2,3,-1,-2,-3)
    my $s = $_[0];
    my $hold = $s;
    my @Frames;
    for (my $i=0; $i<=2; $i++) {
	push(@Frames, $s);
	$s =~ s/^.//;
    }
    $s = $hold;
    $s = revcomp($s);
    for (my $i=0; $i<=2; $i++) {
	push(@Frames, $s);
	$s =~ s/^.//;
    }
    return @Frames;
}

sub translate
{
    ## Input: A nucleotide sequence
    ## Output: A translated peptide sequence using the codon table
    my $s = $_[0];
    my $pep = "";
    for (my $j=0; $j<length($s); $j+=3) {
	my $codon = substr $s, $j, 3;
	if (length($codon) == 3) {
	    if (exists $Codons{$codon}) {
		$pep = $pep . $Codons{$codon};
	    }
	    else { $pep = $pep . "X"; $warn++; }
	}
    }
    return $pep;
}

sub orf_type
{
    ## Input: An open reading frame (nucleotide sequence)
    ## Output: The "type" of ORF it is. Four options: Has both a start and stop codon
    ##         has only a start codon (ATG), has only a stop codon (TAA,TGA,TAG), or
    ##         has neither.
    my $s = $_[0];
    my %type = ("00" => "incomplete", "01" => "lack_start", "10" => "lack_stop", "11" => "complete");
    my $t;
    if ($s =~ m/^ATG/) { $t = "1"; }
    else {$t = "0";}
    if ($s =~ m/TAA$/ || $s =~ m/TGA$/ || $s =~ m/TAG$/) {
	$t = $t . "1";
    }
    else {$t = $t . "0";}
    return $type{$t};
}

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 stop_freq
{
    ## Input: A nucleotide sequence
    ## Output: The frequency with which stop codons appear in that sequence
    my $s = $_[0];
    my $stops = 0;
    my %Stops = ('TAA' => 1, 'TGA' => 1, 'TAG' => 1);
    for (my $j=0; $j<length($s); $j+=3) {
        my $codon = substr $s, $j, 3;
	if (length($codon) == 3) {
	    if (exists $Stops{$codon}) { $stops++; }
	}
    }
    my $frac = $stops / (length($s)/3);
    return $frac;
}

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