#!/usr/bin/env perl

#  BAGEL4 program 
# 	Anne de Jong and Auke van Heel
#	University of Groningen
#	the Netherlands
#	anne.de.jong@rug.nl
#
#  	samll ORF calling routine adapted from Bagel4
#
# v 2.0 
# 


use strict ;
use warnings ;
use lib "/data/molgentools/lib";
use anne_files ;
use anne_genomics ;


# -----------------------------------------------------------------------------------  global vars ---------------------------------------------------------------------------
my $sessiondir = "./" ;
my $query ;
my $min = 60 ;
my $max = 180 ;
my $outPrefix = "result" ;
my $ORF_overlap = 30 ;  # allows 10bp overlap
my $usage = "small_orf_calling.pl
				-s Sessiondir [default=$sessiondir]
				-query queryfile
				-min minimal aa [default=$min]
				-max maximal aa [default=$max]
				-outPrefix prefix for output file
				
		DEPENDS on: query.fna & query.predict
		
		small ORF calling routine 		
				
e.g.  ./small_orf_calling.pl -s . -query NC_0085331.AOI_1.fasta -min 60 -max 240 -outPrefix NC_0085331.AOI_1
	
" ;

&parseparam() ;


# ---------------------------------------------------------- main -------------------------------------------------------------------------
	my $orf_count = 0 ;
	my %GeneTable ;

	my @results_fna ;
	my @results_faa ;
	my @results_table = "scaffold\name\tstart\tend\tstrand\trbs\tgene\tprotein" ;

# 1. Get the sequences
	my %fasta=anne_genomics::fasta2hash("$sessiondir/$query");
	foreach my $dnaKey (keys %fasta) {
		print $dnaKey."\n";
		my $DNA = $fasta{$dnaKey} ;
		orf_calling($DNA, '+');
		orf_calling(anne_genomics::inverse_complement($DNA), '-');
		
		foreach my $geneKey (sort {$GeneTable{$a}{start} <=> $GeneTable{$b}{start}} keys %GeneTable) {
			my $protein = anne_genomics::translate($GeneTable{$geneKey}{gene});
			$protein =~ s/-$// ;  # remove last stopcodon char
			push @results_faa, ">$geneKey\n$protein";
			push @results_fna, ">$geneKey | scaffold=$dnaKey | strand=$GeneTable{$geneKey}{strand} start=$GeneTable{$geneKey}{start} end=$GeneTable{$geneKey}{end} RBS=$GeneTable{$geneKey}{rbs}\n$GeneTable{$geneKey}{gene}\n"  ; 
			push @results_table, "$dnaKey\t$geneKey\t$GeneTable{$geneKey}{start}\t$GeneTable{$geneKey}{end}\t$GeneTable{$geneKey}{strand}\t$GeneTable{$geneKey}{rbs}\t$GeneTable{$geneKey}{gene}\t$protein" ;
			#my $dna_gene = substr $GeneTable{$geneKey}{gene}, 0, 30 ;
			#my $dna_true = substr $DNA, $GeneTable{$geneKey}{start}, 30 ;
			#my $truepos = index($DNA, $GeneTable{$geneKey}{gene});
			#$dna_true = substr anne_genomics::inverse_complement($DNA), $GeneTable{$geneKey}{start}, 30 if ($GeneTable{$geneKey}{strand} eq '-') ;
		}
		%GeneTable =();
		
	}	
	
	
	anne_files::write_lines("$sessiondir/$outPrefix.sOFS.faa", @results_faa) ;
	anne_files::write_lines("$sessiondir/$outPrefix.sOFS.fna", @results_fna) ;
	anne_files::write_lines("$sessiondir/$outPrefix.sOFS.table", @results_table) ;


	

	


# ---------------------------------------------------------- functions -------------------------------------------------------------------------

sub mask_orf_regions {
	# mask all gene regions and the tailing parts
	my $dna = shift ;
	my @lines = anne_files::read_lines("$sessiondir/$query.predict");  # read the glimmer orf positions
	my $first_start = 0 ;
	my $last_end = 0;
	my $first_line = 1 ;
	foreach my $line(@lines) {
		my @items = split /\s+/, $line ;
		if (scalar @items>2) {
			if ($items[1] > $items[2]) { my $tmp = $items[1] ; $items[1] = $items[2]; $items[2] = $tmp; }  # swap if start/end for strand -
			my $start = $items[1] + $ORF_overlap ; # allow ORF overlap
			my $end   = $items[2] - $ORF_overlap ;
			if ($first_line) { $first_line = 0; $first_start = $start; }
			$last_end = $end;
			my $len = $end-$start ;
			my $N_ORF = 'N' x $len ;
			my $ORF = substr $dna, $start, $len ;
			$dna =~ s/$ORF/$N_ORF/ ;  # put N's at the ORF regions
		}	
	}
	
	# Do not allow ORF calling in the tails; remove the tails
	my $N_ORF = 'N' x $first_start ;
	my $ORF = substr $dna, 0, $first_start ;
	$dna =~ s/$ORF/$N_ORF/ ;  # put N's at the left tail
	my $len = length($dna) - $last_end  ;
	#print "$N_ORF\n$ORF\n$len-$last_end\n";
	if ($len>1) {
		$N_ORF = 'N' x $len ;
		$ORF = substr $dna, -$len ;
		$dna =~ s/$ORF/$N_ORF/ ;  # put N's at the right tail
		#print "$N_ORF\n$ORF\n$len-$last_end\n";
	}	
	return $dna ;
}


sub orf_calling {
	my $DNA = shift ;
	my $strand = shift ;
	my $startcodons = 'ATG|GTG|TTG';
	my $index = 1 ;
	my $DNAlen = length($DNA) ;
	my $start ;
	my $end ;
	my $lastgene ="";
	#	while ( (substr $DNA, $index) =~m/(..GG.{5,14})((ATG|GTG)(...){0,$max}?(TGA|TAG|TAA))/) {


	while ( (substr $DNA, $index) =~m/(.....{5,14})(($startcodons)((ATG|GTG|TTG|GCC|AGT|TGT|CGA|ATC|AAC|AGC|TAC|TCG|ACA|CTG|CCG|GCA|AAG|GTT|CAC|AGA|ACC|CCA|TGG|CGC|CTC|CAG|ACG|AAA|GTA|CTT|GGA|GTC|TGC|TCA|ATT|TAT|AAT|ACT|CAA|GAC|GGT|TCC|TTT|AGG|CGT|CGG|CAT|ATA|CCC|GGG|GAG|TTA|CTA|GAT|TCT|TTC|GCG|GGC|GAA|GCT|CCT)){$min,$max}(TGA|TAG|TAA))/) {
		my $RBS = $1 ;
		my $gene = $2 ;
		my $currentIndex = $index ;
		my $zeroIndex = $-[0] ;
		my $genestart = $index + length($RBS) + $-[0] ;
		$index += $-[0] + 3*$min - 3 ;  # goto the next position and check again for a gene
		#my $genestart = $index - 3*$conf{smallorf_min_prot_length} - 3 ;  # the gene start position
		my $DNAstart = substr $DNA, $genestart, 10 ;
		$RBS =~ s/N//g ;  # remove the N's from the RBS after the $_ values because s/ will generate a new $_
		if ($gene ne $lastgene) {
			$orf_count++ ;
			my $key = "sORF_$orf_count" ;
			#print "$key\t$strand\t$currentIndex\t$index\t$zeroIndex\t$genestart\t$DNAstart\t".anne_genomics::translate($gene)."\n";
			$GeneTable{$key}{gene} = $gene ;
			$GeneTable{$key}{rbs} = $RBS ;
			if ($strand eq '+') {
				$GeneTable{$key}{start}	= $genestart + 1 ;
				$GeneTable{$key}{end}	= $GeneTable{$key}{start} + length($gene) -1;
				$GeneTable{$key}{strand} = $strand ;
			} elsif ($strand eq '-') {
				$GeneTable{$key}{end} 	= $DNAlen - $genestart ;
				$GeneTable{$key}{start}	= $GeneTable{$key}{end} - length($gene) + 1 ;
				$GeneTable{$key}{strand} = $strand ;
			}	
		}
		$lastgene = $gene ;
	}
}



sub chk_protein {
	my $prot = shift ;
	my $chk = 0 ;
	if ( (substr $prot, 1, 15) !~ m/.*C.*/g ) {  # should not contain a C in the first 15 aa
		my $mature = substr $prot, 15  ;
		if ( $mature=~ m/[C]/g and $mature=~ m/[TS]/g ) {  # should contain a C or T/S in the leader
			$chk = 1;
		}	
	}
	return $chk ;
}

	
	
sub parseparam {
    my $var ;
    my @arg = @ARGV ;
    while(@arg) {
        $var = shift(@arg) ;
        die $usage if ($var eq '-h' or $var eq '--help') ;
		$sessiondir		= shift(@arg) if($var eq '-s') ;
		$query		= shift(@arg) if($var eq '-query') ;
		$min		= shift(@arg) if($var eq '-min') ;
		$max		= shift(@arg) if($var eq '-max') ;
		$outPrefix		= shift(@arg) if($var eq '-outPrefix') ;
    }
	die $usage if (!$query) ;
}