#!/lab/bin/perl -w

use strict;
use Getopt::Std;

my %opts;
getopt("i:f:", \%opts);
if (!defined $opts{'i'}  &&  !defined $opts{f}  &&  @ARGV != 1) {
    print "Syntax: fa_select <identifier>\n",
	  "        fa_select -i <file of identifiers>\n",
	  "        fa_select -f <fasta file>\n\n";
    exit 1;
}
if (!defined $opts{'f'}) {
    $opts{'f'} = '-';
}

sub BinarySearch($$) {
    my ($target, $listRef) = @_;

    my ($fwd, $rev);

    $fwd = 0;
    $rev = scalar(@$listRef) - 1;
    while ($fwd <= $rev) {
	my $mid = int(($fwd + $rev)/2);
	my $res = substr($target, 0, length($listRef->[$mid])) cmp
	          substr($listRef->[$mid], 0, length($target));
	if ($res < 0) {
	    $rev = $mid - 1;
	} elsif ($res > 0) {
	    $fwd = $mid + 1;
	} else {
	    return $mid;
	}
    }

    return undef;
}

my @ids;
if (!defined $opts{'i'}  &&  @ARGV != 1) {
    $opts{'i'} = '-'; #read from stdin
}
if (defined $opts{'i'}) {
    open IDFILE, $opts{'i'};
    while (<IDFILE>) {
	chomp;
	push @ids, $_;
    }
    close IDFILE;
} else {
    push @ids, $ARGV[0];
}

@ids = sort {$a cmp $b} @ids;

my $c = 0;
my $selected = 0;
open FAFILE, $opts{'f'} or die "$opts{f}: $!";
while (<FAFILE>) {
    if (/^>/) {
#	print STDERR "$c\r" if (++$c % 1000 == 0);
	if (@ids == 0) { #if our list of ids is empty, we're done!
	    last;
	}
	$selected = 0;
	my ($id) = /^>(.+?)\s/;
	if (defined (my $i = BinarySearch($id, \@ids))) {
	    print $_;
	    $selected = 1;
	    splice @ids, $i, 1; #we found it, so remove from list
	}
    } elsif ($selected) {
	print $_;
    }
}
close FAFILE;
