#!/usr/bin/perl
#pk

##########################################################################################
#	  This file is part of proteinortho.
#	  (C) 2009 Marcus Lechner
# 
#	  proteinortho is free software; you can redistribute it and/or modify
#	  it under the terms of the GNU General Public License as published
#	  by the Free Software Foundation; either version 2, or (at your
#	  option) any later version.
#
#	  proteinortho is distributed in the hope that it will be useful, but
#	  WITHOUT ANY WARRANTY; without even the implied warranty of
#	  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
#	  General Public License for more details.
#
#	  You should have received a copy of the GNU General Public License
#	  along with proteinortho; see the file COPYING.  If not, write to the
#	  Free Software Foundation, Inc., 59 Temple Place - Suite 330,
#	  Boston, MA 02111-1307, USA.	
##########################################################################################

##########################################################################################
# About
##########################################################################################
# 
# @author Paul Klemm
# @email klemmp@staff.uni-marburg.de
# @company Bioinformatics, University of Leipzig
# @version 1
# @date 8-04-2019
#
##########################################################################################

use warnings;
use strict;

my $usage = <<'ENDUSAGE';
proteinortho_grab_proteins.pl        greps all genes/proteins of a given fasta file
 
SYNOPSIS
 
proteinortho_grab_proteins.pl (options) QUERY FASTA1 (FASTA2 ...)

	QUERY	proteinortho.tsv FILE or search STRING or '-' for STDIN:
		a)	proteinortho output file (.tsv). This uses by default the -exact option.
		b)	string of one identifier e.g. 'tr|asd3|asd' OR multiple identifier separated by ',' (-F=)
	FASTA*	fasta file(s) (database)

	(options):
		-tofiles, -t  print everything to files instead of stdout files are called OrthoGroup**.fasta for a proteinortho.tsv file
		-E            enables regex matching otherwise the string is escaped (e.g. | -> \|)
		-exact        search patters are extended with a \b, that indicates end of word.
		-source, -s   adds the filename (FASTA1,...) to the found gene-name
		-F=s          char delimiter for multiple identifier if QUERY is a string input (default: ',')

DESCRIPTION
 
	This script finds and extract all given identifier of a list of fasta files. 
	The identifier can be provided as a simple string 'BDNF1', regex string 'BDNF*' 
	or in form of a proteinortho output file (myproject.proteinortho.tsv).
       
	Example:
 
 	# 1. most simple call:

	perl proteinortho_grab_proteins.pl 'BDNF1' *.faa

		STDOUT:
			>BDNF1 Brain derived neurotrophic factor OS=human(...)
			MNNGGPTEMYYQQHMQSAGQPQQPQTVTSGPMSHYPPAQPPLLQPGQPYSHGAPSPYQYG
			>BDNF15 Brain derived neurotrophic factor OS=human(...)
			MAFPLHFSREPAHAIPSMKAPFSRHEVPFGRSPSMAIPNSETHDDVPPPLPPPRHPPCTN

	    The second hit BDNF15 is reported since it also contains 'BDNF1' as a substring. 
	    To prevent such a behaviour use proteinortho_grab_proteins.pl -E 'BDNF1\b'. 
	    The \b marks the end of a word and -E enables regex expressions.

	    Or simply add -exact: perl proteinortho_grab_proteins.pl -exact 'BDNF1' *.faa

 	# 2. multiple ids:

	perl proteinortho_grab_proteins.pl 'BDNF1,BDNF2,BDNF3' *.faa

 	# 3. more complex regex search:

	perl proteinortho_grab_proteins.pl -E 'B?DNF[0-3]3+' *.faa

		This finds: BDNF13, BDNF23, DNF13, DNF033, ... 

 	# 4. proteinortho tsv file and write output to files:

	proteinortho_grab_proteins.pl -tofiles myproject.proteinortho.tsv test/*.faa

		This will produce the files: OrthoGroup0.fasta, OrthoGroup1.fasta, OrthoGroup2.fasta, ...
		Each fasta file contains all genes of one orthology group (one line in myproject.proteinortho.tsv)
 
ENDUSAGE

my $query;
my $help;
my $tofiles=0;
my $justid;
my $prefix=">";
my $doregex=0;
my $source=0;
my $exact=0;
my $del=',';

my @ARGViddone=[];
my $ARGViddone_counter=scalar(@ARGV);
for(my $v = 0 ; $v < scalar @ARGV ; $v++){
	$ARGViddone[$v]=1;
	if($ARGV[$v] =~ m/--?(help|h)$/){$help=1;}
	elsif($ARGV[$v] =~ m/^--?(tofiles|t)$/){$tofiles=1;}
	elsif($ARGV[$v] =~ m/^--?(source|s)$/){$source=1;}
	elsif($ARGV[$v] =~ m/^--?F=(.*)$/){$del=$1;}
	elsif($ARGV[$v] =~ m/^--?E$/){$doregex=1;}
	elsif($ARGV[$v] =~ m/^--?exact$/){$exact=1;}
	elsif($ARGV[$v] =~ m/^-.+/){print $usage; print STDERR "ERROR: invalid option ".$ARGV[$v]."!\n\n";exit(1);}
	elsif(!defined($query)){$query = $ARGV[$v];}
	else{$ARGViddone[$v]=0;$ARGViddone_counter--;}
}
if ($help){
    print $usage;
    exit(0);
}
my $fail="";
if (scalar(@ARGV) == 0){
    $fail.="ERROR: no arguments found!\n";
}
if (!defined($query)){
    $fail.="ERROR: no QUERY provided!\n";
}
if ( $ARGViddone_counter==scalar(@ARGV) ){
    $fail.="ERROR: no FASTA files provided!\n";
}
if($fail){
	print $usage.$fail;
	exit(1);
}

my %qdata;
# my $qdata_count = {};

unless(open(FH,'<',$query)) {
	if($query eq "-"){
		foreach my $line (<>) 
		{
			my $linel="";
			if(!$doregex){
				$linel=quotemeta($line);
			}
			if($exact){
				$linel.='\b';
			}
			$qdata{$linel}=$linel;
			# $qdata_count{$linel}=0;
		}
	}else{
		my @sp = split($del,$query); 
		for(my $v = 0 ; $v < scalar @sp ; $v++){
			if(!$doregex){
				$sp[$v]=quotemeta($sp[$v]);
			}
			if($exact){
				$sp[$v].='\b';
			}
			$qdata{$sp[$v]}=$sp[$v];
			# $qdata_count{$sp[$v]}=0;
		}
	}
}else{
	my $orthogroupcounter=0;
	while(<FH>){
		chomp($_);
		if($_ eq "" || substr($_,0,1) eq "#"){next;}
		my @sp = split(/[\t,]+/,$_);
		if(scalar(@sp)>3){
			for(my $v = 3 ; $v < scalar @sp ; $v++){
				if($sp[$v] eq "*" || $sp[$v] eq ""){next;}
				$qdata{$sp[$v].'\b'}="OrthoGroup".$orthogroupcounter;
				# $qdata_count{$sp[$v]}=0;
			}
		}
		$orthogroupcounter++;
	}
	close(FH);
}

my %outputfasta;

for(my $v = 0 ; $v < scalar @ARGV ; $v++){

	if($ARGViddone[$v]){next;}

	my $basename = $ARGV[$v];
	if($basename =~ m/\/([^\/]+)$/){
		$basename=$1;
	}

	open(FH,'<',$ARGV[$v]);
	my $geneprintswitch = "";

	while(<FH>){

		if($_ eq ""){continue;}
		if(substr($_,0,1) eq $prefix ){
			$geneprintswitch = "";

			foreach my $regexv (keys %qdata) { 
				if( $_ =~ $regexv ){
				
					my $headerstr=$_;
					if($source){
						$headerstr=$_;
						chomp($headerstr);
						$headerstr.=" ".$basename."\n";
					}

					$geneprintswitch = $qdata{$regexv};
					# $qdata_count{$key} = $qdata_count{$key}+1;

					if($tofiles){
						if(exists $outputfasta{$geneprintswitch.".fasta"}){
							$outputfasta{$geneprintswitch.".fasta"}.=$headerstr;
						}else{
							$outputfasta{$geneprintswitch.".fasta"}=$headerstr;
						}
					}else{
						print $headerstr;
					}

					last;
				}
			}
		}else{
			if($geneprintswitch ne ""){
				if($tofiles){
					$outputfasta{$geneprintswitch.".fasta"}.=$_;
				}else{
					print $_;	
				}
			}
		}
	}
	close(FH);
}

foreach my $key (keys %outputfasta) { 
	open(FHOUT,'>',$key);
	print FHOUT $outputfasta{$key};
	close(FHOUT);
}

# print STDERR "# key\toccurences\n";
# foreach my $key (keys %qdata_count) { 
	#if($qdata_count{$key} != 1){
		# print STDERR $key."\t".$qdata_count{$key}."\n";
	#}
# }
