#!/usr/bin/perl

use Bio::Seq;
use Bio::SeqIO;
use Bio::SeqFeature::Generic;
use Bio::SeqIO::genbank;
use DBI;
use DBD::mysql;

use strict;
use warnings;

our $DB_NAME;
our $DB_HOST;
our $DB_USER;
our $DB_PASS;
require 'lib/config.pl';

my $usage = "getGBinfo.pl <infile> <nomeMtg> <outdir>";
my $infile = shift or die $usage;
my $nomeMtg = shift or die $usage;
my $outDir = shift or die $usage;

my $dbh = DBI->connect("DBI:mysql:$DB_NAME:$DB_HOST",$DB_USER,$DB_PASS,{RaiseError => 1});

my $seq_in = Bio::SeqIO->new('-file' => "<$infile",
			     '-format' => 'GenBank');


my %contig;
($contig{'id'}) = ($infile =~ /contig(\d+)/i); ## ignore case
if ($contig{'id'} !~ /\d+/){
	print "Impossible to determine contig number in file $infile. Skipping this file.";
	exit (1);
}

$contig{'idMontagem'} = $nomeMtg;

# print "Adding Contig #".$contig{'id'}."<BR>";

my %validCodes;
my $seq = $seq_in->next_seq();

my $contigSeq = $seq->seq();

our $show=0;
%validCodes = getValidCodes($dbh);

my $featGENE = 0;
my $featCDS = 0;

my $featRBS;
my %rbs;

my %trna;
my $trnaNum = 1;

my $counter = 1;

for my $feat_object ($seq->get_SeqFeatures) {
	$show=0;
	my %infoORFs;
	my %annot;


# 	if (($feat_object->primary_tag =~ /^gene$/) && !$featCDS){
# 		$featGENE = 1;
# 	}elsif (($feat_object->primary_tag =~ /^CDS$/) && !$featGENE){
# 		$featCDS = 1;
# 	}

	if ($feat_object->primary_tag =~ /^RBS$/){
		$featRBS = 1;
		my ($id) = $feat_object->get_tag_values('gene');

		my $strand = $feat_object->strand;

		if ($strand == 1){
			$rbs{'position'} = $feat_object->start;
		}else{
			$rbs{'position'} = $feat_object->end;
		}
		$rbs{'old_position'} = $rbs{'position'};
		$rbs{'pattern'} = substr($contigSeq, $feat_object->start - 1, 5);

	}elsif (
# 	    (($feat_object->primary_tag =~ /^gene$/) && !$featCDS )
# 	    ||
	    (($feat_object->primary_tag =~ /^CDS$/)  && !$featGENE)
	   ){

		$infoORFs{'idMontagem'} = $contig{'idMontagem'};

		my $strand = $feat_object->strand;

		if ($strand == 1){
			$infoORFs{'start'} = $feat_object->start;
			$infoORFs{'end'} = $feat_object->end;
		}else{
			$infoORFs{'start'} = $feat_object->end;
			$infoORFs{'end'} = $feat_object->start;
		}

		$infoORFs{'length'} = $feat_object->end - $feat_object->start  + 1;

		my $id;
		if ($feat_object->has_tag('locus_tag')){
			($id) = $feat_object->get_tag_values('locus_tag');
			$annot{'synonym'} = $id;
		}elsif ($feat_object->has_tag('gene_id')){
			($id) = $feat_object->get_tag_values('gene_id');
		}else{
			$id = sprintf("%04d", $counter);
		}

		$infoORFs{'contigId'} = $contig{'id'};

		my $orfId;
		if ($id =~ /\d{3,4}\.\d{4}$/){
			my ($ctgId) = ($id =~ /(\d{3,4})\./);
			$ctgId = sprintf("%04d", $ctgId);

			($orfId) = ($id =~ /\.(\d{4}.*)$/);
			$orfId = "$ctgId.$orfId";
		}elsif($id =~ /\d+$/){
			($id) = ($id =~ /.*(\d{4})$/);
			my $ctgId = sprintf("%04d", $contig{'id'});
			$orfId = "$ctgId.$id";
		}

		if ($feat_object->has_tag('gene')){
			($annot{'orfName'}) = $feat_object->get_tag_values('gene');
		}else{
			$annot{'orfName'} = $orfId;
		}
print "$id - $orfId<BR>";
		$infoORFs{'id'} = $orfId;

		($infoORFs{'frame'}) = getFrame($infoORFs{'start'}, $strand);

		$annot{'idMontagem'} = $infoORFs{'idMontagem'};
		$annot{'id'} = $infoORFs{'id'};

		my $valid;
		my $cogId;
		my $fs;
		my $annotator;
		my $reviewer;
		my $note;
		if ($feat_object->has_tag('note')){
			($note) = $feat_object->get_tag_values('note');
			$annot{'note'} = $note;
			
			($valid) = ($note =~ /VALIDATION:(\w+\s?\w+)\|/);
			if ($valid){
				$valid = lc $valid;
				$annot{'validation'} = $validCodes{$valid};
			}
			
			($cogId) = ($note =~ /COG-NUMBER:(COG\d{4})\|/);
			if ($cogId){
				$annot{'COGid'} = $cogId;
			}
	
			($fs) = ($note =~ /FRAMESHIFT:(\w{2,3})\|/);
			$annot{'frameshift'} = '';
			if ($fs && $fs =~ /yes/){
				$annot{'frameshift'} = 1;
			}
			
			($annotator) = ($note =~ /ANNOTATOR:(\w+)\|/);
			if ($annotator){
				$annot{'anotador'} = $annotator;
				($reviewer) = ($note =~ /REVIEWER:(\w+)\|/);
				if ($reviewer){
					$annot{'revisor'} = $reviewer;
				}
			}
		}

		if ($feat_object->has_tag('product')){
			($annot{'product'}) = $feat_object->get_tag_values('product');
		}
		if ($feat_object->has_tag('function')){
			($annot{'function'}) = $feat_object->get_tag_values('function');
		}

		my $orfSeq = substr($contigSeq, $feat_object->start - 1, $infoORFs{'length'});
		($infoORFs{'bases_a'},$infoORFs{'bases_c'},$infoORFs{'bases_t'},$infoORFs{'bases_g'}) = calculaNumBases($orfSeq);

		insert($dbh, 'orf', %infoORFs);
		insert($dbh, 'anotacao', %annot);

		if ($featRBS){
			$rbs{'idMontagem'}	= $contig{'idMontagem'};
			$rbs{'id'}		= $annot{'id'};
			$rbs{'start'}		= $infoORFs{'start'};
			$rbs{'stop'}		= $infoORFs{'end'};
			$rbs{'codon'}		= substr($orfSeq, 0, 3);
			$rbs{'old_codon'}	= $rbs{'codon'};
			$rbs{'shift'}		= 0;

			insert($dbh, 'rbs', %rbs);

			$featRBS = 0;
		}else{
			$rbs{'idMontagem'}	= $contig{'idMontagem'};
			$rbs{'id'}		= $annot{'id'};
			$rbs{'start'}		= $infoORFs{'start'};
			$rbs{'stop'}		= $infoORFs{'end'};
			$rbs{'codon'}		= substr($orfSeq, 0, 3);
			$rbs{'old_codon'}	= $rbs{'codon'};
			$rbs{'shift'}		= 0;

			$rbs{'position'} = 0;
			$rbs{'old_position'} = $rbs{'position'};
			$rbs{'pattern'} = "---";

			insert($dbh, 'rbs', %rbs);
		}

		my $codon = substr($orfSeq, 0, 3);
	
		my $q = "INSERT INTO startCodons VALUES(
				\'".$contig{'idMontagem'}."\',
				\'".$annot{'id'}."\',
				0,
				\'$codon\',
				".$infoORFs{'start'}.",
				".$infoORFs{'start'}.",
				0)
		";
		$dbh->do($q);
		
		if ($strand != 1){
			my $tmpSeq = Bio::PrimarySeq->new ('-seq' => $orfSeq);
			$orfSeq = $tmpSeq->revcom->seq;
			
		}

		writeSeq($outDir, $infoORFs{'id'}, $orfSeq);
		
		$counter++;

	}elsif ($feat_object->primary_tag =~ /^tRNA$/i){

		my ($prd) = $feat_object->get_tag_values('product');

		$trna{'idMontagem'}	= $contig{'idMontagem'};
		$trna{'contigId'}		= $contig{'id'};
		$trna{'tRNAnum'}	= $trnaNum;
		($trna{'type'}) = ($prd =~ /tRNA-(\w+)/i);

		my $strand = $feat_object->strand;

		if ($strand == 1){
			$trna{'begin'} = $feat_object->start;
			$trna{'end'} = $feat_object->end;
		}else{
			$trna{'begin'} = $feat_object->end;
			$trna{'end'} = $feat_object->start;
		}

		$trnaNum++;
		
		insert($dbh, 'tRNA', %trna);

	}elsif ($feat_object->primary_tag =~ /^source$/){
		$contig{'tam'} = $feat_object->length;
	}
}
insert($dbh,  'contig', %contig);


###### BLASTs
# opendir(DIR, $outDir."/p");
# my @orfFiles = readdir(DIR);
# closedir(DIR);
# 
# my $blastOutDir = "BlastReports/$nomeMtg";
# mkdir $blastOutDir;
# 
# foreach my $orfFile (@orfFiles) {
# 	next if ($orfFile eq "." or $orfFile eq "..");
# 	system("cd bin && perl runOrfsBlast.pl $DB_NAME $nomeMtg ../$outDir/p/$orfFile ../BlastDatabases/nr nr ../$blastOutDir");
# }

######################
#### SUB ROUTINES ####
######################

sub getValidCodes{
	my ($dbh) = @_;

	my $q = "SELECT * FROM validation";
	my $res = $dbh->prepare($q);
	$res->execute();
	my %validation;

	while (my $data = $res->fetchrow_hashref()) {
		my %data = %{$data};
		my $status = lc $data{'status'};
		$validation{$status} = $data{'code'};
	}

	return %validation;
}

sub getFrame{
	my ($start, $strand) = @_;

	my $frame;
	if ($strand == 1){
		$frame = $start % 3;
		if ($frame == 0){
			$frame = 3;
		}
		$frame = "F".$frame;
	}else{
		$frame = $start % 3;
		if ($frame == 2){
			$frame = 3;
		}elsif ($frame == 0){
			$frame = 2;
		}
		$frame = "R$frame";
	}

	return $frame;
}

sub insert{
	my ($dbh, $tableName, %data) = @_;

	my @table;my @value;my $q;my $tbl;my $vle;

	for my $k (keys %data){
		push @table, $k;
		push @value, $data{$k};
	}

	$q = "INSERT INTO $tableName ";
	$tbl = "(";
	$vle = " VALUES (";

	for (my $i=0; $i<@table; $i++){
		$tbl .= $table[$i].", ";
		$vle .= "\"".$value[$i]."\", ";
	}

	$tbl =~ s/,\s$/)/;
	$vle =~ s/,\s$/)/;

	$q .= $tbl.$vle;

	$dbh->do($q);
}

sub calculaNumBases{
	my ($seq) = @_;

	my %hash = ('A' => 0, 'C' => 1 , 'T' => 2, 'G' => 3);
	my @array = (0,0,0,0);
	my @nt = split(//, $seq);

	foreach my $base (@nt){
		if ($base =~ /A/i || $base =~ /C/i || $base =~ /G/i || $base =~ /T/i){
			$array[$hash{$base}]++;
		}
	}

	return @array;
}

sub writeSeq{
	my ($orfDir, $id, $seq) = @_;

	my $aaDir = $orfDir."/n";
	my $ptDir = $orfDir."/p";

	my $aaSeq = Bio::PrimarySeq->new ('-seq' => $seq);
	my $ptSeq = $aaSeq->translate();

	my $aaFile = $aaDir.'/ORF'.$id.'.fas';
	my $ptFile = $ptDir.'/ORF'.$id.'.fas';

	my $fna = Bio::SeqIO->new('-file' =>">$aaFile", '-format' => 'fasta');
	my $fpt = Bio::SeqIO->new('-file' =>">$ptFile", '-format' => 'fasta');

	$aaSeq->display_id("$id.n");
	$fna->write_seq($aaSeq);
	
	$ptSeq->display_id("$id.p");
	$fpt->write_seq($ptSeq);

}