#/usr/bin/perl -w
# pk

use Getopt::Long; # for parameter specification on the command line
use File::Basename; # basename()


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

	QUERY	identifier FILE or search query STRING:
		a)	string of one identifier e.g. 'tr|asd3|asd' OR multiple identifier separated by ',' OR '-' for STDIN
		b)	(-file QUERY) file with the ids. The file contains either just 
			identifier or tabseparated a mapping:
			e.g. POP1 to the 3 identifier 
			tr|asd|asd, tr|asd2|asd and tr|asd3|asd.
	INFILE	file containing the query ids (database)

	optional:
		-tofiles, -t 	print everything to files instead of files 
		-prefix=s 	prefix of each identifier (is regex interpreted)
				[default: -prefix=">"]
		-E    	enables regex matching otherwise the string is escaped (e.g. | -> \|)
		-source, -s    	adds the filename to the found gene-name
		-F=s 	char delimiter for multiple identifier (default: ',')

DESCRIPTION
 
	This script finds all ids of a list of fasta files with identifier
	provided in a different file and saves the output to seperate files named
	ID_FASTA where ID is the found identifier and FASTA the 
	name of the fasta file that contains this identifier.
	STDERR contains ids and the number of times they are found.
       
	Example:
 
 	# 1. most simple call:

	perl grepfasta.pl 'tr|asd|asd' *.faa

		STDOUT:
			>tr|asd|asd Cell pattern formation-associated protein stuA OS=Acremon(...)
			MNNGGPTEMYYQQHMQSAGQPQQPQTVTSGPMSHYPPAQPPLLQPGQPYSHGAPSPYQYG
			>tr|asd2|asd 1-phosphatidylinositol 3-phosphate 5-(...)
			MAFPLHFSREPAHAIPSMKAPFSRHEVPFGRSPSMAIPNSETHDDVPPPLPPPRHPPCTN

 	# 2. regex search:

	perl grepfasta.pl -E '.*a?sd[0-3].*' *.faa

 	# 3. multiple ids:

	perl grepfasta.pl 'tr|asd|asd,tr|asd2|asd' *.faa

 	# 4. id file and write output to files:

	cat test.idlist | grepfasta.pl - test.faa

		-----test.idlist:---------- (does not need the first column)
		POP1	tr|asd|asd
		POP1	tr|asd1|asd
		POP1	tr|asd2|asd
		POP2	tr|asd3|asd
		
 	# 5. id file and write output to files:

	grepfasta.pl -tofiles test.idlist test.faa

		-----test.idlist:---------- (does not need the first column)
		POP1	tr|asd|asd
		POP1	tr|asd1|asd
		POP1	tr|asd2|asd
		POP2	tr|asd3|asd
		
		------test.faa:------------
		>tr|asd|asd Cell pattern formation-associated protein stuA OS=Acremon(...)
		MNNGGPTEMYYQQHMQSAGQPQQPQTVTSGPMSHYPPAQPPLLQPGQPYSHGAPSPYQYG
		>tr|asd3|asd Histone H4 OS=Acremonium chrysogenum (strain ATCC 11550 (...)
		MTGRGKGGKGLGKGGAKRHRKILKDNIQGITKPAIRRLARRGGVKRISAMIYEETRGVLK
		>tr|A0A086SUI7|A0A086SUI7_ACRC1 Uncharacterized protein OS=Acremonium(...)
		MAFPLHFSREPAHAIPSMKAPFSRHEVPFGRSPSMAIPNSETHDDVPPPLPPPRHPPCTN
		>tr|asd2|asd 1-phosphatidylinositol 3-phosphate 5-(...)
		MAFPLHFSREPAHAIPSMKAPFSRHEVPFGRSPSMAIPNSETHDDVPPPLPPPRHPPCTN

		OUTPUT (2 files)

		------POP1_test:------------
		>tr|asd|asd Cell pattern formation-associated protein stuA OS=Acremon(...)
		MNNGGPTEMYYQQHMQSAGQPQQPQTVTSGPMSHYPPAQPPLLQPGQPYSHGAPSPYQYG
		>tr|asd2|asd 1-phosphatidylinositol 3-phosphate 5-(...)
		MAFPLHFSREPAHAIPSMKAPFSRHEVPFGRSPSMAIPNSETHDDVPPPLPPPRHPPCTN
		
		------POP2_test:------------
		>tr|asd3|asd Histone H4 OS=Acremonium chrysogenum (strain ATCC 11550 (...)
		MTGRGKGGKGLGKGGAKRHRKILKDNIQGITKPAIRRLARRGGVKRISAMIYEETRGVLK

		STDERR: 	
		tr|asd|asd	1
		tr|asd1|asd	1
		tr|asd2|asd	0
		tr|asd3|asd	1
 
ENDUSAGE

my $query,$help,$tofiles=0,$justid,$prefix=">",$doregex=0,$source=0,$del=',';
GetOptions('help!'=>\$help,'h!'=>\$help,
	'tofiles!'=>\$tofiles,
	't!'=>\$tofiles,
	'prefix=s'=>\$prefix,
	'source'=>\$source,
	's'=>\$source,
	'T=s'=>\$del,
	'E'=>\$doregex);

$query = $ARGV[0];
shift(@ARGV);

if ($help || scalar(@ARGV) == 0 || !$query ){
    print $usage;
    exit(1);
}

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

unless(open(FH,'<',$query)) {
	if($query eq "-"){
		foreach $line (<>) 
		{
			$linel="";
			if(!$doregex){
				$linel=quotemeta($line);
			}
			$qdata{$linel}=$linel;
			$qdata_count{$linel}=0;
		}
	}else{
		my @sp = split($del,$query);
		for($v = 0 ; $v < scalar @sp ; $v++){
			if(!$doregex){
				$sp[$v]=quotemeta($sp[$v]);
			}
			$qdata{$sp[$v]}=$sp[$v];
			$qdata_count{$sp[$v]}=0;
		}
	}
}else{
	while(<FH>){
		chomp($_);
		my @sp = split("\t",$_);
		if(scalar(@sp)==2){
			$qdata{$sp[1]}=$sp[0];
			$qdata_count{$sp[1]}=0;
		}elsif(scalar(@sp)==1 && $_ ne ""){
			$qdata{$_}=basename($query);
			$qdata_count{$_}=0;
		}
	}
	close(FH);
}

for($v = 0 ; $v < scalar @ARGV ; $v++){
	open(FH,'<',$ARGV[$v]);
	my $geneprintswitch = "";

	while(<FH>){
	
		#chomp($_);
		if($_ eq ""){
			continue;
		}
		if(substr($_,0,1) eq $prefix ){
			$geneprintswitch = "";

			foreach my $key (keys %qdata) { 
				my $regexv = $key;
				if( $_ =~ $regexv ){
					$geneprintswitch = $qdata{$key};
					$qdata_count{$key}=$qdata_count{$key}+1;
					if($tofiles){
						chomp($_);
						if($qdata_count{$key} == 1){
							system("echo '$_' >".$geneprintswitch."_".basename($ARGV[$v])).".fasta";
						}else{
							system("echo '$_' >>".$geneprintswitch."_".basename($ARGV[$v])).".fasta";
						}
					}else{
						if($source){
							chomp($_);
							print $_." ".basename($ARGV[$v])."\n";
						}else{
							print $_;
						}
					}
				}
			}
		}else{
			if($geneprintswitch ne ""){
				if($tofiles){
					chomp($_);
					system("echo '$_' >>".$geneprintswitch."_".basename($ARGV[$v])).".fasta";
				}else{
					print $_;	
				}
			}
		}
	}
	close(FH);
}
print STDERR "# key\toccurences\n";
foreach my $key (keys %qdata_count) { 
	#if($qdata_count{$key} != 1){
		print STDERR $key."\t".$qdata_count{$key}."\n";
	#}
}
