#!/usr/bin/python
'''Currently SplitFreqSuppliervia a file-based wrapper around MRBAYES3.1.1's mcmc/sumt.'''
from PIPRes.util.server_import import *
from PIPRes.cipres_types import * 
from PIPRes.basic import writeTreesBlock
import tempfile, os
from PIPRes.wrap.nexus_program_server import *
from PIPRes.wrap.output_parser import readMBSplitFreqs
from PIPRes.util.io import tryRemove, tryRmdir
from PIPRes.service_impl.split_freq_supplier import SplitFreqSupplier

class MBWrap(object, SplitFreqSupplier, SimpleNexusWrapperServer):
    debugging = False
    def __init__(self, registry, mbPath):
        self.debugging = MBWrap.debugging
        SimpleNexusWrapperServer.__init__(self, registry)
        self.convertTreeFunc = self.orb is None and toTreeBridge or toIDLTree
        self.registry = registry
        self.mat = None
        self.defaultCmdPre = '''
        lset nst = 1 rates=equal;
        mcmcp stoprule = yes;\n'''
        self.defCmdPost = '''
        mcmc;
        sumt;\n'''
        
        if not os.path.exists(mbPath):
            raise ValueError, 'The specified path to MRBAYES (%s) does not exist.' % mbPath
        self.mbPath = mbPath
        self.createdSuffixes = ['.con', '.mcmc', '.parts', '.trprobs' ]
        self.nRuns = 2
        self.perRunSuffixes = ['.p', '.t']
        self.mbVersion = (3, 1, 1) #@ should get this from arg or property
        self.touch()

    def _calcSplitsIfEmpty(self):
        if len(self.splitFreq) == 0:
            if not self.mat:
                raise RuntimeError, 'Failed to call setMatrix before requesting split frequency'
            self._mcmcAnalysis(self.getDefaultCmdPrefix() + self.userCommands() + self.defCmdPost)

    def getDefaultCmdPrefix(self):
        return self.defaultCmdPre

    def getSplitsTable(self):
        self._calcSplitsIfEmpty()
        return self.splitFreq

    def setMatrix(self, mat):
        _LOG.debug('setMatrix called\n')
        # _LOG.debug('matrix = %s\n' % mat)
        self.mat = mat
        self.touch()

    def touch(self):
        self.splitFreq = {}
        
    
    def _getExtraPaths(self, pref = 'data.nex'):
        extraPaths = ['log.out']
        extraPaths.extend([pref + s for s in self.createdSuffixes])
        if self.nRuns > 1:
            for r in range(self.nRuns):
                t = pref + 'run' + str(1 + r)
                extraPaths.extend([t + s for s in self.perRunSuffixes])
        else:
            extraPaths.extend([pref + s for s in self.perRunSuffixes])
        return extraPaths
        
    def _mcmcAnalysis(self, configCmd = ''):
        extInvContext = self._createTempDataFile(
                matrix=self.mat, 
                supportsTaxaSubsets=False, 
                writeTree=False, 
                dirPref='mbWrap')
        try:
            # compose the mb block
            mbBlock = '''BEGIN mrbayes;
            log start append;
            %s
            log stop;
            quit;
            END;
            ''' % (configCmd)
            extInvContext.fileObj.write(mbBlock)
            extInvContext.fileObj.close()

            self._invokeMB(extInvContext.tempDir, extInvContext.filename)

            splitsFile = os.path.join(extInvContext.tempDir, extInvContext.filename + '.parts')
            hasMissingWS = (self.mbVersion == (3,1,1))
            self.splitFreq = readMBSplitFreqs(splitsFile, missingWhitespaceBug=hasMissingWS)
        finally:
            if not self.debugging:
                self.cleanupTemporaries(extInvContext.tempDir, [extInvContext.filename] + self._getExtraPaths(extInvContext.filename))
            else:
                _LOG.debug('leaving the temporary directory %s (running in debug mode)' % extInvContext.tempDir)

    def _invokeMB(self, tempDirectory, fileName):
        self._callProcess(tempDirectory, [self.mbPath, str(fileName)])


class MBTwoExpWrap(MBWrap):
    '''Wrapper around Ziheng Yang and Bruce Rannala's tweak of MB (expects a single line of input:
    <internal branch length prior mean> <external prior mean>
    This wraps MB 3.0, so the stoprule is not implemented.'''
    def __init__(self, registry, mbPath, internalMean=0.001, externalMean=1.0):
        MBWrap.__init__(self, registry, mbPath)
        self._internalMean = internalMean
        self._externalMean = externalMean
        self.defaultCmdPre = ' lset nst = 1 rates=equal;'
        self.stdInFile = 'priorMeans.txt'
        self.mbVersion = (3, 0, 'YR')
        self.touch()
        self._nGenerations = 1000000

    def getDefaultCmdPrefix(self):
        return self.defaultCmdPre + ' mcmcp ngen = %d ; '% self._nGenerations

    def setInternalMean(self, x):
        self._internalMean = x
        self.touch()

    def setExternalMean(self, x):
        self._externalMean = x
        self.touch()

    def setNGenerations(self, x):
        self._nGenerations = x
        self.touch()

    def getInternalMean(self):
        return self._internalMean

    def getExternalMean(self):
        return self._externalMean
        
    def getNGenerations(self):
        return self._nGenerations
        
    externalMean = property(getExternalMean, setExternalMean)
    internalMean = property(getInternalMean, setInternalMean)
    nGenerations = property(getNGenerations, setNGenerations)

    def _getExtraPaths(self, pref = 'data.nex'):
        return MBWrap._getExtraPaths(self, pref) + [self.stdInFile]

    def _invokeMB(self, tempDirectory, fileName):
        pm = open(os.path.join(tempDirectory, self.stdInFile), 'w')
        pm.write('%lf %lf\nexecute data.nex;\nquit;\n' % (self.internalMean, self.externalMean))
        pm.close()
        self._callProcess(tempDirectory, 
                [self.mbPath, str(fileName), '<', self.stdInFile],
                invocationStyle=InvocationStyleEnum.kUseSystem) # 


_LOG = cipresGetLogger('pipres.service_impl.server.mb_wrap')

def _getPathToMBFromArgv(argv):
    MBWrap.debugging = '-debug' in argv
    if len(argv) < 2:
        raise ValueError, 'Expecting the path to mb as an argument'
    return argv[1]
    
def createMBWrap(registry, argv):
    return MBWrap(registry, _getPathToMBFromArgv([0] + argv)) # we add an empty element to act like the script name in a sys.argv list

def defaultMBWrap():
    '''Returns a MBWrap object with a getCipresRegistry, and path to mb from the environment.'''
    return MBWrap(getCipresRegistry(), os.environ['PATH_TO_MB'])

def defaultMBTwoExpWrap():
    '''Returns a MBTwoExpWrap object with a getCipresRegistry, and path to mbTwoExp from the environment.'''
    return MBTwoExpWrap(getCipresRegistry(), os.environ['PATH_TO_YR_MB'])



if __name__=='__main__':
    argv = sys.argv
    mbPath = _getPathToMBFromArgv(argv)
    
    def MBWrapFactory(registry, pathToMB = mbPath):
        return MBWrap(registry, pathToMB)
    try:
        interfaceDict = {}
        for k in ['SplitFreqSupplier']:
            interfaceDict[k] = MBWrapFactory
        cipresServe(argv, interfaceDict)
    except:
        logException(_LOG)