#!/usr/bin/perl
 
#Program to convert nucleotide sequences and a protein alignment into a nucleotide alignment
#Usage: nucleotide_translation_alignment.pl <nucleotide fasta> <protein alignment fasta> [<translation table>]
#outputs to stdout
use strict;
use warnings;
use lib '/Users/jkwan/scripts/sequence-toolkit';
use Sequence_toolkit;
use Data::Dumper;

my $nucleotide_fasta = $ARGV[0];
my $protein_alignment = $ARGV[1];
my $translation_table = $ARGV[2];
unless ( $translation_table )
{
	$translation_table = 11; # Default is the bacterial code
}

#my %codon_translations =  ('TTT' => 'F', 'TTC' => 'F', 'TTA' => 'L', 'TTG' => 'L', 'CTT' => 'L', 'CTC' => 'L', 'CTA' => 'L', 'CTG' => 'L', 'ATT' => 'I', 'ATC' => 'I', 'ATA' => 'I', 'ATG' => 'M', 'GTT' => 'V', 'GTC' => 'V', 'GTA' => 'V', 'GTG' => 'V', 'TCT' => 'S', 'TCC' => 'S', 'TCA' => 'S', 'TCG' => 'S', 'CCT' => 'P', 'CCC' => 'P', 'CCA' => 'P', 'CCG' => 'P', 'ACT' => 'T', 'ACC' => 'T', 'ACA' => 'T', 'ACG' => 'T', 'GCT' => 'A', 'GCC' => 'A', 'GCA' => 'A', 'GCG' => 'A', 'TAT' => 'Y', 'TAC' => 'Y', 'CAT' => 'H', 'CAC' => 'H', 'CAA' => 'Q', 'CAG' => 'Q', 'AAT' => 'N', 'AAC' => 'N', 'AAA' => 'K', 'AAG' => 'K', 'GAT' => 'D', 'GAC' => 'D', 'GAA' => 'E', 'GAG' => 'E', 'TGT' => 'C', 'TGC' => 'C', 'TGG' => 'W', 'CGT' => 'R', 'CGC' => 'R', 'CGA' => 'R', 'CGG' => 'R', 'AGT' => 'S', 'AGC' => 'S', 'AGA' => 'R', 'AGG' => 'R', 'GGT' => 'G', 'GGC' => 'G', 'GGA' => 'G', 'GGG' => 'G',);

my $nucl_seqHashRef = Sequence_toolkit::fastaToHash( $nucleotide_fasta );
my $prot_alignHashRef = Sequence_toolkit::fastaToHash( $protein_alignment );

my $align_length;
foreach my $prot (keys %$prot_alignHashRef)
{
	unless (defined $align_length)
	{
		$align_length = length ( $prot_alignHashRef->{$prot} );
	}
	
	if ( length( $prot_alignHashRef->{$prot} ) != $align_length )
	{
		print "Error! Proteins not aligned\n";
		exit 1;
	}
	
	my $prot_seq = $prot_alignHashRef->{$prot};
	my @prot_seq_charArray = split (//, $prot_seq);
	my $nucl_seq = $nucl_seqHashRef->{$prot};
	my @nucl_seq_charArray = split (//, $nucl_seq);
	# Print seq header
	print ">$prot\n";
	
	foreach my $char (@prot_seq_charArray)
	{
		if ($char eq "-")
		{
			print "---";
		} else {
			my @codonArray = splice( @nucl_seq_charArray, 0, 3 );
			my $codon = join( "", @codonArray);
			my $codon_translation = Sequence_toolkit::getTranslationForCodon( $codon, $translation_table );
			unless ( $codon_translation )
			{
				$codon_translation = 'X';
			}

			if ($codon_translation ne uc($char))
			{
				print "Error! Translation of $codon doesn't equal $char\n";
				exit 1;
			} else { 
				print "$codon";
			}
		}
	}
	print "\n";
}