#!/usr/bin/perl -w
# Script for calculate ka/ks

use strict;

use Bio::Tools::Run::Alignment::Clustalw;
use Bio::Tools::Run::Phylo::PAML::Yn00;
use Bio::Tools::Run::Phylo::PAML::Codeml;
use Bio::Align::Utilities qw(aa_to_dna_aln);
use Bio::SimpleAlign;
use Bio::AlignIO;


# 1.  Load the DNA sequences into 2 variables:
my %dna_seq0;
my $seq1=Bio::Seq->new(-seq =>"TCTCTCTGGCCCAAAATCCGGGTTCCATTAAAAGTTGTGAGGACTGCTGAAAACAAGTTAAGTAACCGTTTCTTCCCTTATGATGAAATCGAGACAGAAGCTGTTCTGGCCATTGATGATGATATCATTATGCTGACCTCTGACGAGCTGCAATTTGGTTATGAG", -id =>"seq1"); # Homo sapiens
$dna_seq0{$seq1->id} = $seq1;
my $seq2=Bio::Seq->new(-seq =>"TCACTGTGGCCCAAAGTCGCAGTGCCTCTTAAAGTGGTCCGCACCAAAGAAAACAAGCTCAGCAATCGATTCTTTCCGTTTGATGAGATCGAGACAGAAGCTGTCCTGGCCATTGACGATGACATCATCATGTTAACCTCAGATGAGCTACAGTTTGGATATGAG", -id =>"seq2"); # Bos taurus
$dna_seq0{$seq2->id} = $seq2; 

#2.  Load the corresponding amino acid sequences:
my @aa_seq;
$aa_seq[0] = $seq1->translate();
$aa_seq[1] = $seq2->translate();

#3.  Load the align amino acid sequences:
my $aln_factory = Bio::Tools::Run::Alignment::Clustalw->new('quiet'=>1);
my $aa_aln = $aln_factory->align(\@aa_seq);


#4.  All the codons in the DNA sequence to correspond to the AA alignment

my $dna_aln = aa_to_dna_aln($aa_aln, \%dna_seq0);

# =====================
# Yn00
# =====================

  my $aln = $dna_aln;

  my $yn = new Bio::Tools::Run::Phylo::PAML::Yn00();
  $yn->alignment($aln);
  my ($rc,$parser) = $yn->run();
  while( my $result = $parser->next_result ) {
    my @otus = $result->get_seqs();
    my $MLmatrix = $result->get_MLmatrix();
    # 0 and 1 correspond to the 1st and 2nd entry in the @otus array
    my $dN = $MLmatrix->[0]->[1]->{dN};
    my $dS = $MLmatrix->[0]->[1]->{dS};
    my $kaks =$MLmatrix->[0]->[1]->{omega};
    print "Yn00:\tKa = $dN\tKs = $dS\tKa/Ks = $kaks\n";
  }

# =====================
# PAML
# =====================

   my $kaks_factory = Bio::Tools::Run::Phylo::PAML::Codeml->new
                   ( -params => { 'runmode' => -2,
                                  'seqtype' => 1,
                                },
		     -save_tempfiles => 0,  # remember to cleanup the temp files/folders
		     -verbose => 0 );
 
	# set the alignment object
   $kaks_factory->alignment($dna_aln);
 
	# run the KaKs analysis
   my ($rc2,$parser2) = $kaks_factory->run();
   unless(defined $parser2) {
	print "Undef parser!\n"; return (undef,undef,undef,undef,undef);
   }
   my $result = $parser2->next_result;
   my $MLmatrix = $result->get_MLmatrix();
   my @otus = $result->get_seqs();

   my $dN = $MLmatrix->[0]->[1]->{dN};
   my $dS = $MLmatrix->[0]->[1]->{dS};
   my $kaks =$MLmatrix->[0]->[1]->{omega};
   print "Codeml:\tKa = $dN\tKs = $dS\tKa/Ks = $kaks\n";