#!/usr/bin/perl -w

use strict;
use Getopt::Long;
use lib "../Perl/modules";
use Genome;
use Gene;
use GenomicsUtils;
use lib "../stat";
use Vector;
use lib "../seq";
use CodonBias;

{
    my $taxon;
    my $genesFile;
    my $col1="Gene1";
    my $col2="Gene2";
    my $keep=1; # keep the additional columns?
    my $host;
    my $debug=0;

    (GetOptions('in=s' => \$genesFile,
		'taxon=i' => \$taxon,
		'col1=s'=>\$col1,
		'col2=s'=>\$col2,
		'debug=i'=>\$debug,
		'host=s'=>\$host,
		'keep=i'=>\$keep)
     && @ARGV==0 && $genesFile && $taxon)
	|| die "Usage: CodonBiasDotProduct.pl -taxon id -in genePairs [-col1 Gene1 -col2 Gene2 -keep 0]"
	. "   [-debug 10 -host name] > out\n";

    my $pairs = Vector::readDelimList($genesFile);
    die "No input" if scalar @$pairs < 1;
    die "Missing column $col1" if !exists $pairs->[0]{$col1};
    die "Missing column $col2" if !exists $pairs->[0]{$col2};
    GenomicsUtils::setDebug($debug);
    GenomicsUtils::connect($host);

    my %codonTable = ( Gene::codonUsageTable() );
    my %aaCodons = ( CodonBias::aaToCodons(%codonTable) );
    my @list = Gene::codonUsageTable();
    my %codonTotal = map { $_ => 0 } keys %codonTable;
    my %codonGene = (); # locusId -> codon -> count
    my %hasCodons = (); # try if gene is not a frameshift
    my $genome = Genome::new(taxonomyId=>$taxon);
    my @genes = $genome->genes();

    foreach my $gene (@genes) {
	my %codonThis = map { $_ => 0 } keys %codonTable;
	$hasCodons{$gene->locusId()} = $gene->addCodonUsage(\%codonThis);
	$codonGene{$gene->locusId()} = \%codonThis;
	map { $codonTotal{$_} += $codonThis{$_} } keys %codonThis;
    }
    my %aaTotal = ();
    CodonBias::sumAA(\%codonTotal, \%aaTotal, \%aaCodons);
    if ($keep) {
	print join("\t" , keys %{ $pairs->[0] },"BiasSim","Bias1","Bias2"),"\n";
    } else {
	print join("\t", "Gene1", "Gene2", "BiasSim","Bias1","Bias2"), "\n";
    }
    foreach my $row (@$pairs) {
	my $locus1 = $row->{$col1};
	my $locus2 = $row->{$col1};
	die "Cannot find $col1=$locus1 for taxon $taxon" if !exists $codonGene{$locus1};
	die "Cannot find $col2=$locus2 for taxon $taxon" if !exists $codonGene{$locus2};
	my $sumBias = 0;
	my $sum1 = 0;
	my $sum2 = 0;
	if ($hasCodons{$locus1} && $hasCodons{$locus2}) {
	    my $c1 = $codonGene{$locus1};
	    my $c2 = $codonGene{$locus2};
	    my %aa1 = (); sumAA($c1, \%aa1, \%aaCodons);
	    my %aa2 = (); sumAA($c2, \%aa2, \%aaCodons);
	    foreach my $aa (keys %aaCodons) {
		next if $aa eq "*";
		foreach my $codon (@{$aaCodons{$aa}}) {
		    my $fhat = $codonTotal{$codon}/$aaTotal{$aa};
		    my $f1 = ($c1->{$codon} + $fhat)/($aa1{$aa} + 1); #pseudocount 1
		    my $f2 = ($c2->{$codon} + $fhat)/($aa2{$aa} + 1); #pseudocount 1
		    $sumBias += ($f1-$fhat)*($f2-$fhat);
		    $sum1 += abs($f1-$fhat);
		    $sum2 += abs($f2-$fhat);
		    #print STDERR "aa $aa codon $codon c1 $c1->{$codon} c2 $c2->{$codon} aa1 $aa1{$aa} aa2 $aa2{$aa} fhat $fhat $f1 f1 $f2 f2 biasTerm " . ($f1-$fhat)*($f2-$fhat) . "\n" if $locus1==14146;
		}
	    }
	}
	if ($keep) {
	    print join("\t", values %$row, $sumBias,$sum1,$sum2), "\n";
	} else {
	    print join("\t", $locus1, $locus2, $sumBias,$sum1,$sum2), "\n";
	}
    }
}
