#!/usr/bin/perl

#a script to take a fasta alignment and sumarize sequence divergence between a specified pair of sequences (termed the reference and the outgroup)

use strict;
use warnings;

my $usage = "\nUSAGE: $0 fastaFile referenceName outgroupName\n\n";

my $file = shift or die ($usage);
my $refName = shift or die ($usage);
my $outgroupName = shift or die ($usage);

my %fastaHash = arrays2hash(get_fasta_names_and_seqs($file));

my $length;
my $gene;
my $refPresent;
my $outgroupPresent;
foreach my $header (keys %fastaHash){
	if ($length){
		$length == length ($fastaHash{$header}) or die ("\nERROR: Sequences not of equal length in $file\n\n");
	}
	$length = length ($fastaHash{$header});
	my @splitHeader = split (/\_/, $header);
	if ($gene){
		$gene eq $splitHeader[1] or die ("\nERROR: gene names do not match in $file\n\n");
	}
	$gene = $splitHeader[1];
	if ($splitHeader[0] eq $refName){
		$refPresent = 1;
	}elsif ($splitHeader[0] eq $outgroupName){
		$outgroupPresent = 1;
	}
}

print "Type\tRef\tOutgroup\tGene\tPosition\tCodonPosition\tRefNuc\tOutgroupNuc\tRefCodon\tOutgroupCodon\tOutgroupCodon_oneChange\tRefAA\tOutgroupAA\tOutgroupAA_oneChange\tSynNon-Syn\n";

$length % 3 == 0 or die ("\nERROR: Sequence length not divisible by 3 in $file\n\n");
$refPresent or die ("\nERROR: Could not find the reference $refName in $file\n\n");
$outgroupPresent or die ("\nERROR: Could not find the outgroup $outgroupName in $file\n\n");

my $ungappedPos = 0;
my @gaps;
for (my $i = 0; $i < length ($fastaHash{"$refName\_$gene"}); $i += 3){
	my $refCodon = substr ($fastaHash{"$refName\_$gene"}, $i, 3);
	my $outgroupCodon = substr ($fastaHash{"$outgroupName\_$gene"}, $i, 3);
	my $refAA = codon2aa($refCodon);
	my $outgroupAA = codon2aa($outgroupCodon);
	$refAA eq '?' and die ("\nERROR: improper codon in $refName\_$gene after position $i: $refCodon\n\n");
	$outgroupAA eq '?' and die ("\nERROR: improper codon in $outgroupName\_$gene after position $i: $outgroupCodon\n\n");
	$refAA eq '-' and next;
	if ($outgroupAA eq '-'){
		push (@gaps, $ungappedPos+1);
		push (@gaps, $ungappedPos+2);
		push (@gaps, $ungappedPos+3);
		$ungappedPos += 3;
		next;
	}
	
	foreach my $j (0..2){
		substr($refCodon,$j,1) eq substr($outgroupCodon,$j,1) and next;
		my $outgroupCodon_1change = $refCodon;
		substr($outgroupCodon_1change,$j,1,substr($outgroupCodon,$j,1));
		my $outgroupAA_1change = codon2aa($outgroupCodon_1change);
		my $syn = "N";
		$refAA eq $outgroupAA_1change and $syn = "S";
		print "SNP\t$refName\t$outgroupName\t$gene\t", $ungappedPos+$j+1,"\t", $j+1, "\t", substr($refCodon,$j,1), "\t", substr($outgroupCodon,$j,1), "\t$refCodon\t$outgroupCodon\t$outgroupCodon_1change\t$refAA\t$outgroupAA\t$outgroupAA_1change\t$syn\n";
	}
	
	$ungappedPos += 3;
	
}

foreach my $pos (@gaps){
	print "Gap\t$refName\t$outgroupName\t$gene\t$pos\n";
}

exit;



sub arrays2hash {
	use strict;
	use warnings;

	(my $keyarray, my $valuearray) = @_;
	if (scalar(@$keyarray) != scalar(@$valuearray)) {
		die "Arrays differ in size: Mismatched number of keys and values"; 
	}
	
	my %newhash = ( );
	
	@newhash{ @$keyarray } = @$valuearray;


	return (%newhash);

}


sub get_fasta_names_and_seqs {
	use strict;
	use warnings;

	my ($inputfilename) = @_;
	my @fasta_names = ();
	my @fasta_seqs= ();

		   
	unless ( open(FILEDATA, $inputfilename) ) {
		print STDERR "Cannot open file \"$inputfilename\"\n\n"; #print error message
		exit; #exit the program
	}	

	my @filedata = <FILEDATA>; #Read the lines of the file into an array
	close FILEDATA;
	
	my $seq_count = 0; #this will be used to keep track of the number of sequences
	foreach my $line (@filedata){
		chomp $line;
		if ($line =~ /^\s*$/) {next;} #ignore line if it is blank
		
		elsif ($line =~ /^>/) { #if the line is a header line (begins with ">")...
			if ($line =~ /^>.*[\w]+/){
				my $partialLine = substr ($&, 1);
				push (@fasta_names, $partialLine); #add that line to an array of fasta names
				push (@fasta_seqs, ""); #and add a new blank element to an array of sequences
				++$seq_count; #also increment our counter which keeps track of sequence number
			}
		}	
		
		else { #if the line's not blank or a header, add it to the current sequence 
			$fasta_seqs[$seq_count-1] .= $line;
		}
		
		$fasta_seqs[$seq_count-1] =~s/\s//g; #remove all whitespace from the current  sequence
	}
	
	return (\@fasta_names, \@fasta_seqs);

}


sub codon2aa {
	use strict;
	use warnings;

    my($codon) = @_;

    $codon = uc $codon;
 
    my(%genetic_code) = (
    
    'TCA' => 'S',    # Serine
    'TCC' => 'S',    # Serine
    'TCG' => 'S',    # Serine
    'TCT' => 'S',    # Serine
    'TTC' => 'F',    # Phenylalanine
    'TTT' => 'F',    # Phenylalanine
    'TTA' => 'L',    # Leucine
    'TTG' => 'L',    # Leucine
    'TAC' => 'Y',    # Tyrosine
    'TAT' => 'Y',    # Tyrosine
    'TAA' => '*',    # Stop
    'TAG' => '*',    # Stop
    'TGC' => 'C',    # Cysteine
    'TGT' => 'C',    # Cysteine
    'TGA' => '*',    # Stop
    'TGG' => 'W',    # Tryptophan
    'CTA' => 'L',    # Leucine
    'CTC' => 'L',    # Leucine
    'CTG' => 'L',    # Leucine
    'CTT' => 'L',    # Leucine
    'CCA' => 'P',    # Proline
    'CCC' => 'P',    # Proline
    'CCG' => 'P',    # Proline
    'CCT' => 'P',    # Proline
    'CAC' => 'H',    # Histidine
    'CAT' => 'H',    # Histidine
    'CAA' => 'Q',    # Glutamine
    'CAG' => 'Q',    # Glutamine
    'CGA' => 'R',    # Arginine
    'CGC' => 'R',    # Arginine
    'CGG' => 'R',    # Arginine
    'CGT' => 'R',    # Arginine
    'ATA' => 'I',    # Isoleucine
    'ATC' => 'I',    # Isoleucine
    'ATT' => 'I',    # Isoleucine
    'ATG' => 'M',    # Methionine
    'ACA' => 'T',    # Threonine
    'ACC' => 'T',    # Threonine
    'ACG' => 'T',    # Threonine
    'ACT' => 'T',    # Threonine
    'AAC' => 'N',    # Asparagine
    'AAT' => 'N',    # Asparagine
    'AAA' => 'K',    # Lysine
    'AAG' => 'K',    # Lysine
    'AGC' => 'S',    # Serine
    'AGT' => 'S',    # Serine
    'AGA' => 'R',    # Arginine
    'AGG' => 'R',    # Arginine
    'GTA' => 'V',    # Valine
    'GTC' => 'V',    # Valine
    'GTG' => 'V',    # Valine
    'GTT' => 'V',    # Valine
    'GCA' => 'A',    # Alanine
    'GCC' => 'A',    # Alanine
    'GCG' => 'A',    # Alanine
    'GCT' => 'A',    # Alanine
    'GAC' => 'D',    # Aspartic Acid
    'GAT' => 'D',    # Aspartic Acid
    'GAA' => 'E',    # Glutamic Acid
    'GAG' => 'E',    # Glutamic Acid
    'GGA' => 'G',    # Glycine
    'GGC' => 'G',    # Glycine
    'GGG' => 'G',    # Glycine
    'GGT' => 'G',    # Glycine
    '---' => '-',    # gap
    );

    if(exists $genetic_code{$codon}) {
        return $genetic_code{$codon};
    }else{
            return "?";
    }
}