import types,string,cStringIO,sys,os
from Tree import Tree
from Node import Node,NodePart,NodeBranchPart
from Trees import Trees
from Nexus import Nexus,NexusData
from Glitch import Glitch
#from NexusToken import nextTok,safeNextTok,nexusSkipPastNextSemiColon
import NexusToken   # needed for cStrings
#import NexusToken2  # Faster, in c
import func
from Var import var

longMessage1 = """
This table shows, for splits that were used in the cons tree (ie not
necessarily all splits), which of those splits had the root on it in
the original bifurcating input trees.  The tree as returned from this
method is rooted on a node on the branch with the highest biRootCount
(or if there were several highest, it is rooted on a node arbitrarily
on the first of those).
"""

longMessage2 = """
This table shows, for nodes derived from splits that were used in the
cons tree (ie not necessarily all splits), which of those nodes were
roots in the original input trees.  The tree as returned from this
method is rooted on the node with the highest rootCount (or if there
are several highest, it is rooted arbitrarily on the first of
those).
"""


class SplitModelUsage(object):
    def __init__(self, TPMI):
        if not isinstance(TPMI, TreePartitionsModelInfo):
            gm = ["SplitModelUsage init."]
            gm.append("Expecting a TreePartitionsModelInfo instance.")
            raise Glitch, gm
        self.nParts = TPMI.nParts
        self.parts = []
        for i in range(self.nParts):
            SMTP = SplitModelUsagePart()
            SMTP.compCounts = [0] * TPMI.parts[i].nComps
            SMTP.rMatrixCounts = [0] * TPMI.parts[i].nRMatrices
            SMTP.gdasrvCounts = [0] * TPMI.parts[i].nGdasrvs
            self.parts.append(SMTP)

    def dump(self):
        print "                modelUsage nParts=%s" % self.nParts
        for pNum in range(self.nParts):
            print "                    part %s" % pNum
            print "%35s = %s" % ('compCounts', self.parts[pNum].comp
Counts)
            print "%35s = %s" % ('rMatrixCounts', self.parts[pNum].rM
atrixCounts)
            print "%35s = %s" % ('gdasrvCounts', self.parts[pNum].gdasrvCounts)

    def zero(self):
        for p in self.parts:
            for i in range(len(p.compCounts)):
                p.compCounts[i] = 0
            for i in range(len(p.rMatrixCounts)):
                p.rMatrixCounts[i] = 0
            for i in range(len(p.gdasrvCounts)):
                p.gdasrvCounts[i] = 0
                


class SplitModelUsagePart(object):
    def __init__(self):
        self.compCounts = None
        self.rMatrixCounts = None
        self.gdasrvCounts = None


class Split(object):
    #def __init__(self, nParts, nModels):
    def __init__(self):
        # self.set is a 1-based set of integers representing taxa on
        # one side of the split.  So naturally the taxa have to be
        # ordered.  It is 1-based because when I make a consensus, the
        # first thing I do is make a comb tree, and the terminal node
        # numbers in that tree and the resulting cons tree will
        # correspond to these 1-based set numbers.
        self.set = None
        self.key = None
        self.string = None
        self.count = 0.0
        self.rootCount = 0.0
        # rootCount2 is needed only for non-biRoot trees that are
        # rooted on the parent of the first taxon.
        self.rootCount2 = 0.0  
        self.proportion = 0.0
        self.cumBrLen = 0.0
        self.biRootCumBrLen = 0.0
        self.modelUsage = None
        self.biRootModelUsage = None
        self.rootModelUsage = None

    def dump(self):
        print "    Split dump:"
        print "%25s = %s" % ('set', self.set)
        print "%25s = %s" % ('key', self.key)
        print "%25s = %s" % ('string', self.string)
        print "%25s = %s" % ('count', self.count)
        print "%25s = %s" % ('rootCount', self.rootCount)
        print "%25s = %s" % ('proportion', self.proportion)
        print "%25s = %s" % ('cumBrLen', self.cumBrLen)
        if self.modelUsage:
            self.modelUsage.dump()
        if self.rootModelUsage:
            print "                rootModelUsage.dump()"
            #self.rootModelUsage.dump()
            for pNum in range(self.rootModelUsage.nParts):
                print "                    part %s" % pNum
                print "%35s = %s" % ('compCounts', self.rootModelUsage.parts[pNum].compCounts)

    def combineWith(self, otherSplitObject):
        # no need to do proportion -- that is handled by _finishSplits()
        # model stuff is ignored -- actually, wiped! -- due to lazy programming.
        oso = otherSplitObject
        assert self.key == oso.key
        self.count += oso.count
        self.rootCount += oso.rootCount
        self.rootCount2 += oso.rootCount2
        self.cumBrLen += oso.cumBrLen
        self.biRootCumBrLen += oso.biRootCumBrLen
        self.modelUsage = None
        self.biRootModelUsage = None
        self.rootModelUsage = None
        



class TreePartitionsModelInfo(object):
    def __init__(self):
        self.nParts = 0
        self.parts = []  # TreePartitionsModelInfoPart objects

    def dump(self):
        print "TreePartitionsModelInfo.dump()  nParts=%s" % self.nParts
        for i in range(self.nParts):
            self.parts[i].dump()

    def check(self):
        gm = ['TreePartitionsModelInfo.check()']
        if self.nParts < 1:
            gm.append("No parts.")
            raise Glitch, gm
        isHetero = 0
        for pNum in range(self.nParts):
            p = self.parts[pNum]
            if p.partNum != pNum:
                self.dump()
                gm.append("partNum is wrong.")
                raise Glitch, gm
            if p.nComps < 1:
                gm.append("Incomplete model info.  Part %s, no nComps." % pNum)
                raise Glitch, gm
            elif p.nComps > 1:
                isHetero = 1

            if p.nRMatrices < 1:
                gm.append("Incomplete model info.  Part %s, no nRMatrices." % pNum)
                raise Glitch, gm
            elif p.nRMatrices > 1:
                isHetero = 1

            if p.nGdasrvs < 0:
                gm.append("Incomplete model info.  Part %s, nGdasrvs not specified." % pNum)
                raise Glitch, gm
            elif p.nGdasrvs > 1:
                isHetero = 1
        if isHetero:
            return 2
        else:
            return 1



class TreePartitionsModelInfoPart(object):
    def __init__(self):
        self.partNum = -1
        self.nComps = 0
        self.nRMatrices = 0
        self.nGdasrvs = -1

    def dump(self):
        #print "    TreePartitionsModelInfoPart.dump  partNum=%s" % self.partNum
        print "    partNum=%s" % self.partNum
        print "        nComps     = %s" % self.nComps
        print "        nRMatrices = %s" % self.nRMatrices
        print "        nGdasrvs   = %s" % self.nGdasrvs



def _getModelInfo(theComment):

    # Eg [&&p4 models p1 c0.2 r0.2 g0.1]
    # or [&&p4 models p2 c0.2 r0.2 g0.1 c1.1 r1.1 g1.1]
    gm = ['TreePartitions._getModelInfo()']
    #print "_getModelInfo(), from theComment = %s" % theComment
    flob = cStringIO.StringIO(theComment)
    flob.seek(1, 0) # past the [

    tok = NexusToken.nextTok(flob)
    if not tok or tok != '&&p4':
        print "a got tok=%s" % tok
        flob.close()
        return  # not an error, just the wrong kind of comment
    tok = NexusToken.nextTok(flob)
    if not tok or tok != 'models':
        print "b got tok=%s" % tok
        flob.close()
        gm.append("Expecting 'models'.  Got %s" % tok)
        raise Glitch, gm
    tok = NexusToken.nextTok(flob)
    if tok[0] != 'p':
        gm.append("Expecting 'pN'.  Got %s" % tok)
        raise Glitch, gm
    try:
        nParts = int(tok[1:])

        if nParts > 0:
            theModelInfo = TreePartitionsModelInfo()
            theModelInfo.nParts = nParts
            for i in range(nParts):
                mp = TreePartitionsModelInfoPart()
                mp.partNum = i
                theModelInfo.parts.append(mp)
        else:
            theModelInfo = None
    except ValueError:
        