####################################################################
# parallelCommon.py
#   routines used by: divide_data.py
#                     combine_replicates.py
#                     combine_regions.py
####################################################################
# system imports
import fileinput
import getopt
import random
import re
import sys
import os.path
from xml.dom.minidom import parse, Document

seedlim = 1000000


# share
def usage():
    print '''Usage:
        --lamarcfile REQUIRED <path to lamarc infile to fragment>
        --outdir     REQUIRED <path to output directory>
'''


# share
def getOptionsAndVerify(barfIfOutdirExists):
    lamarcfile = None
    outdir = None

    try:
        opts, args = getopt.getopt(sys.argv[1:],"hl:o:",["help","lamarcfile=","outdir="])
    except getopt.GetoptError, err:
        usage()
        sys.exit(2)

    for o, a in opts:
        if o in ("-h","--help"):
            usage()
            sys.exit()
        elif o in ("-l", "--lamarcfile"):
            lamarcfile = a
        elif o in ("-o", "--outdir"):
            outdir = a
        else:
            assert False, "unhandled option"

    # lamarc input file and output directory both must exist
    if lamarcfile == None or outdir == None:
        usage()
        sys.exit()

    # make sure lamarcfile exists and is a file
    if not os.path.isfile(lamarcfile):
        print "%s doesn't exist or isn't a file. exiting" % lamarcfile
        sys.exit(1)

    # make sure outdir is a new, empty directory
    if os.path.exists(outdir):
        if barfIfOutdirExists:
            print "directory %s already exists. exiting" % outdir
            sys.exit(1)
    else:
        os.makedirs(outdir)

    return (lamarcfile,outdir,os.getcwd())


# gets first sub-tag named "tagname" of element "domElem"
def getFirstTag(domElem,tagname):
    elems = domElem.getElementsByTagName(tagname)
    for elem in elems:
        return elem
    return None

# gets only sub-tag named "tagname" of element "domElem"
# EWFIX -- not actually doing any error checking
def getSingleTag(elem,tagname):
    return getFirstTag(elem,tagname)

# gets the long value embedded in a tag
def getLongVal(elem):
    for child in elem.childNodes:
        # EWFIX -- should check only one
        if child.nodeType == child.TEXT_NODE:
            return long(child.data)
    return None
        

# replaces data in the first text child node of "elem" with "strData"
def setVal(elem,strData):
    for child in elem.childNodes:
        # EWFIX -- should check only one
        if child.nodeType == child.TEXT_NODE:
            child.data = strData
            return
        
# when creates or replaces child tag "childTagName" of parent element
# "parentTag" with text of value "newValue"
def replaceOrSet(domTop,parentTag,childTagName,newValue):

    childTag = getSingleTag(parentTag,childTagName)
    if childTag:
        setVal(childTag,newValue)
    else:
        childTag = domTop.createElement(childTagName)
        parentTag.appendChild(childTag)
        parentTag.appendChild(domTop.createTextNode("\n"))
        childTag.appendChild(domTop.createTextNode(newValue))
        

# change the name (or set if not present) of all input and outputs specified
# in the format tag to have an infix of "idStr"
#
# also sets seed to an appropriate random value
#
# input "useIn" when true, controls input summary file reading
def fixFormatTag(domTop,fmtTag,idStr,useIn):
    replaceOrSet(domTop,fmtTag,"seed", "%d" % (random.randint(0,seedlim)*4+1))
    replaceOrSet(domTop,fmtTag,"results-file", "outfile_%s.txt" % idStr)
    if useIn:
        replaceOrSet(domTop,fmtTag,"use-in-summary", "true")
        replaceOrSet(domTop,fmtTag,"in-summary-file", "insumfile_%s.xml" % idStr)
    else:
        replaceOrSet(domTop,fmtTag,"use-in-summary", "false")
    replaceOrSet(domTop,fmtTag,"use-out-summary", "true")
    replaceOrSet(domTop,fmtTag,"out-summary-file", "outsumfile_%s.xml" % idStr)
    replaceOrSet(domTop,fmtTag,"curvefile-prefix", "curve_%s.xml" % idStr)
    replaceOrSet(domTop,fmtTag,"tracefile-prefix", "trace_%s.xml" % idStr)
    replaceOrSet(domTop,fmtTag,"newicktreefile-prefix", "newick_%s.xml" % idStr)
    replaceOrSet(domTop,fmtTag,"out-xml-file", "menusettings_%s.xml" % idStr)
    replaceOrSet(domTop,fmtTag,"xml-report-file", "report_%s.xml" % idStr)


#
#profilePattern = re.compile('\s*(\S+)\s+(.*)')
#def replaceEachProfile(profileText):
#    if not profileText:
#        return
#    remainingText = profileText
#    outText = ""
#    m = profilePattern.match(remainingText)
#    while m:
#        outText = "%s %s" % ( outText, "none")
#        remainingText = m.group(2)
#        m = profilePattern.match(remainingText)
#    return outText

# get population count from a top level lamarc XML document
def getPopCount(lamDom):
    dataTag = getSingleTag(lamDom,"data")
    regionTag = getFirstTag(dataTag,"region")
    return len(regionTag.getElementsByTagName("population"))

# makes a string containing the word "none" "numProfiles" times
# spaces are between each occurence and before the first and after
# the last
def makeProfilesString(numProfiles):
    outstr = ""
    for i in range(numProfiles):
        outstr = "%s none" % outstr
    outstr = "%s " % outstr
    return outstr

# turn off the profile for a given force
def turnOneProfileOff(domTop,forcesTag,forceName,numProfiles):
    thisForce = getSingleTag(forcesTag,forceName)
    if thisForce:
        newProfilesString = makeProfilesString(numProfiles)
        replaceOrSet(domTop,thisForce,"profiles",newProfilesString)


# turn off all profiles -- we use this because we don't want to
# do any profiling until the last run of lamarc on the brought-together
# data
def turnProfilesOff(domTop,lamTag):
    popCount = getPopCount(domTop)
    tag = getSingleTag(lamTag,"forces")
    turnOneProfileOff(domTop,tag,"coalescence",popCount)
    turnOneProfileOff(domTop,tag,"migration",popCount*popCount)
    turnOneProfileOff(domTop,tag,"growth",popCount)
    turnOneProfileOff(domTop,tag,"recombination",1)

           
# python outputs XML with an initial tag <?xml version="1.0" ?>
# even though this is perfectly correct, the lamarc sum file reading
# routines cannot handle it. This is a kludge to work around that
linePattern = re.compile('(.*)<\?.*\?>(.*)')
def stripXmlInfo(filename):
    firstLine = True
    for line in fileinput.input(filename,inplace=1): 
        if firstLine:
            m = linePattern.match(line)
            if m:
                line = "%s%s" % (m.group(1),m.group(2))
            firstLine = False
            print line      # prints line to file via fileinput inplace option
        else:
            line.rstrip()
            print line ,    # prints line to file via fileinput inplace option

def sumfileCombineWarn(where):
    print "Combining sumfiles in %s. This can take a while..." % where

def dividerLine():
    print "****************************************************"

def describeThisScript(scriptName,purpose,lamfile,lamdir):
    dividerLine()
    print "running python script: %s" % scriptName
    print "purpose:               %s" % purpose
    print "arguments:             -l %s" % lamfile
    print "                       -o %s" % lamdir

def nextStep(lamdir,filelist,nextIsLast):
    dividerLine()

    print "next step: run lamarc on",
    if not nextIsLast:
        print "each of:",
    print

    for fileName in filelist:
        print "               %s" % fileName
    print "           if you are running on a different machine, copy"
    print "           all files in each of the above directories there"
    if not nextIsLast:
        print "           and copy all files generated by the run back to"
        print "           the same directory"
    print

def finalStep(pydir,pyscript,lamfile,lamdir):
    print "then: change back to this directory"
    print "      (%s)" % pydir
    print "      and run python script %s with arguments" % pyscript
    print "           -l %s" % lamfile
    print "           -o %s" % lamdir
    dividerLine()