<filename>genomics/seqExpts/core.py
#!/usr/bin/env python

import os, datetime, commands, pdb, sys
import arrayUtils, cmdlineProgs, utility
import xml.etree.ElementTree as ET
import genomics.formats.conversion as conversion


class LibrarySequencing:
    """Represents a set of sequences resulting from
    the sequencing of a single library. Can store the sequences
    as a fastq file and/or the corresponding alignments as a bam file."""

    # FIXME: I wonder if I should introduce some kind of "mapping" class too,
    # and have it point to the library sequencing and/or vice-versa. At the
    # moment this seems like overkill.

    def __init__(self, library):
        self.fastqFilename = None
        self.bamFilename = None

        # The library that was sequenced:
        self.library = library

    def getBam(self):
        return self.bamFilename

    def getFastq(self):
        return self.fastqFilename

    def setBam(self, bamFilename):
        self.bamFilename = bamFilename

    def setFastq(self, fastqFilename):
        self.fastqFilename = fastqFilename

    def getName(self):
        """Returns a string representing the name of this library sequencing."""

        # Returns a name based on the library that was sequenced, and returns
        # that string:
        return self.library.getName()

    def getLibrary(self):
        return self.library


class SequencingExpt:
    """Class (can be concrete, to represent a generic sequencing experiment)
    representing a generic experiment that produces high-throughput sequence
    data."""

    def __init__(self):
        # String key to LibrarySequencing value dictionary; holds all of the
        # sets of sequences comprising this experiment:
        self.libseqs = {}

    def setLibSeq(self, name, libseq):
        """Stores the specified LibrarySequencing object under the given name."""

        self.libseqs[name] = libseq

    def getLibSeq(self, name):
        """Returns the library sequencing denoted by the given name."""

        return self.libseqs[name]

    # FIXME: My design seems a bit broken here. Introducing this method for
    # cases where the experiment just has a single library sequencing:
    def getLibSeqUnique(self):
        """Returns the first and only library sequencing."""
        assert len(self.libseqs.values()) == 1

        return self.libseqs.values()[0]


class Protocol:
    """Abstract class representing a generic experimental protocol that
    produces a library."""

    def __init__(self):
        pass

    def getName(self):
        """Returns name of this protocol. Must be implemented by all concrete
        inheriting classes."""

        pass


class GenericProtocol(Protocol):
    """A generic protocol. This can be used when a protocol object is required
    to represent the fact that some protocol was carried out, but where it's not
    important what that protocol was. Introduced for use in a generic library
    complexity analysis (August 12th 2013)."""
    def __init__(self):
        Protocol.__init__(self)

    def getName(self):
        return "GenericProtocol"


class Library:
    """A library of sequences generated by a specified protocol."""

    def __init__(self, protocol):
        self.protocol = protocol

        # The date that the library was produced. Starts as unknown,
        # but could be set by (e.g.) parsing from XML metadata of
        # an experiment from SRA/GEO etc.:
        self.date = None

    def setDate(self, date):
        # A python standard library's date object
        self.date = date

    def getName(self):
        """Returns name of the library."""

        # Currently just based on the protocol that generated the library:
        return self.protocol.getName()

    def getProtocol(self):
        return self.protocol


class SeqAnalysis:
    """Abstract class representing a generic analysis on a single
    sequencing experiment."""

    def __init__(self, expt, name, outDir):
        """SequencingExpt instance and name of experiment as inputs."""
        self.expt = expt # The analysed experiment
        self.name = name # The name of the analysis
        self.outDir = outDir # The output directory for the analysis
        self.runDate = None

    def run(self):
        """Must be implemented by all extending sub-classes. This
        super-method just sets the date of the analysis."""

        self.runDate = datetime.datetime.now()

    def getRunDate(self):
        return self.runDate

    def beenRun(self):
        return self.runDate != None

    def getExpt(self):
        return self.expt

    def getName(self):
        return self.name

    def editXMLelem(self, xmlElem):
        """Edits the specified XML ElemenTree object, adding a node
        for this analysis."""

        # Make root node:
        analysisNode = ET.SubElement(xmlElem, 'Analysis')

        # Make name node:
        nameNode = ET.SubElement(analysisNode, 'Name')
        nameNode.text = self.getName()

        # Make analysis type node:
        typeNode = ET.SubElement(analysisNode, 'Type')
        typeNode.text = str(self.__class__.__name__)

        # Make date of running node:
        runDateNode = ET.SubElement(analysisNode, 'RunDate')
        runDateNode.text = str(self.getRunDate())

        # Node ID of experiment analysed:
        # FIXME (when XML requirements are clearer): Not sure how to do this.

        # Content that is specific to individual analyses:
        self.editIO_XML(analysisNode)
        self.editParams_XML(analysisNode)

    # def writeXML(self, outfile):
    #     """Writes a (sub-)tree of XML to the specified output file."""

    #     xmlTree = ET.ElementTree()
    #     self.editXMLelem(xmlElem)
    #     xmlTree.write(outfile)
    #     outfile.flush()
        
    def editIO_XML(self, xmlElem):
        """Must be implemented by subclasses."""
        pass

    def editParams_XML(self, xmlElem):
        """Must be implemented by subclasses."""
        pass


class SingleLibAnalysis(SeqAnalysis):
    """Abstract class representing an analysis of a single sequencing
    library."""

    # FIXME: Just keeping this basic at this point. Not 100% sure if this
    # design makes sense.

    def __init__(self, expt, name, outDir):
        SeqAnalysis.__init__(self, expt, name, outDir)


class ComplexityAna(SingleLibAnalysis):
    """An analysis of library complexity."""

    def __init__(self, expt, name, outDir, run_cCurve=True, run_lcExtrap=True, 
                 preseqOptionsStr=""):
        SingleLibAnalysis.__init__(self, expt, name, outDir)

        # Record complexity curve and extrapolation output files:
        self.cCurve_filename = self.outDir + "/" + name + "_cCurve.txt"
        self.lcExtrap_filename = self.outDir + "/" + name + "_lcExtrap.txt"
        self.run_cCurve = run_cCurve
        self.run_lcExtrap = run_lcExtrap
        self.preseqOptionsStr = preseqOptionsStr
        self.cCurve_matrix = None
        self.lcExtrap_matrix = None
        self.cCurvePlotFile = None
        self.lcExtrapPlotFile = None

    def run(self):
        """Uses preseq to perform the complexity analysis."""

        self.runPreseq()
        self.generatePlots()

    def runPreseq(self):
        """Runs c_curve and lc_extrap."""

        # Generate temporary bed filename, and convert the bam file to bed
        # format:
        tmpBedFilename = utility.makeTempFilename("TmpBedFile")
        inputBamFilename = self.expt.getLibSeqUnique().getBam()
        print >> sys.stderr, "Converting bam to temporary bed..."
        conversion.bam2bed(inputBamFilename, tmpBedFilename)
        print >> sys.stderr, "Done."

        # Run c_curve and lc_extrap on the bed file if requested...
        if self.run_cCurve:
            print >> sys.stderr, "Running c_curve..."
            cmdlineProgs.run_c_curve(tmpBedFilename,
                                     self.cCurve_filename,
                                     optionsStr = self.preseqOptionsStr,
                                     verbose=True)
            print >> sys.stderr, "Done."

            # Check that an output file exists and quit, reporting an error,
            # if it does not:
            try:
                with open(self.cCurve_filename): pass
            except IOError:
                print 'Error: c_curve failed to produce a file.' + \
                    self.cCurve_filename + ' Quitting.'
                sys.exit(1)

            # Parse the output matrix:
            self.cCurve_matrix = \
                arrayUtils.parseMatrix(open(self.cCurve_filename),
                                       header=True)

        if self.run_lcExtrap:
            print >> sys.stderr, "Running lc_extrap..."
            cmdlineProgs.run_lc_extrap(tmpBedFilename,
                                       self.lcExtrap_filename,
                                       optionsStr = self.preseqOptionsStr,
                                       verbose=True)
            print >> sys.stderr, "Done."

            # Check that an output file exists and quit, reporting an error,
            # if it does not:
            try:
                with open(self.lcExtrap_filename): pass
            except IOError:
                print 'Error: lc_extrap failed to produce a file.' + \
                    self.lcExtrap_filename + ' Quitting.'
                sys.exit(1)

            # Parse the output matrix:
            self.lcExtrap_matrix = \
                arrayUtils.parseMatrix(open(self.lcExtrap_filename),
                                       header=True)

        cmdlineProgs.deleteFiles([tmpBedFilename])

    def generatePlots(self):
        """Generate plots of the complexity analyses, and save as png
        files."""

        if self.cCurve_matrix != None:
            self.cCurvePlotFile = self.outDir + "/" + self.name + "_cCurve.png"
            utility.makePlot(self.cCurve_matrix, self.cCurvePlotFile,
                             xlabel="Number of reads sampled",
                             ylabel="Number of unique reads recovered",
                             title="Complexity estimation plot: " + self.name)

        if self.lcExtrap_matrix != None:
            self.lcExtrapPlotFile = self.outDir + "/" + self.name + "_lcExtrap.png"
            utility.makePlot(self.lcExtrap_matrix, self.lcExtrapPlotFile,
                             xlabel="Number of reads sampled",
                             ylabel="Number of unique reads recovered",
                             title="Complexity estimation plot: " + self.name)


class SeqExptAnalyser:
    """An analyser for a single sequencing experiment. Carries out zero
    or more analyses on that experiment."""

    def __init__(self):
        # The analyses carried out by this analyser.
        # Keys are names of analyses, values are analyses:
        self.analyses = {}

    def getAnalysis(self, analysisName):
        return self.analyses[analysisName]

    def getAnalyses(self):
        return self.analyses

    def addAnalysis(self, analysis):
        # Ensure that the experiment is the same as for any other
        # analyses:
        for currAnalysis in self.analyses.values():
            if not analysis.getExpt() == currAnalysis.getExpt():
                pdb.set_trace()
                dummy = 1
            assert(analysis.getExpt() == currAnalysis.getExpt())

        self.analyses[analysis.getName()] = analysis

    def getExpt(self):
        if len(self.analyses.values()) > 0:
            return self.analyses.values()[0].getExpt()
        else:
            return None

    def runAnalyses(self):
        for analysis in self.getAnalyses().values():
            if (not analysis.beenRun()):
                print >> sys.stderr, "Running analysis", \
                    analysis.getName(), "..."
                analysis.run()
                print >> sys.stderr, "Completed analysis", analysis.getName(), \
                    "."

    def edit_XML(self, xmlElem):
        """Makes an XML node for the analyser itself within an existing XML
        node."""

        analyserNode = ET.SubElement(xmlElem, 'Analyser')
        addXMLcontent(analyserNode)

    def makeXML(self):
        """Makes a new XML node for the analyser."""

        # Add XML content for the analyser itself:
        analyserNode = ET.Element('Analyser')
        self.addXMLcontent(analyserNode)
        return analyserNode

    def addXMLcontent(self, analyserXmlElem):
        """Add the XML content for and all this analyser's analyses and
        experiment(s)."""

        # Add XML content for the experiment:
        expt = self.getExpt()
        if expt != None:
            # FIXME: When XML requirements become clearer, then write general
            # and specific xml output methods for the chip-seq experiment class
            # and parent class. Currently not used:
            #expt.editXMLelem(analyserXmlElem)
            pass
        
        # Add XML content for the analyses:
        for analysis in self.analyses.values():
            analysis.editXMLelem(analyserXmlElem)

        # FIXME: Perhaps add sub-nodes linking to the analyses?

    def writeXML(self, outfile):
        """Writes a tree of XML to the specified output file."""

        rootElem = self.makeXML()
        xmlTree = ET.ElementTree(rootElem)
        xmlTree.write(outfile)
        outfile.flush()


class AnchoredSeqExptAnalyser(SeqExptAnalyser):
    """An analyser with a specified base directory that must be
    consistent for all analyses run."""

    def __init__(self, baseDir):
        SeqExptAnalyser.__init__(self)

        self.baseDir = baseDir
