#!/usr/bin/perl -w
# Bruno Contreras-Moreira, Pablo Vinuesa, ...
# 2005-7 UNAM, Mexico

# Fasta4phylogenomics.pl

# program that takes a .faa/.fna pair of sequence files and runs phylogenomic tests with them

use strict;
use Getopt::Std;
use lib "lib";
use phyTools;

set_phyTools_env();

my $progname = $0; # Fasta4phylogenomics.pl
my $VERSION  = 0.2; # Oct 1st, 2011; PV adjusted code to run new versions of marfil_prot.pl and marfil_amp.pl, 
                    #     wich in turn make calls to run muscle3.8.31, phyml3.0.1 and protTest2.4, 
		    #     New subs were correspondingly added in phyTools.pm).

my $MARFILPROTEXE = "marfil_prot.pl -s -P -D -i ";
my $MARFILAMPSEXE = "marfil_amps.pl -n ";

my $DUMMYCUTOFF = 999;

########################################################################

my (%opts,$marfilEXE,$cwd);
my ($fnafile,$faafile,$exclude_list,$keep_files,$use_marfil_prot,$only_unique_seqs,$bootstrap,$user_subst_model,$skip_alignments,$only_align);

getopts('haskuMBP:n:p:e:T:m:', \%opts);

if(($opts{'h'})||(scalar(keys(%opts))==0)) 
{ 
	print   "\nusage: $progname [options]\n\n";
	print   "-h \t this message\n";
	print   "-p \t file with protein sequences in FASTA format                         (by default unaligned sequences)\n";
	print   "-n \t file with nucleotide sequences in FASTA format                      (optional, required if -M not used\n";
	print   "-e \t exclude sequence identifiers contained in this file                 (optional)\n";
	print   "-u \t use only unique sequences                                           (optional)\n";
	print   "-M \t use marfil_prot.pl for phylogenetical evaluation                    (optional, default marfil_amps.pl)\n";
	print   "-B \t calculate WAG bootstrap support for amplicon topologies             (optional)\n";
	print   "-k \t keep all temporary files                                            (optional)\n";
	print   "-m \t use substitution model provided by user                             (optional,[WAGG|JTTG|...])\n";  
	print   "-s \t skip protein alignments as input sequences are already aligned      (optional)\n";
   print   "-a \t only align sequences                                                (optional)\n";#BCM
	exit; 
}

if(defined($opts{'p'})){ $faafile = $opts{'p'}; }
else{ die "# $0 : need a protein file in FASTA format\n" }

if(defined($opts{'e'})){ $exclude_list = $opts{'e'}; }
else{ $exclude_list = '' }
if(defined($opts{'k'})){ $keep_files = 1; }
else{ $keep_files = 0 }
if(defined($opts{'u'})){ $only_unique_seqs = 1; }
else{ $only_unique_seqs = 0; }
if(defined($opts{'a'})){ $only_align = 1; }
else{ $only_align = 0; }

if(defined($opts{'M'}))
{ 
	$use_marfil_prot = 1; 
	$marfilEXE = $MARFILPROTEXE;
}
else
{ 
	$use_marfil_prot = 0; 
	$marfilEXE = $MARFILAMPSEXE;
}

if(defined($opts{'n'})){ $fnafile = $opts{'n'}; }
elsif(!$use_marfil_prot){ die "# $0 : need a nucleotide file in FASTA format\n" }

if(defined($opts{'B'})){ $bootstrap = 1; }
else{ $bootstrap = 0; }

if(defined($opts{'m'})){ $user_subst_model = $opts{'m'}; }
else{ $user_subst_model = ''; }

if(defined($opts{'s'})){ $skip_alignments = 1; }
else{ $skip_alignments = 0; }

print "# $progname -n $fnafile -p $faafile -e $exclude_list -M $use_marfil_prot -k $keep_files -u $only_unique_seqs -B $bootstrap -m $user_subst_model -s $skip_alignments -a $only_align\n\n";

$cwd = `pwd`; chomp $cwd;

print "# $progname : cwd = $cwd\n\n";

############################################################################

my (@EXCLIST,@marfil_files,$header,$seq,$exclterm,$lastcodon,$isSTOPcodon);
my ($nt_amp,$aa_amp,$marfil_command,$marfil_input,$marfil_no_input,$alpha,$aln_stats,$alignment_quality,$n_of_sequences);
my ($substitution_model,$full_marfil_stats,$phyfile,$PHYMLtree,$puzzleparamsfile,$modelfile,$puzzledistfile,$tarfile);
#my ($FRQ, $n_seq_comb_bias); # PV 2 Jun 07

# -1) read exclusion file that contains taxa that should be removed from the analysis ###############
if($exclude_list)
{
	print "# Reading exclude list...\n";

	open(EXCLIST,$exclude_list) || die "# $progname : cannot read $exclude_list\n";
	while(<EXCLIST>)
	{
		next if(/^#|^$/);
		chomp;
		push(@EXCLIST,$_);
		print "# to be excluded: $EXCLIST[$#EXCLIST]\n"; 
	}
	close(EXCLIST);
}

# 0) remove identical and sequences included in @EXCLIST #########################################

my %fullDNAfasta;
if(!$use_marfil_prot){ %fullDNAfasta = read_FASTA_sequence($fnafile, 0, 0, $only_unique_seqs); } # 0 ,0 = remove gaps, skip bad CDSs 
my %fullAAfasta = read_FASTA_sequence($faafile);
my (%DNAfasta,%AAfasta);

SEQS: foreach $seq (sort {$a<=>$b} (keys(%fullAAfasta)))
{
	next if(!$fullDNAfasta{$seq});

	if(@EXCLIST && $fullAAfasta{$seq}{'NAME'} =~ /\[((\S+\s?){1,4})/)   
	{ 
		$header = $1;  #print ">$header<\n";
		foreach $exclterm (@EXCLIST){ next SEQS if($header =~ /$exclterm/); }
	}	
	# 1.1) this sequence should remain
	if(!$use_marfil_prot){ $DNAfasta{$seq} = $fullDNAfasta{$seq}; }
	$AAfasta{$seq} = $fullAAfasta{$seq};
}

if(!$use_marfil_prot && scalar(keys(%DNAfasta)) != scalar(keys(%AAfasta)))
{
	die "# $progname : input FASTA files contain different numbers of sequences\n";	
}
else
{ 
	if(scalar(keys(%AAfasta)) < 2)
	{
		die "# $progname : need at least 2 sequences to run this program...\n";
	}
	else 
	{ 
		$n_of_sequences = scalar(keys(%AAfasta));
		print "# $progname : number of sequences considered = $n_of_sequences\n"; 
	}
}

# 1) check sequence quality and length ###########################################################
if(!$use_marfil_prot)
{
	$fnafile = reverse $fnafile;
	$fnafile = (split(/\//,$fnafile))[0];
	$fnafile = reverse $fnafile;
}

$faafile = reverse $faafile;
$faafile = (split(/\//,$faafile))[0];
$faafile = reverse $faafile;

my $myDNAfile;
if(!$use_marfil_prot)
{ 
	$myDNAfile = "my_" . (split(/\./,$fnafile))[0] . ".fna"; 
	push(@marfil_files,$myDNAfile);
}
my $myAAfile  = "my_" . (split(/\./,$faafile))[0]  . ".faa";
push(@marfil_files,$myAAfile);

if(!$use_marfil_prot)
{ 
	open(MYDNA,">$myDNAfile") || die "# $progname : cannot write to $myDNAfile\n"; 
	#print "# FASTA header equivalences (DNA):\n";

	foreach $seq (sort {$a<=>$b} (keys(%DNAfasta)))
	{	
		$DNAfasta{$seq}{'NAME'} =~ s/>//;
		#print "# $seq => $DNAfasta{$seq}{'NAME'}";

		$DNAfasta{$seq}{'SEQ'} = lc($DNAfasta{$seq}{'SEQ'});	

		$isSTOPcodon = 0;
		$lastcodon = substr($DNAfasta{$seq}{'SEQ'},length($DNAfasta{$seq}{'SEQ'})-3,3);
		if($lastcodon eq 'tga' || $lastcodon eq 'taa' || $lastcodon eq 'tag'){ $isSTOPcodon = 1 }
		#print "#$lastcodon#$isSTOPcodon#\n";

		# remove STOP codon if found
		if($isSTOPcodon){ $DNAfasta{$seq}{'SEQ'} = substr($DNAfasta{$seq}{'SEQ'},0,-3); }

		print MYDNA ">$seq " . $DNAfasta{$seq}{'NAME'}; 
		print MYDNA $DNAfasta{$seq}{'SEQ'},"\n"; 
	}

	close(MYDNA);
}

open(MYAA,">$myAAfile") || die "# $progname : cannot write to $myAAfile\n";	

print "\n# FASTA header equivalences (AA):\n";

foreach $seq (sort {$a<=>$b} (keys(%AAfasta)))
{
	$AAfasta{$seq}{'NAME'} =~ s/>//;
	$AAfasta{$seq}{'SEQ'} =~ s/\*|#|\+//; ## OPA,OCH,AMB in SEQ.pm: not required if STOP codons are present in nt sequence
		
	print "# $seq => $AAfasta{$seq}{'NAME'}";
		
	print MYAA ">$seq " . $AAfasta{$seq}{'NAME'}; 
	print MYAA $AAfasta{$seq}{'SEQ'},"\n";
}

close(MYAA);
print "\n";

if(!$use_marfil_prot)
{
	foreach $seq (sort {$a<=>$b} (keys(%AAfasta)))
	{
		if(length($AAfasta{$seq}{'SEQ'})*3 != length($DNAfasta{$seq}{'SEQ'}))
		{
			print "# sequence $seq : amino acid and nt lengths do not match (".
				(length($AAfasta{$seq}{'SEQ'})*3).
				") and (".length($DNAfasta{$seq}{'SEQ'}).") ...\n";
			die "#$AAfasta{$seq}{'SEQ'}\n$DNAfasta{$seq}{'SEQ'}\n"; 
		}
	}
}

# 2) multiply align sequences ##############################################
my $AA_aln_file;
unless($skip_alignments)
{
	$AA_aln_file = run_muscle($myAAfile,!$keep_files); 
	if($AA_aln_file eq $ERROR)
	{
		die "# $progname : sorry, muscle alignment failed\n";
	}
	else{ $aa_amp = $AA_aln_file; }
}
else{ $aa_amp = $myAAfile; }

if(!$use_marfil_prot)
{
	push(@marfil_files,$AA_aln_file); #PV 19 Feb 2008; to include in tar -cvz

	# 2.1) get a nucleotide alignment based on the one produced by muscle
	my $DNA_aln_file = convert_AAaln2NTaln($myDNAfile,$AA_aln_file); 
	if($DNA_aln_file eq $ERROR)
	{
		die "# $progname : sorry, convert_AAaln2NTaln failed\n";
	}
	else{ $nt_amp = $DNA_aln_file; }
}
else{ $nt_amp = '' }

# 3) run marfil ################################################
if(!$use_marfil_prot)
{
	print "## alignment $fnafile\n";        
}
else{ print "## alignment $faafile\n"; }

if($only_align)
{
   if(!$keep_files)
	{
		unlink(@marfil_files);
	}
	exit;
}


$alignment_quality = $marfil_command = '';
					
if($use_marfil_prot)
{ 
	$marfil_command = "$marfilEXE $aa_amp"; 
	if($user_subst_model){ $marfil_command .= " -m $user_subst_model" }  
	$marfil_input = $aa_amp;
	$marfil_no_input = $nt_amp;
}
else
{ 
	$marfil_command = "$marfilEXE $nt_amp"; 
	if($user_subst_model){ $marfil_command .= " -m $user_subst_model" }
	$marfil_input = $nt_amp;
	$marfil_no_input = $aa_amp;
} #die "$marfil_command\n";
	
			
$substitution_model = $full_marfil_stats = '';
open(MARFIL,"$marfil_command |") || die "# $progname : cannot execute $marfil_command\n";
while(<MARFIL>)
{               
	#($FRQ, $n_seq_comb_bias) = (0.0,0.0); # PV 02 Junio 07
	if(/^# output : subs_model = (\S+) stats = (.+)/)
	{
		$substitution_model = $1;
		$full_marfil_stats = $2; #print "$substitution_model $full_marfil_stats\n";
		
		# alignID \t alignname \t n_of_sequences \t nt_sites \t %const_sites \t n_of_site_patterns \t %const_site_patterns \t 
		# percent_gaps \t ChiSQ_failed \t Acc_No_ChiSQ_perc \t %fullyresolved_quartets \t %partiallyresolved_quartets \t
		# %nonresolved_quartets \t normalized_FRQ
		# wanted : 2,3,8,10,11,12,13?

		$alignment_quality .= "# subs.model = $substitution_model\n";
		$alignment_quality .= "# n_of_seqs = ".(split(/\t/,$full_marfil_stats))[2].
			" n_alignments_sites = ".(split(/\t/,$full_marfil_stats))[3]."\n";
		$alignment_quality .= "# n_of_seqs_with_composition_bias = ".
			(split(/\t/,$full_marfil_stats))[8]."\n";
		$alignment_quality .= "# \%fully resolved quartets = ".
			(split(/\t/,$full_marfil_stats))[10]." \%partially resolved quartets = ".
			(split(/\t/,$full_marfil_stats))[11]." \%non resolved quartets = ".
			(split(/\t/,$full_marfil_stats))[12]."\n";
			
		#$n_seq_comb_bias = (split(/\t/,$full_marfil_stats))[8]; print "##n_seq_comb_bias is $n_seq_comb_bias\n";# PV 02 Junio 07
		#$FRQ = (split(/\t/,$full_marfil_stats))[10];    print "##FRQ is $FRQ\n";        # PV 02 Junio 07
	}	
	elsif(/^# outfiles : (\S+) (\S+) (\S+) (\S+) (\S+)/)
	{
		# guardar para poesteriores analisis (al igual que .tab), posiblemente con tar
		($phyfile,$PHYMLtree,$puzzleparamsfile,$modelfile,$puzzledistfile) = ($1,$2,$3,$4,$5); 
		#$alignment_quality .= "# outfiles = $marfil_input $phyfile $PHYMLtree $puzzleparamsfile\n";
	}
	# alpha = 0 alignment stats : length = 207 %gaps = 3.8 %constant = 47.3 %site_patterns = 61.8
	elsif(/^# alpha = (\d+.\d+) alignment stats : (.+)/)
	{
		($alpha,$aln_stats) = ($1,$2);
		$alignment_quality .= "# alpha = $alpha\n";
		$alignment_quality .= "# aln_stats : $aln_stats\n";
	}
	elsif(/^# distance stats : (.+)/)
	{
		$alignment_quality .= "# dist_stats : $1\n";
	}
	elsif(/^# mean aLRT = \d+\.\d+ median aLRT = \d+\.\d+/)
	{
		$alignment_quality .= $_;
	}
	#elsif(/^: sorry,/){}
}
close(MARFIL);

## 2) write a sel file with best alns and corresponding prams and tree files # Pablo 2 June 07
# print "## phyfile is $phyfile\n## puzzleparamsfile is $puzzleparamsfile\n## PHYMLtree is $PHYMLtree\n\n";
	# open(SEL, ">>Selected_alns.sel") or die "can't append to Selected_alns.sel: $!\n";
# if ( $n_seq_comb_bias <= $ENV{"MAXCHISQFAIL"} && $FRQ >= $ENV{"MINPERFRQ"} )
# {
#     print SEL $phyfile ."\t" . $puzzleparamsfile . "\t" . $PHYMLtree . "\n";
# }
# close (SEL);

if($full_marfil_stats eq ''){ $alignment_quality = "# marfil $marfil_input failed\n";  }
else
{
	# tar outfiles
	$tarfile = $faafile;
	$tarfile =~ s/\.faa/\.tgz/; #print ">> $marfil_input | $marfil_no_input |";
	system("tar czf $tarfile $marfil_input $marfil_no_input $phyfile $PHYMLtree $puzzleparamsfile $modelfile $puzzledistfile"); # ojo estaba tar cfz PV Oct08
	
	push(@marfil_files,$phyfile,$PHYMLtree,$puzzleparamsfile,$modelfile,$puzzledistfile,$marfil_input);
	$alignment_quality .= "# compressed outfile = $tarfile\n";
}

if(!$bootstrap){ $alignment_quality .="# end_of_alignment\n"; }
	
elsif($n_of_sequences > 3) # must have 4+ sequences at least to build trees
{
	if($bootstrap)
	{
		my $ampphyfile = convert_FAS2PHY($aa_amp,!$keep_files); 
		my %ampFASTA = read_FASTA_sequence($aa_amp);				
		my $bootphyfile = $ampphyfile;
		$bootphyfile =~ s/\.phy/_boot.phy/;
		rename($ampphyfile,$bootphyfile);

		my ($mean_boot_support,$median_boot_support,$bootreefile) = 
			get_mean_boot4WAG_dist_tree($bootphyfile,\%ampFASTA,!$keep_files);
	
		$alignment_quality .= "# mean bootstrap value = $mean_boot_support median bootstrap value = $median_boot_support\n";
		push(@marfil_files,$bootphyfile,$bootreefile);
	
		$alignment_quality .= "# end_of_alignment\n"; 
	}
}
elsif($bootstrap){ print "# end_of_alignment (less than 4 sequences for -L/-B)\n"; }	



print "# phylogenetic alignment evaluation:\n$alignment_quality\n";			
			

# 4) clean tmp files ################################################################
if(!$keep_files)
{
	unlink(@marfil_files);
}

