#!/usr/bin/perl -w

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

my $helpStr=qq{

    codonBias.pl <params> <options>

	<params>:
	-ref: 
	     cog
	     top100
	
	-taxId: taxonomyId, then the script will get all genomic scaffolds 
	
	<options>
	-host: DB host (default is local host) 	
	           
    Note:


};

unless (@ARGV >= 4){die "$helpStr";}

my ($taxId, $refOption, $host, $printFrameShift);

(GetOptions( 
	    'taxId=i' => \$taxId,
	    'ref=s' => \$refOption,
	    'host=s' => \$host,
	    'printFrameShift=i' => \$printFrameShift
) && scalar @ARGV == 0) || die $helpStr;


connect($host);
my $cogFile = "/usr2/people/khuang/operon/getOperons/genomeCluster/lowCBIcogs"; #should store this in the database
my %codonTable = Gene::codonUsageTable();
my %aaCodons = CodonBias::aaCodons(\%codonTable);

my $checkTax = GenomicsUtils::query(qq{SELECT * FROM Taxonomy WHERE taxonomyId=$taxId});
die " \n** Taxonomy ID: $taxId is an invalid genome **\n " unless defined($checkTax->[0][0]);

my $genome = Genome::new(taxonomyId => $taxId); # genome object
my @genes = $genome->genes(); # get an array of gene objects in the genome 
my ( %geneCodonTables, %hasData );

my %locusToCOG = map { $_->[0] => { "COG" => $_->[1], "code" => $_->[2], "desc" => $_->[3] } }
@{ GenomicsUtils::query("SELECT o.locusId, COG.cogInfoId, COGInfo.funCode, COGInfo.description"
			. " FROM Locus o, COG, COGInfo, Position p, Scaffold s"
			. " WHERE s.taxonomyId=$taxId AND s.isGenomic AND s.isActive"
			. "  AND p.scaffoldId=s.scaffoldId AND o.posId=p.posId AND o.priority=1"
			. "  AND o.locusId=COG.locusId AND COG.cogInfoId=COGInfo.cogInfoId") };
my %cogToGenes = ();
foreach my $gene (@genes) {
    my $locusId = $gene->locusId();
    addToHashList(\%cogToGenes, $locusToCOG{$locusId}->{"COG"}, $gene) if exists $locusToCOG{$locusId};
}

# compute codon tables
foreach my $gene (@genes) {
    my $locusId= $gene->locusId();
    my %codon = map { $_ => 0 } keys %codonTable;
    $hasData{$locusId} = $gene->addCodonUsage(\%codon); # adds codons
    $geneCodonTables{$locusId} = \%codon;
}


# compute CBI
my %cbiORF = ();
foreach my $locusId (keys %geneCodonTables) {
    $cbiORF{$locusId} = CodonBias::CBI($geneCodonTables{$locusId}, \%aaCodons) if $hasData{$locusId};
}

# Compute the reference codon table and best codons list
my @refORFs = ();
my %refCodonTable = map { $_ => 0 } keys %codonTable;

if ($refOption eq 'cog'){
    my $nCOG = 0;
    open COG, $cogFile || die "$cogFile:$!\n";
    while(<COG>) {
	s/\n//;
	if (exists $cogToGenes{$_}) {
	    foreach my $gene (@{ $cogToGenes{$_} }) {
		# at least 300 aa, not a frameshift
		push @refORFs, $gene->locusId() if $hasData{$gene->locusId()} && $gene->length() >= 300;
	    }
	}
	$nCOG++;
	last if $nCOG >= 500;
    }
    close COG || die "$cogFile:$!\n";

    # now filter @refORFs to be size 100, only the most biased ones
    @refORFs = sort {$cbiORF{$a} <=> $cbiORF{$b}} @refORFs;
    $#refORFs = 99 if scalar(@refORFs) > 100;
} elsif ($refOption eq 'top100'){
    my $count = 0;
    foreach my $locusId (sort {$cbiORF{$a} <=> $cbiORF{$b}} keys %cbiORF) {
	push @refORFs, $locusId;
	last if @refORFs >= 100;
    }
}

print STDERR "Reference Locuses from $refOption (" . scalar @refORFs . ") -- " . join(",", @refORFs) . "\n";
foreach my $locusId (@refORFs) {
    map { $refCodonTable{$_} += $geneCodonTables{$locusId}{$_} } keys %codonTable;
}

my %bestCodonCount = CodonBias::BestCodonCount(\%refCodonTable, \%aaCodons);
my %sumCodonCount = ();
CodonBias::sumAA(\%refCodonTable, \%sumCodonCount, \%aaCodons);

print "\#".(join "\t", "gi", "CAI", "CAI_Alm", "CBI", "COG", "COGanno", "COGfun", "Description"."\n");

my $anno = GenomicsUtils::queryHashList("SELECT locusId,description FROM Description WHERE locusId IN ("
					. join(",", keys %hasData) . ")");

foreach my $gene (@genes) {
    my $locusId=$gene->locusId();
    my ($caiLiSharp, $caiAlm);
    if ($hasData{$locusId}) {
	$caiLiSharp = CodonBias::CAI($geneCodonTables{$locusId}, \%aaCodons, \%refCodonTable, \%bestCodonCount);
	$caiAlm = CodonBias::CAI($geneCodonTables{$locusId}, \%aaCodons, \%refCodonTable, \%sumCodonCount);
    } else {
	next unless $printFrameShift;
	$caiLiSharp = "";
	$caiAlm = "";
	$cbiORF{$locusId} = "";
    }
    my ($cogId, $cogDesc, $cogFun);
    if (exists $locusToCOG{$locusId}) {
	$cogId = $locusToCOG{$locusId}{"COG"};
	$cogDesc = $locusToCOG{$locusId}{"desc"};
	$cogFun = $locusToCOG{$locusId}{"code"};
    } else {
	$cogId = "-";
	$cogDesc = "-";
	$cogFun = "-";
    }
    print join("\t", $locusId, $caiLiSharp, $caiAlm, $cbiORF{$locusId},
	       $cogId, $cogDesc, $cogFun, (exists $anno->{$locusId} ? $anno->{$locusId}[0] : "")) , "\n";
}

#sub chisq(\%\%\%$){
#
#    my ($codonR, $aaCR, $binR, $tCodons) = @_;
#    my $chi = 0;
#   
#    foreach my $aa (keys %aaCodons){
#	foreach my $sCodon (@{$aaCodons{$aa}}){
#	    if ( $aaCR->{$aa} ){
#		$chi += ( ( $codonR->{$sCodon} - ($aaCR->{$aa}/$binR->{$aa}) )**2 ) / ( $aaCR->{$aa}/$binR->{$aa} ) ;
#	    }
#	}
#    }
#    return $chi/($tCodons - 1);
#
#}
