#!/usr/bin/env perl

# MANUAL FOR ambiguous_2_unambiguous.pl

=pod

=head1 NAME

ambiguous_2_unambiguous.pl -- takes in a multifasta of NTs with ambiguous bases and will produce all (or some) versions of that sequence.

=head1 SYNOPSIS

 ambiguous_2_unambiguous.pl --fasta=/Path/to/infile.fasta --out=/Path/to/output.fasta [ -n 10 ] 
                     [--help] [--manual]

=head1 DESCRIPTION

 This script will take in a multi-fasta file and for sequences with ambiguous bases it will
 produce all iterations of that sequence. IF YOU WOULD LIKE you can tell the program to never
 produce more than -n sequences. You may want this if you just want to save time and space
 and just want a quick glance at the sequence.
 
=head1 OPTIONS

=over 3

=item B<-f, --fasta>=FILENAME

Input file in FASTA format. (Required) 

=item B<-o, --out>=FILENAME

Output file in FASTA format. (Required) 

=item B<-n, --num>=INT

Max number of unambiguous sequences to spit out for each ambiguous sequence. (Default = 10)

=item B<-h, --help>

Displays the usage message.  (Optional) 

=item B<-m, --manual>

Displays full manual.  (Optional) 

=back

=head1 DEPENDENCIES

Requires the following Perl libraries.



=head1 AUTHOR

Written by Daniel Nasko, 
Center for Bioinformatics and Computational Biology, University of Maryland.

=head1 REPORTING BUGS

Report bugs to dnasko@umiacs.umd.edu

=head1 COPYRIGHT

Copyright 2017 The S2FAST Team.
License GPLv3+: GNU GPL version 3 or later <http://gnu.org/licenses/gpl.html>.  
This is free software: you are free to change and redistribute it.  
There is NO WARRANTY, to the extent permitted by law.  

=cut


use strict;
use warnings;
use Getopt::Long;
use File::Basename;
use Pod::Usage;

#ARGUMENTS WITH NO DEFAULT
my($fasta,$outfile,$help,$manual);
#ARGS WITH DEFAULT
my $max_seqs = 10; ## by default we'll only give back 10 unambiguous sequences for each ambiguous sequence

GetOptions (	
				"f|fasta=s"	=>	\$fasta,
				"o|out=s"	=>	\$outfile,
                                "n|num=s"       =>      \$max_seqs,
				"h|help"	=>	\$help,
				"m|manual"	=>	\$manual) || pod2usage({-exitval => 2, -verbose => 1, -output => \*STDERR});

# VALIDATE ARGS
pod2usage(-verbose => 2)  if ($manual);
pod2usage( {-exitval => 0, -verbose => 2, -output => \*STDERR} )  if ($help);
pod2usage( -msg  => "\n\n ERROR!  Required argument --fasta not found.\n\n", -exitval => 2, -verbose => 1)  if (! $fasta );
pod2usage( -msg  => "\n\n ERROR!  Required argument --out not found.\n\n", -exitval => 2, -verbose => 1)  if (! $outfile);
$max_seqs = int($max_seqs);
if ($max_seqs < 1) { die "\n Error: You cannot select a -n < 1 !!! \n"; }
if ($max_seqs > 10000) { die "\n Error: You've selected a high --num parameter, >10000. Select a more reasonable value.\n" }

### GLOBALS
my ($header,$seq) = ("","");
my $fasta_fh;

## Initialize
my %AmbigBases = construct_hash();
check_nt_fasta($fasta); ## Check that the file is a valid FASTA file

open(my $ofh, '>', "$outfile") || die "\n Cannot open the file: $outfile\n";
if ($fasta =~ m/\.gz$/) { ## if a gzip compressed infile
    open($fasta_fh, '-|', 'gunzip', '-c', $fasta) || die "\n\n Cannot open the input file: $fasta\n\n";
}
else { ## If not gzip comgressed
    open($fasta_fh,'<', "$fasta") || die "\n\n Cannot open the input file: $fasta\n\n";
}
while(<$fasta_fh>) {
    $_ =~ s/\r[\n]*/\n/gm; ## Convert other newline characters to unix '\n'
    chomp; ## strip off new line
    if ($_ =~ m/^>/) { ## If line starts with a '>'
    	unless ($seq eq "") { ## If this is NOT the first time we've seen a header line
    	    my @Unambig = find_ambiguous($seq, $max_seqs); ## Get $max_seqs unambiguous sequences in the Array
    	    for (my $i=0; $i<scalar(@Unambig); $i++) {
    		if ($i > $max_seqs-1) { last; }
    		my $id = $i+1;
    		print $ofh ">" . $header . "_unambig_" . $id . "\n" . $Unambig[$i] . "\n";
    	    }
    	}
    	$header = $_; ## reset header to match the line we're on
    	$header =~ s/^>//; ## strip off '>'
    	$header =~ s/ .*//; ## strip off everything after the first space. BLAST does this.
    	$seq = ""; ## reset sequence to empty
    }
    else {
    	$seq = $seq . $_; ## collect sequence information on sequence lines
    }
}
close($fasta_fh);
## Need to run this one more time for the last sequence in the file
my @Unambig = find_ambiguous($seq, $max_seqs);
for (my $i=0; $i<scalar(@Unambig); $i++) {
    if ($i > $max_seqs-1) { last; }
    my $id = $i+1;
    print $ofh ">" . $header . "_unambig_" . $id . "\n" . $Unambig[$i] . "\n";
}
close($ofh);


exit 0;

### Subroutines
sub check_nt_fasta
{
    ## Input: a FASTA file
    ## Output: nothing, fail if this isn't a valid nucleotide FASTA file
    my $f = $_[0];
    my $fh;
    ## Dictionary of valid NT characters
    my %Valid = ('A'=>1, 'T'=>1, 'G'=>1, 'C'=>1, 
		 'W'=>1, 'S'=>1, 'M'=>1, 'K'=>1, 
		 'R'=>1, 'Y'=>1, 'B'=>1, 'D'=>1, 
		 'H'=>1, 'V'=>1, 'N'=>1);
    my $h;
    if ($fasta =~ m/\.gz$/) { ## if a gzip compressed infile
	open($fh, '-|', 'gunzip', '-c', $fasta) || die "\n\n Cannot open the input file: $fasta\n\n";
    }
    else {
	open($fh ,'<', $fasta) || die "\n Cannot open the file: $fasta\n";
    }
    while(<$fh>) {
	$_ =~ s/\r[\n]*/\n/gm;
	chomp;
	if ($_ =~ m/^>/) {
	    $h = $_;
	}
	else {
	    my $s = uc($_); ## make all characters uppercase
	    my @a = split(//, $s); ## Read each character in the string $s into the array @a
	    foreach my $i (@a) {
		unless (exists $Valid{$i}) { ## If this character isn't valid
		    unless ($i eq "") { ## If this character isn't just an empty string (usually a space)
			close($fh); die "\n Error: It looks like the input FASTA you gave is an amino acid FASTA. Input needs to be a nucleotide FASTA\n Offending sequence: $h\n Offending base: --> $i <--\n\n\n Its ASCII encoding is: " . ord($i) . "\n\n";
		    }
		}
	    }
	}
    }
    close($fh);
}
sub find_ambiguous
{
    ## Input: one nulceotide sequence that may contain ambiguous characters
    ## Output: No more than $max unambiguous sequences in an array
    my $s = uc($_[0]); ## upper case the sequence
    my $max = $_[1];
    $s =~ s/ //g; ## get rid of all spaces in the sequence
    my %Uniq;
    my @Seq = split(//, $s); ## Read each character in the string $s into the array @Seq
    my $combos = 1;
    foreach my $base (@Seq) {
	if (exists $AmbigBases{$base}) {
	    $combos *= scalar(@{$AmbigBases{$base}});
	}
    }
    if ($combos < $max) { $max = $combos; }
    my $sequences = 0;
    while ($sequences < $max) {
	my $temp_seq = "";
	foreach my $base (@Seq) {
	    if ($base =~ m/A|T|G|C/) { $temp_seq = $temp_seq . $base; } ## If the base isn't an ambiguous one.
	    else { $temp_seq = $temp_seq . random_base($base); } ## Else, grab an unambiguous one for the ambig one we have
	}
	$Uniq{$temp_seq} = 1;
	$sequences = keys %Uniq;
    }
    my @tmp;
    foreach my $unambig_seq (keys %Uniq) {
	push (@tmp, $unambig_seq);
    }
    return(@tmp);
}
sub random_base
{
    ## Input: an ambiguous base character
    ## Output: a random base associated with the ambiguous character
    my $b = $_[0];
    my $types = scalar(@{$AmbigBases{$b}}); ## count how many ambiguous bases there are for the input character
    my $random_base = ${$AmbigBases{$b}}[rand $types]; # pick a random base!
    return $random_base;
}
sub construct_hash
{
    ## Input: nothing
    ## Output: a hash of the ambiguous base characters
    my %hash = (
        W => [ 'A', 'T' ],
        S => [ 'C', 'G' ],
        M => [ 'A', 'C' ],
        K => [ 'G', 'T' ],
        R => [ 'A', 'G' ],
        Y => [ 'C', 'T' ],
        B => [ 'C', 'G', 'T' ],
        D => [ 'A', 'G', 'T' ],
        H => [ 'A', 'C', 'T' ],
        V => [ 'A', 'C', 'G' ],
        N => [ 'A', 'C', 'G', 'T' ]
        );
    return(%hash);
}
