#!/usr/bin/env perl
# Author: Eric Wafula
# Email: ekw10@psu.edu
# Institution: Penn State University, Biology Dept, Claude dePamphilis Lab
# Date: 01-31-2015

use strict;
use warnings;
use File::Spec;
use File::Basename;
use Getopt::Long qw(:config no_ignore_case);
use FindBin;

my $home =  "$FindBin::Bin/..";

my $usage = <<__EOUSAGE__;

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#
#                                  GENE FAMILY CLASSIFICATION
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#  Required Options:
#
#  --proteins <string>             : Amino acids (proteins) sequences fasta file (proteins.fasta)
#
#  --scaffold <string>             : Orthogroups or gene families proteins scaffold.  This can either be an absolute
#                                    path to the directory containing the scaffolds (e.g., /home/scaffolds/22Gv1.1)
#                                    or just the scaffold (e.g., 22Gv1.1).  If the latter, ~home/data is prepended to
#                                    the scaffold to create the absolute path.
#                                    the scaffold to create the absolute path.
#                                    If Monocots clusters (version 1.0): 12Gv1.0
#                                    If Angiosperms clusters (version 1.0): 22Gv1.0
#                                    If Angiosperms clusters (version 1.1): 22Gv1.1
#                                    If Green plants clusters (version 1.0): 30Gv1.0
#                                    If Other non PlantTribes clusters: XGvY.Z, where "X" is the number species in the scaffold,
#                                    and "Y.Z" version number such as 12Gv1.0. Please look at one of the PlantTribes scaffold
#                                    data on how data files and directories are named, formated, and organized.	
#
#  --method <string>               : Protein clustering method
#                                    If GFam: gfam
#                                    If OrthoFinder: orthofinder
#                                    If OrthoMCL: orthomcl
#                                    If Other non PlantTribes method: methodname, where "methodname" a nonempty string of
#                                    word characters (alphanumeric or "_"). No embedded special charaters or white spaces.
#
#  --classifier <string>           : Protein classification method 
#                                    If BLASTP: blastp
#                                    If HMMScan: hmmscan
#                                    If BLASTP and HMMScan: both
#
# # # # # # # # # # # # # # # # # # 
#  Others Options:
#
#  --config_dir <string>           : (Optional) Absolute path to the directory containing the default configuration files
#                                    for the selected scaffold defined by the value of the --scaffold parameter (e.g.,
#                                    /home/configs/22Gv1.1). If this parameter is not used, the directory containing the
#                                    default configuration files is set to ~home/config.
#
#  --num_threads <int>             : number of threads (CPUs) to used for HMMScan, BLASTP, and MAFFT
#                                    Default: 1 
#
#  --super_orthogroups <string>    : SuperOrthogroups MCL clustering - blastp e-value matrix between all pairs of orthogroups
#                                    If minimum e-value: min_evalue (default) 
#                                    If average e-value: avg_evalue
#
#  --single_copy_custom <string>   : Configuration file for single copy orthogroup custom selection - incompatible with
#                                    "--single_copy_taxa".  This must be an absolute path to the configuration file
#                                    (e.g., '~/home/scaffolds/22Gv1.1.singleCopy.config') or the string 'default'.   If the
#                                    latter, the config is defined to be ~/config_dir/~scaffold.singleCopy.config.  See the
#                                    single copy configuration files the config sub-directory of the installation on how to
#                                    customize the single copy selection.
#
#  --single_copy_taxa <int>        : Minimum single copy taxa required in orthogroup - incompatible with "--single_copy_custom"
#
#  --taxa_present <int>            : Minimum taxa required in single copy orthogroup - requires "--single_copy_taxa"
#
#  --orthogroup_fasta              : Create orthogroup fasta files - requires "--coding_sequences" for CDS orthogroup fasta
#                                    
#  --coding_sequences <string>     : Corresponding coding sequences (CDS) fasta file (cds.fasta)
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
#  Example Usage:
#
#  GeneFamilyClassifier --proteins proteins.fasta --scaffold 22Gv1.1 --method orthomcl --classifier blastp 
#
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
 
__EOUSAGE__
    ; 

# Declare and initialize variables;
my $proteins;
my $scaffold;
my $method;
my $classifier;
my $scaffold_dir;
my $config_dir;
my $num_threads;
my $super_orthogroups;
my $single_copy_custom;
my $single_copy_taxa;
my $taxa_present;
my $orthogroup_fasta;
my $coding_sequences;

my $options = GetOptions (  'proteins=s' => \$proteins,
              'scaffold=s' => \$scaffold,
              'method=s' => \$method,
              'classifier=s' => \$classifier,
              'scaffold_dir=s' => \$scaffold_dir,
              'config_dir=s' => \$config_dir,
              'num_threads=i' => \$num_threads,
              'super_orthogroups=s' => \$super_orthogroups,
              'single_copy_custom=s' => \$single_copy_custom,
              'single_copy_taxa=i' => \$single_copy_taxa,
              'taxa_present=i' => \$taxa_present,
              'orthogroup_fasta' => \$orthogroup_fasta,
              'coding_sequences=s' => \$coding_sequences,
              );
              
if (File::Spec->file_name_is_absolute($scaffold)) {
	$scaffold_dir = $scaffold;
	$scaffold = basename($scaffold);
} else {
	$scaffold_dir = "$home/data/$scaffold";
}

if (!$config_dir || !File::Spec->file_name_is_absolute($config_dir)) {
	$config_dir = "$home/config";
}

if ($single_copy_custom && $single_copy_custom eq "default") {
    $single_copy_custom = "$config_dir/$scaffold.singleCopy.config";
}

# validate options            
unless ( $options ) { die $usage; }
unless ( $proteins and $scaffold and $method and $classifier ) { print "\nOne or more required options not set\n"; die $usage; }
if ( $scaffold !~ /^\d+Gv\d+\.\d+$/) { print "\nOrthogroups or gene families proteins scaffold name $scaffold is not in the required format";
	print " i.e. XGvY.Z, where X is number species in the scaffold, and Y.Z version number such as 12Gv1.0.\n"; die $usage; }
if ( $method !~ /^\w+$/) { print "\nProtein clustering method name $method is not in the required format. The name should be a nonempty string";
	print " of word characters (alphanumeric or _). No embedded special charaters or white spaces.\n"; die $usage; }
if ( $single_copy_custom and $single_copy_taxa ) { 
	print "\nSingle copy custom selection is incompatibale with minumum single copy taxa required\n"; die $usage; 
}
if ( $taxa_present and !$single_copy_taxa ) { print "\nMinumum taxa required requires minumum single copy taxa required\n"; die $usage; }
$scaffold =~ /(\d+)(.*)/;
if ( $single_copy_taxa ) {
	unless ( $single_copy_taxa <= $1 ) { print "\nMinumum single copy taxa required in orthogroup exceeds scaffold taxa ($1 genomes)\n"; die $usage; }
	unless ( $taxa_present <= $1 ) { print "\nMinumum taxa required in single copy orthogroup exceeds scaffold taxa ($1 gemomes)\n"; die $usage; }
}

# defaults
if (!$num_threads) { $num_threads = 1; }

print "\n";
print localtime()." - Starting gene family classification\n";
print "-- Scaffold: $scaffold\n-- Method: $method\n-- Classifier: $classifier\n\n";

# create output directory
my $dirname ="./geneFamilyClassification_dir";
if (-d $dirname) { die "Exiting...!\nGene family classification output directory ($dirname) already exists!\n\n"; }
make_directory($dirname);

if ( $classifier eq "blastp" ) {
	sort_sequences ( $classifier, $proteins, $scaffold, $method, $num_threads, $super_orthogroups, $dirname, $scaffold_dir );
}
elsif ( $classifier eq "hmmscan" ) {
	sort_sequences ( $classifier, $proteins, $scaffold, $method, $num_threads, $super_orthogroups, $dirname, $scaffold_dir );
}
elsif ( $classifier eq "both" ) {
	sort_sequences ( $classifier, $proteins, $scaffold, $method, $num_threads, $super_orthogroups, $dirname, $scaffold_dir );
	get_blast_hmmscan_orthos ( $dirname, $scaffold, $method, $super_orthogroups, $scaffold_dir );
}
else { die "Unknown protein classification method.\n\n$usage"; }

if ( $single_copy_custom or $single_copy_taxa) {
	get_single_copy_orthogroups ( $single_copy_custom, $single_copy_taxa, $taxa_present, $classifier, $scaffold, $method, $dirname, $config_dir );
}

if ( $orthogroup_fasta ) {
	get_orthogroup_fasta ( $proteins, $coding_sequences, $single_copy_custom, $single_copy_taxa, $classifier, $scaffold, $dirname );
} 

print localtime()." - Completed gene family classification\n\n";

exit(0);

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # #  sub-routines # # # # # # # # # # # # # # # # # # # # # # # # # # #

sub make_directory {
        my ( $new_dir ) = @_;
        if (!-d $new_dir) {
                mkdir($new_dir, 0755);
        }
}

sub sort_sequences {
	my ( $classifier, $proteins, $scaffold, $method, $num_threads, $super_orthogroups, $dirname, $scaffold_dir ) = @_;
	print localtime()." - Sorting protein sequences\n";
	if ( (!$super_orthogroups) or ($super_orthogroups ne "min_evalue") or ($super_orthogroups ne "avg_evalue") ) { $super_orthogroups = "min_evalue" } 
	if ( ($classifier eq "blastp") or ($classifier eq "both") ) {
		print "-- ".localtime()." - Running BLASTP\n\n";
		system "blastp -db_soft_mask 21 -outfmt 6 -evalue 1e-5 -num_threads $num_threads -query $proteins -db $scaffold_dir/db/blast/$method -out $dirname/proteins.blastp.$scaffold  >/dev/null";
		print "-- ".localtime()." - Getting best BLASTP hits\n\n";
		my $results = "proteins.blastp.$scaffold";
		get_best_blastp_orthos ( $classifier, $results, $dirname, $scaffold, $method, $super_orthogroups, $scaffold_dir );
	}
	if ( ($classifier eq "hmmscan") or ($classifier eq "both") ) {
		print "-- ".localtime()." - Running HMMScan\n\n";
		system "hmmscan -E 1e-5 --cpu $num_threads --noali --tblout $dirname/proteins.hmmscan.$scaffold -o $dirname/hmmscan.log $scaffold_dir/db/hmm/$method $proteins >/dev/null";
		print "-- ".localtime()." - Getting best HMMScan hits\n\n";
		my $results = "proteins.hmmscan.$scaffold";
		get_best_hmmscan_orthos ( $classifier, $results, $dirname, $scaffold, $method, $super_orthogroups, $scaffold_dir );
	}	
}

sub  get_best_blastp_orthos {
	my ( $classifier, $blast_results, $dirname, $scaffold, $method, $super_orthogroups, $scaffold_dir ) = @_;
	my (%best, %max, %list);
    open (IN, "$dirname/$blast_results") or die "can't open $dirname/$blast_results file\n";
	while (<IN>) {
    	chomp;
    	my @F=split(/\t/, $_);
    	if ($F[0] eq $F[1]) { next; }
    	if (!$best{$F[0]}) {
        	$best{$F[0]} = $_;
        	$max{$F[0]} = $F[11];
    	}
    	else {
        	if ($F[11] > $max{$F[0]}) {
            	$best{$F[0]} = $_;
            	$max{$F[0]} = $F[11];
        	}
    	}
	} 
	close IN;
	open (IN, "$scaffold_dir/annot/$method.list") or die "can't open $scaffold_dir/annot/$method.list file\n";
	while (<IN>) {
		chomp;
		my @F=split(/\t/, $_);
		$list{$F[1]} = $F[0];
	}
	close IN;	
	open (OUT, ">$dirname/$blast_results.bestOrthos") or die "can't open $dirname/$blast_results.bestOrthos file\n";
	print OUT "Gene ID\tOrthogroup ID\n"; 
	foreach (keys %best) {
    	my @F = split(/\t/, $best{$_});
    	print OUT "$F[0]\t$list{$F[1]}\n";
	}
	close OUT;
	if ($classifier eq "blastp") {
		my $orthogroup_assignment = "$blast_results.bestOrthos";
		print "-- ".localtime()." - Getting orthogroup annotation summaries\n\n";
		get_annot_summary ( $orthogroup_assignment, $dirname, $scaffold, $method, $super_orthogroups, $scaffold_dir );
	}
}

sub  get_best_hmmscan_orthos {
	my ( $classifier, $hmmscan_results, $dirname, $scaffold, $method, $super_orthogroups, $scaffold_dir ) = @_;
	my %hits;
	open (IN, "$dirname/$hmmscan_results") or die "can't open $dirname/$hmmscan_results file\n";
	while (<IN>) {
    	if (/^#/){next;}
        my @F = split(/\s+/, $_);
        $hits{$F[2]}{$F[0]} = $F[5];
	}    
	close IN;
	open (OUT, ">$dirname/$hmmscan_results.bestOrthos") or die "can't open $dirname/$hmmscan_results.bestOrthos file\n";
	print OUT "Gene ID\tOrthogroup ID\n"; 
	for my $hit (keys %hits) {
    	my $score = 0;
    	my $best_target;
    	for my $target (keys %{$hits{$hit}}) {
        	if ($hits{$hit}{$target} >= $score) {
            	$score = $hits{$hit}{$target};
            	$best_target = $target; 
        	}
    	}    
    	print OUT "$hit\t$best_target\n";
 	}
	close OUT;
	if ($classifier eq "hmmscan") {
		my $orthogroup_assignment = "$hmmscan_results.bestOrthos";
		print "-- ".localtime()." - Getting orthogroup annotation summaries\n\n";
		get_annot_summary ( $orthogroup_assignment, $dirname, $scaffold, $method, $super_orthogroups, $scaffold_dir );
	}
}

sub get_blast_hmmscan_orthos {
	my ( $dirname, $scaffold, $method, $super_orthogroups, $scaffold_dir ) = @_;
	my (%blastp, %hmmscan, %genes);	
	if ( (!$super_orthogroups) or ($super_orthogroups ne "min_evalue") or ($super_orthogroups ne "avg_evalue") ) { $super_orthogroups = "min_evalue" }
	opendir (DIR, "$dirname") or die "can't open $dirname directory\n";
	while (my $filename = readdir(DIR)) {
		if ($filename =~ /^proteins\.blastp\.\w+\.\d+\.bestOrthos$/){ 
			open (IN, "$dirname/$filename") or die "can't open $dirname/$filename file\n";
			while (<IN>) {
				chomp;
				if (/^Gene/) {next;}
				my @F = split(/\t/, $_);
				$blastp{$F[0]} = $F[1];
				$genes{$F[0]} = $F[0];
			}
			close IN;
		}
		if ($filename =~ /^proteins\.hmmscan\.\w+\.\d+\.bestOrthos$/){ 
			open (IN, "$dirname/$filename") or die "can't open $dirname/$filename file\n";
			while (<IN>) {
				chomp;
				if (/^Gene/) {next;}
				my @F = split(/\t/, $_);
				$hmmscan{$F[0]} = $F[1];
				$genes{$F[0]} = $F[0];
			}
			close IN;
		}
	}
	closedir DIR;
	open (OUT, ">$dirname/proteins.both.$scaffold.bestOrthos") or die "can't open $dirname/protein.both.$scaffold.bestOrthos file\n";
	print OUT "Gene ID\tOrthogroup ID\n";
	foreach (sort keys %genes) {
		if (!$blastp{$_} and $hmmscan{$_}) { print OUT "$_\t$hmmscan{$_}\n"; next; }
		elsif ($blastp{$_} and !$hmmscan{$_}) { print OUT "$_\t$blastp{$_}\n"; next; }
		elsif ($blastp{$_} == $hmmscan{$_}) { print OUT "$_\t$blastp{$_}\n"; next }
		else { print OUT "$_\t$hmmscan{$_}\n"; }
	}
	close OUT;
	my $orthogroup_assignment = "proteins.both.$scaffold.bestOrthos";
	print "-- ".localtime()." - Getting orthogroup annotation summaries\n\n";
	get_annot_summary ( $orthogroup_assignment, $dirname, $scaffold, $method, $super_orthogroups, $scaffold_dir );	
}

sub get_annot_summary {
	my ( $orthogroup_assignment, $dirname, $scaffold, $method, $super_orthogroups, $scaffold_dir ) = @_;
	my %annot;
	open (OUT, ">$dirname/$orthogroup_assignment.summary") or die "can't open $dirname/$orthogroup_assignment.summary file\n";	
	open (IN, "$scaffold_dir/annot/$method.$super_orthogroups.summary") or die "can't open $scaffold_dir/annot/$method.$super_orthogroups.summary file\n";
	while (<IN>) {
		chomp;
		if ($_ =~ /^Orthogroup/) { print OUT "Gene ID\t$_\n"; next; }
		$_ =~ /(\d+)\s+(.*)/;
		$annot{$1} = $_;
	}
	close IN;
	open (IN, "$dirname/$orthogroup_assignment") or die "can't open $dirname/$orthogroup_assignment file\n";
	while (<IN>) {
		chomp;
		if ($_ =~ /^Gene/) { next; }
		my @F = split(/\t/, $_);
		if ($annot{$F[1]}) { print OUT "$F[0]\t$annot{$F[1]}\n"; }
	}
	close IN;
	close OUT;	
}

sub get_single_copy_orthogroups {
	my ( $single_copy_custom, $single_copy_taxa, $taxa_present, $classifier, $scaffold, $method, $dirname, $config_dir ) = @_;
	print "-- ".localtime()." - Determining single copy orthogroups\n\n";	
	my $annot_summary = "proteins.$classifier.$scaffold.bestOrthos.summary";
	open (OUT, ">$dirname/$annot_summary.singleCopy") or die "can't open $dirname/$annot_summary.singleCopy file\n";
	open (ANNOT, "$dirname/$annot_summary") or die "can't open $dirname/$annot_summary file\n";
	if ($single_copy_custom) {
		my %genes;
		my $annot_field = 1;
		open (IN, "$single_copy_custom") or die "can't open file $single_copy_custom\n";
		while (<IN>) {
			chomp; 
			if ($_ !~ /^\w+/) { next; }
			my @F = split(/\=/, $_);
			$annot_field++;
			$genes{$annot_field} = $F[1];
		}
		close IN;
		while (<ANNOT>) {
			chomp;
			if ($_ =~ /^Gene\s+ID(.*)/) { print OUT "$_\n"; next; }
			my @F = split(/\t/, $_);
			my $matches = 0;
			for my $field (2..$annot_field) {
				if ( ($genes{$field} == -1) or ($genes{$field} == $F[$field]) ) { $matches++; }
			}
			if ($matches == keys(%genes)) { print OUT "$_\n"; } 
		}
	}
	if ($single_copy_taxa and $taxa_present) {
		while (<ANNOT>) {
			chomp;
			if ($_ =~ /^Gene\s+ID(.*)/) { print OUT "$_\n"; next; }
			my @F = split(/\t/, $_);
			$scaffold =~ /(\d+)(.*)/;
			my $annot_field = $1 + 1;
			my $single_copy = 0;
			my $taxa = 0;
			for my $field (2..$annot_field) {
				if ($F[$field] == 1) { $single_copy++; }
				if ($F[$field] != 0) { $taxa++; }
			}
			if ($single_copy >= $single_copy_taxa and $taxa >= $taxa_present) { print OUT "$_\n"; } 
		}
	}
	if ($single_copy_taxa and !$taxa_present) {
		while (<ANNOT>) {
			chomp;
			if ($_ =~ /^Gene\s+ID(.*)/) { print OUT "$_\n"; next; }
			my @F = split(/\t/, $_);
			$scaffold =~ /(\d+)(.*)/;
			my $annot_field = $1 + 1;
			my $single_copy = 0;
			for my $field (2..$annot_field) {
				if ($F[$field] == 1) { $single_copy++; }
			}
			if ($single_copy >= $single_copy_taxa) { print OUT "$_\n"; } 
		}		
	}
	close ANNOT;
	close OUT;
}

sub get_orthogroup_fasta {
	my ( $proteins, $coding_sequences, $single_copy_custom, $single_copy_taxa, $classifier, $scaffold, $dirname ) = @_;
	print "-- ".localtime()." - Retrieving orthogroup fasta files\n\n";
	my (%orthos, %pep, %cds);
	my $orthogroups = "$dirname/orthogroups_fasta";
	make_directory($orthogroups);
	my $orthogroup_assignment = "proteins.$classifier.$scaffold.bestOrthos";
	open (IN, "$dirname/$orthogroup_assignment") or die "can't open $dirname/$orthogroup_assignment file\n";
	while (<IN>) {
		chomp;
		if ($_ =~ /^Gene/) { next; }
		my @F = split(/\t/, $_);
		$orthos{$F[1]}{$F[0]} = $F[0];
	}
	close IN;
	%pep = get_sequences ($proteins);
	if ($coding_sequences) { %cds = get_sequences ($coding_sequences); }
	my ($ortho_id, $seq_id);
	foreach $ortho_id (keys %orthos) {
		open (PEP, ">$orthogroups/$ortho_id.faa") or die "can't open $orthogroups/$ortho_id.faa file\n";
		if ($coding_sequences) { open (CDS, ">$orthogroups/$ortho_id.fna") or die "can't open $orthogroups/$ortho_id.fna file\n"; }
		foreach $seq_id (sort keys %{$orthos{$ortho_id}}) {
			$pep{$seq_id} =~ s/.{80}(?=.)/$&\n/g;
			print PEP ">$seq_id\n$pep{$seq_id}\n";
			if ($coding_sequences) {
				$cds{$seq_id} =~ s/.{80}(?=.)/$&\n/g;
				print CDS ">$seq_id\n$cds{$seq_id}\n";
			}				
		}
		close PEP;
		close CDS;
	}
	if ($single_copy_custom or $single_copy_taxa) {
		my $single_copy = "$dirname/single_copy_fasta";
		make_directory($single_copy);
		my $single_copy_selection = "proteins.$classifier.$scaffold.bestOrthos.summary.singleCopy";
		my (%single_copy_orthos, $single_copy_ortho_id);
		open (IN, "$dirname/$single_copy_selection") or die "can't open $dirname/$single_copy_selection file\n";
		while (<IN>) {
			chomp;
			if ($_ =~ /^Gene\s+ID(.*)/) { next; }
			my @F = split(/\t/, $_);
			$single_copy_orthos{$F[1]} = $F[1];
		}
		close IN;
		foreach $single_copy_ortho_id (keys %single_copy_orthos) {
			system "cp $orthogroups/$single_copy_ortho_id.*  $single_copy/";
		}
	}	
}

sub get_sequences {
	my ( $file ) = @_;
	my (%sequences, $id);
	open (IN, "$file") or die "can't open $file file\n";
	while (<IN>) {
		if ($_ =~ />(\S+)/){ $id = $1; }
		else { s/\s+//g; $sequences{$id} .= $_; }
	}
	close IN;
	return %sequences; 
}


