#!/usr/bin/perl -w

##################################################################################################
# This script retrieves sequences and creates a formatted flat file on the fly
# Sequences are retrieved by seqname (locus id).
# Format and Output mode (STDOUT or FILE) must be specified.
# If the GI number is available from the database and the genbank format is selected, then      
#     the flat file is retreived from the remote NCBI genbank database.                              
# Note: This script follows steps and ideas as seen in the script 'updatedb.pl'
# -This script relies on use of unique IDs established in the database. If any unique ID that is critical
# to this program is altered in any way (which normally shouldn't occur), this script will not function properly. (You will need to update
# this program to fix it)
#
# UTR regions not supported yet
##################################################################################################
# Last update: 4/18/06
##################################################################################################

use strict;
use DBI;
use Bio::Seq;
use Bio::Seq::RichSeq;
use Bio::SeqIO;
use Bio::Species;
use Bio::DB::GenBank;
use Bio::SeqFeature::Generic;
use Bio::Location::Simple;

my $dir = "$ENV{HOME}/sequences";
my $usage= "perl getFormatSequences.pl format output seqname";
my $format = lc(shift) or die "USAGE: $usage $!\n";
# check for help
if( $format =~ /(-|-h|--help)/ ) {
	print "\n\t*** Usage: $usage ***\n\n";
	print "\tArguments:\n";
	print "\t-format\t\tSelect desired format ('fasta' or 'genbank')\n";
	print "\t-output\t\tChoose from 'stdout' or 'file'. Choosing 'file' will output the flat file to the directory '$dir'.\n\t\t\tPlease make sure that this directory exists\n";
	print "\t-seqname\tProvide the seqname/locus of the sequence you wish to obtain from the database\n\n";
	exit;
}
my $output = uc(shift) or die "USAGE: $usage $!\n"; # has to be either stdout or file
my $seqname = uc(shift) or die "USAGE: $usage $!\n";
# terminate if incorrect output type is given
die "Incorrect output type: '$output'...need 'stdout' or 'file'\nCheck usage: $usage" if ( ($output ne 'STDOUT') && ($output ne 'FILE') );

#my $seqname = &getSeqname; # get the seqname from standard input

# set up date
my @months = ('JAN', 'FEB', 'MAR', 'APR', 'MAY', 'JUN', 'JUL', 'AUG', 'SEP', 'OCT', 'NOV', 'DEC');
my ($second, $minute, $hour, $day, $month, $year, $weekday, $dayofyear, $daylightSavings) = localtime(time);
$year = $year - 100; # localtime returns $year as number of years since 1900 (ie: 2006 is 106, subtract 100 to get 06)
$month = $months[$month]; # localtime returns month from [0-11]. Get the month in letters using our array
my $date = sprintf('%d-%s-20%02d', $day, $month, $year);

my $dbname = 'elastin';  # make sure dbname is correct!
my $server = 'localhost';
my $user = 'webuser';
my $dbh = DBI->connect("DBI:Pg:dbname=$dbname", "$user", "", { RaiseError=>1 });

# retrieve the sequence from database
my $sequence = &getSequence($dbh, $seqname); 
#print $sequence."\n";



my $richseq = new Bio::Seq::RichSeq; # initialize RichSeq object

# get seqinfo
my ($taxid, $desc, $orgname, $comname, $lineage, $dbxref, $strctypeid, $seqtypeid, $geneid, $locprefix, $genename) 
	= &getSeqInfo($dbh, $seqname); 
#print "$taxid\t$desc\t$orgname\t$comname\n";

if( defined($dbxref) ) {
	# access the remote NCBI to retrieve the genbank record if we have a gi number
	my $gb = new Bio::DB::GenBank;
	$richseq = $gb->get_Seq_by_gi($dbxref) or die "Can't retreive sequence by gi $dbxref: $!";
	#print "\n/--- Retrieved GenBank record from NCBI ---/\n";
} 
else { #else, create our own richseq object and build the flat file
	my @result;


	# parse the seqname
	my ($seqtype_letter, $seqnumber);
	if( $seqname =~ /^(\w{5})(\w)?(\d+)/ ) {
		if($1 eq $locprefix) {
			$seqtype_letter = $2;
			$seqnumber = sprintf("%02d", $3);;
			#print "$seqname matched!!! $seqnumber\n";
		}
	} # else if the seqname does not match the above naming convention, we won't annotate any extra features (ie: exons, introns)

	# construct RichSeq object
	$richseq = Bio::Seq::RichSeq->new( -seq => $sequence,
						-id => $seqname,
						-accession_number => $seqname,
						-desc => $desc
	);


	if ( $format eq 'genbank' ) { # skip this section if the desired format is simply fasta

	    #########################
	    # add a species feature #
	    #########################
	    # use subroutine to create the species object
	    my $species = &parseSpecies($taxid, $orgname, $comname, $lineage); 
	    # add species to the RichSeq object
	    $richseq->species($species);


	    ########################
	    # add a source feature #
	    ########################
			my ($seqstart, $seqend) = (1, $richseq->length);
	    my $source = Bio::SeqFeature::Generic->new(  -start => 1,
	    						-end => $richseq->length,
							-primary => 'source');
	    
	    $source->add_tag_value('mol_type', 'genomic DNA') if ( $strctypeid == 3 || $strctypeid == 2 || $strctypeid == 4); 
			# 3 is unique id for 'genomic', 2 for 'exon', 4 for 'intron'
			$source->add_tag_value('mol_type', 'mRNA') if ( $strctypeid == 6 ); # 6 for transcript
	    $source->add_tag_value('organism', $orgname); # binomial('FULL') gets the genus, species and subspecies (if available)
	    $source->add_tag_value('dbxref', "taxon:$taxid");
	    $richseq->add_SeqFeature($source); # add this source feature to the richseq

	    #add annotations
	    $richseq->division(' ');
	    $richseq->add_keyword($genename);
	    $richseq->add_keyword('ELN') if ($genename eq 'elastin');
	    $richseq->add_date($date);

	    #####################
	    # add a CDS feature #
	    #####################
			# get CDS seqstart and seqend from database
			@result = $dbh->selectrow_array("SELECT seqstart, seqend
																												FROM splicedinfo
																												WHERE seqname=?", undef, $seqname);
			# if a CDS seqstart and seqend exist, assign these values
			# otherwise, use the defaults (1 and length)
			if(defined($result[0])) {
				$seqstart = $result[0];
				$seqend = $result[1];
			}
	    my $cds = Bio::SeqFeature::Generic->new( -start => $seqstart,
						-end => $seqend,
						-primary => 'CDS' );
	   	# annotate this feature 
	    $cds->add_tag_value('dbxref', "GeneID:$geneid") if ( defined($geneid) ); # add geneID reference if available
			$cds->add_tag_value('gene', $genename); # annotate the translation if sequence is genomic / nucleotide (6) or cDNA (3) 
	    $cds->add_tag_value('translation', $richseq->translate->seq) if ( $strctypeid == 2 ); # translate exons (2)
			
	    ###################################
	    # create SplitLocation CDS feature #
	    ###################################
			# if number exists and is a genomic sequence
			my $exonfeat;
			if( defined $seqnumber && $seqtype_letter ne 'P' && $seqnumber == 00) {
				# grab exon info from database
					# also get the domaintype by matching the exon with its corresponding domain (if available)
				my $exonsref = $dbh->selectall_arrayref("
					SELECT d1.number, d1.seqstart, d1.seqend, dt.type
					FROM splicedinfo d1
						, splicedinfo d2
						, domaintype dt
					WHERE 
						d1.number=d2.number
						AND d2.dtypeid=dt.dtypeid
						AND d1.seqname like ?
						AND d2.seqname like ?
						AND d1.number <> 0
					ORDER BY d1.seqstart, d1.seqname
				", undef, $locprefix."__", $locprefix."P__"); # bind values: match with exons; match with domians

				# create a SplitLocation obj
				my $splitloc = new Bio::Location::Split();
				foreach(@$exonsref) {
					# check if a seqstart and seqend exists
					if( !defined($_->[1]) or !defined($_->[2])) {
						$splitloc = undef; # nullify the object
						last; # exit this loop
					}
					# insert Location objects into the SplitLocation
					$splitloc->add_sub_Location(new Bio::Location::Simple(	-start=>$_->[1],
																																	-end=>$_->[2] ));
				}

				if(defined($splitloc)) { # add the splitlocation to CDS feature if it is defined
					# set the CDS location
					$cds->location($splitloc);

					# translate the genomic sequence
					my @sublocs = $splitloc->sub_Location(1); # retrieve array of sublocations; 1 for forward sort
					@sublocs = map { $richseq->subseq($_->start, $_->end) } @sublocs; # map the subsequences for each exon
					my $splicedseq = Bio::Seq->new(	-seq=> join('', @sublocs)); # create new Seq obj with subsequences joined
					$cds->add_tag_value('translation', $splicedseq->translate->seq); # annotate with translation
				}

				$richseq->add_SeqFeature($cds);

				########################
				# create exon features #
				########################
				foreach(@$exonsref) {
					# make sure a seqstart and seqend coordinate is defined
					if( defined($_->[1]) and defined($_->[2]) ) {
						$exonfeat = Bio::SeqFeature::Generic->new(	-start => $_->[1],
																													-end =>	$_->[2],
																													-primary => 'exon');
						# annotate this feature
						$exonfeat->add_tag_value('number', $_->[0]);
						$exonfeat->add_tag_value('note', $_->[3]) if (defined($_->[3])); # add note about domain type if available
						$richseq->add_SeqFeature($exonfeat);
					}
				}

			} # end if seqnumber
			else {
				# add the CDS feature
				$richseq->add_SeqFeature($cds);
			} # end else
	} # end if format eq 'genbank'
} # end else

# write the sequence to output
&writeSeq($richseq, $orgname, $dir, $seqname, $format);

$dbh->disconnect;

# return the sequence from the database
sub getSequence {
	my ($dbh, $seqname) = @_;	
	#print "\n/--- Searching for $seqname ---/\n";

	my @row_ary = $dbh->selectrow_array("SELECT sequence FROM sequence WHERE seqname=?", undef, $seqname);

	# terminate the program if the sequence is not found in the database
	die "$seqname not found!! Aborting...\n" unless ( defined($row_ary[0]) );
	#print "$seqname ... found! :)\n"; # on success!
	return $row_ary[0]; # return the value of the first element (ie: the sequence)
} # end sub

# get sequence info including taxonomic information
# return all the information as a list
sub getSeqInfo {
	my ($dbh, $seqname) = @_;
	#print "\n/--- Getting seqinfo ---/\n";

	my @row_ary = $dbh->selectrow_array("
					SELECT g.taxid, s.description, o.name, o.comname, o.lineage, s.dbxref, s.strctypeid, s.seqtypeid, g.dbxref, g.locprefix, gn.name
					FROM seqinfo s, genemodel g, organism o, genename gn
					WHERE s.seqname = ?
					  AND s.modid = g.modid
					  AND g.taxid = o.taxid
						AND g.genenameid = gn.genenameid
					", undef, $seqname);
	return @row_ary; # return the entire array
} # end sub

# create a Bio::Species object from the organism name and common name
sub parseSpecies {
	my ($taxonid, $name, $common, $lineage) = @_;

	# split the lineage into elements of an array
	my @class_ary = split(/; /, $lineage); # delimiter: '; ' (semi-colon followed by a space)
	@class_ary = reverse @class_ary; # bioperl needs classification inputted in the reverse order to the Species object

	# use regex to split up the organism name into its components: genus species subspecies
	$name =~ /(^[A-Z][\w.]+)\s([\w.'",]+)\s*(.*)/ ;
	my ($genus, $species, $subspecies) = ($1, $2, $3);

	my $spe_obj = Bio::Species->new( -classification => \@class_ary );

	# annotate the species object
	$spe_obj->sub_species($subspecies) unless ( $3 eq '' ); # don't annotate if no subspecies
	$spe_obj->common_name($spe_obj->binomial() . " (".$common.")");
	$spe_obj->ncbi_taxid($taxonid);
	return $spe_obj; 
} # end sub

# construct SeqIO object and write the sequence to output
sub writeSeq {
	#arguments: RichSeq object, organism name, output directory, seqname, format
	my ($seq_obj, $orgname, $dir, $seqname, $format) = @_;
	my $seqout;
	#print "\n/--- Outputting sequence to " . uc($output) . " ---/\n";

	if ( $output eq 'STDOUT' ) {
		$seqout = Bio::SeqIO->new(  -fh => \*STDOUT, 
						-format => $format );
	} else {
		# check if the parent directory ($dir) has a folder already created for this organism
		#opendir DH, $dir or die "cannot open $dir: $!";
		#my $folder = undef;
		my $org_folder = $orgname;
		$org_folder =~ s/ /_/g ; # generate the folder name by replacing all spaces with underscore

		# check to see if this folder exists
		if(!-d $dir) {
			#print "Directory $org_folder does not exist...creating...\n";

			# create the directory
			mkdir $dir, 0777 or die "cannot mkdir for $dir: $!";
		}

		# set directory to location where file will be outputted
		chdir $dir or die "cannot chdir to $dir: $!";

		# create the org_folder if it does not exist
		if(!-d $org_folder) {
			mkdir $org_folder, 0777 or die "cannot mkdir for $org_folder: $!";
		}

		#print "Changing directory to $dir/$org_folder ...\n";
		chdir "$dir/$org_folder" or die "cannot chdir to $org_folder: $!";

		my $filenameFormat = $seqname . ".$format";
		$seqout = Bio::SeqIO->new( -file => ">$filenameFormat",
						-format => $format );
		#print "Written to file $filenameFormat...\n";
	}
	# write the output
	$seqout->write_seq($seq_obj) or die "Failed writing to " . uc($output);
} # end sub
