#! -w
#
# select-by-key-residues  -residues foo.res
#	[-mismatch m.a2m]	[-err m.err]
#	 < x.a2m > y.a2m
#
# selects those sequences in the a2m file in standard input
#	that have the residues specified in foo.res
#	and outputs them to standard out.
# If the optional -mismatch argument is provided, then rejected
#	sequences are sent there.
# If the optional -err argument is provided, then the reasons for
#	rejection are sent there
#
# The format for the foo.res file is a white-space separated
# list of residues that must be present.
# For example,
#	N150 AS180 P-225
# would mean that alignment column 150 (with 1 as the first column)
# must have Asparagine, 180 must have Alanine or Serine, and 225 must
# have Proline or a gap.


{   use FileHandle;
    STDOUT->autoflush(1);
    STDERR->autoflush(1);
    use English;
    use File::Basename;
    use lib dirname($PROGRAM_NAME);
    use SamT2k;
    

    sub process_command_line;
    sub read_res_file($);
    sub check_seq($$);
    sub get_alignment_columns($);    
    # globals
    undef $mis_a2m;
    undef $mis_why;
    undef $res_file;
    undef @locs;
    undef @res;
    
    &process_command_line;
    
    &read_res_file($res_file);
    # now @locs has a list of key residue locations (alignment columns)
    # $res[$loc] has string of legal residues at location $loc.
    
    open (MIS_A2M, "> $mis_a2m") if (defined $mis_a2m);
    open (MIS_WHY, "> $mis_why") if (defined $mis_why);

    
    while (<STDIN>)
    {   if (/[>]/) 
    	{   &check_seq($id, $seq) if (defined($seq));
	    $id = $_;
	    undef $seq;
	}
	else
	{    $seq .= $_;
	}
    }
    &check_seq($id, $seq) if (defined($seq));
    close MIS_A2M if (defined $mis_a2m);
    close MIS_WHY if (defined $mis_why);
    
}

sub process_command_line
{   local($i) = 0;

    while ($i <= $#ARGV) 
    {	$_ = $ARGV[$i++];
	if (/^-mis/)	{   $mis_a2m  = absolute( $ARGV[$i++]);
			}
	elsif (/^-err/)	{   $mis_why  = absolute( $ARGV[$i++]);
			}
	
	elsif (/^-res/)	{   $res_file  = absolute( $ARGV[$i++]);
			    (-e $res_file) ||
				fatal("$res_file doesn't exist");
			}
	else	{ fatal("Invalid option '$ARGV[$i-1]'");
		}
    }

    # Make certain that required parameters were provided.
    defined($res_file) 
    	|| fatal("Must specify '-residues res_file' with $PROGRAM_NAME");

}

# create array @locs of important column positions
# and array @res, where $res[$loc] is a string of legal
# characters for alignment column $loc.
sub read_res_file($)
{
    my ($res_file) = @_;
    open (RES, "<$res_file");
    
    LINE: while(<RES>)
    {   local @words= split;	# split on white space
    	foreach $word (@words)
	{   $word = uc($word);
	    next LINE if $word =~ /[#]/;
	    $word =~ /([-A-Z]+)(\d+)/
	    	|| fatal("$word not a recognized residue specification.");
	    $loc = $2;
	    push @locs, $loc;
	    $res[$loc] = $1;
	}
    }
    close RES;
}

# check to see if a sequence has the required residues.
# if so, output its id line and the seq to standard out,
# otherwise, output to MIS_A2M (when $mis_a2m defined).
sub check_seq($$)
{
    my ($id, $seq) = @_;
    my $cols = &get_alignment_columns($seq);
    foreach $loc(@locs)
    {   $has = substr($cols,$loc,1);
         if (!defined($has) ||
	    	index($res[$loc], $has) < 0)
    	{ # mismatch
	    if (defined $mis_why)
	    {	print MIS_WHY "$id has $has$loc, not $res[$loc]$loc\n";
	    }
	    if (defined $mis_a2m)
	    {    print MIS_A2M $id . $seq;
	    }
	    return;
	}
    }
    # no mismatches
    print $id . $seq;
}

# converts a string into a shorter string containing only alignment columns,
# with column zero added and always being "-"
sub get_alignment_columns($)
{
   my ($seq) = @_;
   return join("","-", split(/[^-A-Z]+/, $seq));
}
