#!/usr/bin/env python
__description__ = \
"""
A set of objects and functions for blasting sequences against local and 
remote servers.  Also functions for filtering sequences.  Uses Biopython.  It
also requires ncbi's blast+ tools and cdhit to be installed and visible in 
the path.
"""

__author__ = "Michael J. Harms"
__date__ = "110528"
__usage__ = "orthologBlast.py seed_fasta_file"
__version__ = "0.1"

# Modules for running sundry processes
import subprocess, shlex, re, sys, os, string

# Modules for blasting, etc. 
from Bio import SeqIO, Entrez
from Bio.Seq import Seq
from Bio.Blast.Applications import NcbiblastpCommandline as blastp
from Bio.Blast import NCBIWWW, NCBIXML

# Modules for parsing XML
from xml.etree import ElementTree as ET

# Set email address here for NCBI/Entrez server
Entrez.email = "PUT_YOUR_EMAIL_ADDRESS_HERE"

# Global unique name counter
homolog_counter = 0

class BlastToolsError(Exception):
    """
    General error class for this module.
    """

    pass


class Homolog:
    """
    A class that holds homologs (minimally, accession and name; maximally,
    sequence information and name of ortholog).
    """

    def __init__(self,definition,accession):
        """
        Create a new homolog.
        """   

        global homolog_counter

        # Record homolog definition and accession number 
        self.definition = str(definition)
        self.accession = str(accession)
        self.rank = sys.maxint

        # Assign this homolog a unique name.
        self.unique_name = "%i" % (homolog_counter)
        self.unique_name = self.unique_name.zfill(8)
        self.unique_name = "XX%s" % self.unique_name
        homolog_counter += 1

        # Initialize other variables that we don't yet know
        self.sequence = None
        self.length = None
        self.taxid = None
        self.organism = None

        self.ortholog_name = None

    def loadSequence(self,sequence,length,taxid,organism):
        """
        Load sequence data for this homolog.  
        """

        self.sequence = Seq(sequence)
        self.length = length
        self.taxid = taxid
        self.organism = organism

    def formatFasta(self,seq_name=None):
        """
        Return a fasta formatted string with either "unique_name" or seq_name
        as the name.
        """

        if seq_name != None:
            name = seq_name
        else:
            seq_name = self.unique_name

        return ">%s\n%s\n" % (seq_name,self.sequence)

    def formatTabDelim(self):
        """
        Write the data for this ortholog out onto a single, tab-delimited 
        line.
        """

        pretty_name = "%s-%s-%s" % (self.ortholog_name,self.organism,
                                    self.accession)
            
        to_write = (self.unique_name,
                    self.ortholog_name,
                    self.organism,
                    self.rank,
                    self.accession,
                    self.length,
                    str(self.sequence),
                    self.definition,
                    pretty_name)

        to_write = "\t".join(["%r" % w for w in to_write])

        return "%s\n" % to_write


def environmentCheck():
    """
    Make sure that all accessory programs are available in the path.
    """

    print "Checking for required external programs...", 

    to_check = ["cdhit","blastp","tblastn","makeblastdb"]

    failed = []
    for c in to_check:
        args = shlex.split(c)
        try:
            out = subprocess.Popen(args,stdout=subprocess.PIPE,
                                        stderr=subprocess.PIPE)
        except OSError:
            failed.append(c)
        except:
            print "Unexpected error when running %s:" % c, sys.exc_info()[0]
            raise 

    if len(failed) > 0:
        print "FAILURE\n"
        return failed
    else:
        print "SUCCESS\n"

    if Entrez.email == None:
        print "No email address has been set!  To avoid typing this in the"
        print "future, edit the line 'Entrez.email = None' to point to your"
        print "email address (e.g. Entrez.email = \"A.N.Other@example.com\")."
        print "\nPlease visit the entrez website for more information about"
        print "Entrez usage rules and why providing an email address is useful."
        print ""
        print "http://www.ncbi.nlm.nih.gov/entrez/query/static/eutils_help.html#UserSystemRequirements"
        print ""

        email = raw_input("Please enter your email address:\n")

        Entrez.email = email

        print ""

    return [] 


def flattenConcatenatedXML(input_file,key_tag):
    """
    Clean up naively concatenated XML files by deleting begin/end tags that
    occur at the place where the two files were concatenated.

    NOTE: This will break and break royally if the key_tags are on the same
    lines as other important entries. 
    """

    f = open(input_file,'r')
    input = f.readlines()
    f.close()

    set_start = re.compile("<%s>" % key_tag)
    set_end =   re.compile("</%s>" % key_tag)

    # Find all beginning and end tags...
    starts = [i for i, l in enumerate(input) if set_start.search(l) != None]

    # If this tag occurs more than once...
    if (len(starts) != 1):
    
        # Keep the first start reverse so we are chewing from the bottom.
        starts.pop(0)
        starts.reverse()

        # Remove all lines between each end and start, again chewing from the
        # bottom.
        for i in range(len(starts)):
            e = starts[i]
            while set_end.search(input[e]) == None:
                input.pop(e),
                e = e - 1
            input.pop(e)

    # Return freshly minted, clean XML
    return "".join(input)


def parseFastaXML(sequence_file):
    """
    Load a set of sequences from an NCBI xml/fasta file into a set of Homolog
    objects.
    """
   
    homolog_list = []
 
    print "Parsing sequences in %s" % sequence_file, 

    # Fix screwed up XML because sequences downloaded and output concatenated 
    sequence_input = flattenConcatenatedXML(sequence_file,"TSeqSet")

    # Now we should have valid XML...
    sequence_input = ET.XML(sequence_input)
    for i, sequence in enumerate(sequence_input):
        properties = dict([(s.tag,str(s.text)) for s in sequence]) 
       
        definition = properties["TSeq_defline"].strip()
        definition = re.sub("\t"," ",definition) 
        accession = properties["TSeq_gi"].strip()
       
        homolog_list.append(Homolog(definition,accession))
        homolog_list[-1].loadSequence(properties["TSeq_sequence"].strip(),
                                      int(properties["TSeq_length"]),
                                      int(properties["TSeq_taxid"]),
                                      properties["TSeq_orgname"].strip())

    print "DONE."
    
    return homolog_list

def parseBlastXML(blast_file,tag_list=("Hit_def","Hit_id")):
    """
    Parse BLAST xml output, extracting tags specified in tag_list and putting
    into a list.  E-value is always appended after the last requested tag.  
    """

    # Fix screwed up XML if blasts were done in series...
    blast_input = flattenConcatenatedXML(blast_file,"BlastOutput_iterations")
   
    # Read blast file properties (in tag_list) into a list to dump out 
    blast_input = ET.XML(blast_input)

    all_hits = []
    for blast_output in blast_input:
        if blast_output.tag != "BlastOutput_iterations":
            continue

        for iteration in blast_output:
            if iteration.tag != "Iteration":
                continue

            for hits in iteration:
                if hits.tag != "Iteration_hits":
                    continue

                for hit in hits:
                    properties = dict([(h.tag,str(h.text)) for h in hit]) 
                    all_hits.append([properties[t] for t in tag_list])

                    for property in hit:
                        if property.tag == "Hit_hsps":
                            for hsp in property:
                                hsp_properties = dict([(p.tag,str(p.text))
                                                       for p in hsp])
                                all_hits[-1].append(hsp_properties["Hsp_evalue"])