#!/usr/bin/python3
'''
This script generates a list of runs based on newick trees
1. Leaf node (N)
* BRH: N.todo -> N.brh
* run brhclus on this BRH file

2. Inner node (N)
* pair file: get_pairs.py -f N.todo -c C1.todo [C2.todo C3.todo...] -> N.pairs
* BRH: N.todo -> N.brh
* input to BRH clus: N.brh, N.pairs, C1.og_raw,C2.og_raw,...,
'''
import sys
import argparse

import newick

class RunData:
    def __init__(self,root):
        self.root = root
        
    def is_leaf(self):
        return self.root.is_leaf()

    def name(self):
        return self.root.name
    
    def cmd(self,options):
        cmdstr='{}'.format(self.root.name)
        if (options.rebase):
            elems = cmdstr.split('_')
            if len(elems)==2:
                count = int(elems[1])
                if count==0:
                    raise Exception('already rebased')
                count=count-1
                cmdstr='{}_{}'.format(elems[0],count)
                
        if options.nokids:
            return cmdstr
        
        delim=' '
        if options.docsv:
            delim=','

        for child in self.root.descendants:
            if child.name is not None:
                if (not options.leaves) and child.is_leaf:
                    continue
                cmdstr=cmdstr+delim+'{}'.format(child.name)

        # compatible with orthologer arguments
        if options.nthreads is not None and options.docsv:
            cmdstr = cmdstr+',{}'.format(options.nthreads)

        return cmdstr
    
allNodeData=[]
allLeafData=[]

class GenCmds:
    def __init__(self):
        self.allLeaves=[]
        self.allLastNodes=[]
        self.allNodes=[]
        self.rootNode=None
        
    def go(self,level,node,root,isRootDescendant,nolast,maxLevel):

        if node.name is None:
            return isRootDescendant

        # if not yet in root branch, check node if it is
        if not isRootDescendant:
            if root is not None:
                # root is specified
                if node.name==root:
                    # current node is root
                    self.rootNode = root
                    isRootDescendant = True
            else:
                isRootDescendant = True
                if self.rootNode is None:
                    self.rootNode = node.name

        if isRootDescendant:
            # if part of root tree -> save nodes/leaves
            # create separate list of leaves
            if node.is_leaf:
                self.allLeaves.append(RunData(node))
                return isRootDescendant

            # check if it is a last internal node before a leaf
            lastnode=True
            for child in node.descendants:
                if not child.is_leaf:
                    lastnode=False
                    break

            # if all kids are leaves, its the last node
            if lastnode:
                self.allLastNodes.insert(0,RunData(node))

            # add to all nodes
            skip = lastnode and nolast
            if not skip:
                self.allNodes.insert(0,RunData(node))
            
        # check descendants if root is not found
        descendants = []
        #if root is not None and not isRootDescendant:
        #    # root branch not yet found - search for it
        #    for child in node.descendants:
        #        if child.name==root:
        #            # found root!
        #            print('FOUND ROOT2 {}'.format(root))
        #            isRootDescendant = True
        #            self.rootNode = root
        #            descendants.append(child)
        #            break

        # if no descendants here, set to all available
        if len(descendants)==0:
            descendants = node.descendants

        # recursive parse of tree
        levelOK = ( maxLevel == 0 ) or ( level < maxLevel )
        if levelOK:
            if isRootDescendant:
                level = level+1
            for child in descendants:
                foundRoot = self.go(level,child,root,isRootDescendant,nolast,maxLevel)
                if not isRootDescendant and foundRoot:
                    isRootDescendant = True
                    break
            
        return isRootDescendant
    
def process(options):
    output = []
    outputLeaves = options.leaves
    outputNodes = options.nodes
    outputLastNodes = options.last
    if not options.leaves and not options.nodes and not outputLastNodes:
        outputLeaves = True
        outputNodes = True
    if options.root:
        outputLeaves = False
        outputNodes = False

    with open(options.nwtree,'r') as fh:
        nwstr = ''
        for line in fh:
            line = line.strip().rstrip('\n')
            if len(line)==0:
                continue
            if line[0]=='#':
                continue
            nwstr = nwstr+line
        tree = newick.loads(nwstr)
        if tree is None:
            return
        gen = GenCmds()
        gen.go(0,tree[0],options.setroot,False,options.nolast,options.maxLevel)

        if options.root:
            output.append(gen.rootNode)
        else:
            useName = options.level is not None
            if outputLeaves:
                for leaf in gen.allLeaves:
                    if (useName and (leaf.name() == options.level)) or not useName:
                        output.append(leaf.cmd(options))

            if outputNodes:
                for node in gen.allNodes:
                    if (useName and (node.name() == options.level)) or not useName:
                        output.append(node.cmd(options))

            if outputLastNodes:
                for node in gen.allLastNodes:
                    if (useName and (node.name() == options.level)) or not useName:
                        output.append(node.cmd(options))
    
    return output

class Options:
    def __init__(self):
        self.default()
        
    def default(self):
        self.nwtree = None
        self.nthreads = None
        self.docsv = False
        self.leaves = False
        self.nodes = False
        self.last = False
        self.nolast = False
        self.level = None
        self.nokids = False
        self.root = False
        self.setroot = None
        self.rebase = False
        self.maxLevel = 0

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate runs from a newick tree')

    parser.add_argument('-t', '--nwtree', dest='nwtree', help='tree in newick format', metavar='FILE', required=True)
    parser.add_argument('-m', '--maxlevel', dest='maxLevel', help='max level in tree, 0->full tree, 1-> children of root, etc', required=False, default=0)
    parser.add_argument('-c', '--csv', dest='docsv', help='output csv', action='store_true', required=False, default=False)
    parser.add_argument('-L', '--leaves', dest='leaves', help='output leaves', action='store_true', required=False, default=False)
    parser.add_argument('-N', '--nodes', dest='nodes', help='output nodes', action='store_true', required=False, default=False)
    parser.add_argument('-M', '--last', dest='last', help='output last nodes', action='store_true', required=False, default=False)
    parser.add_argument('-x', '--nolast', dest='nolast', help='exclude last nodes (with -N)', action='store_true', required=False, default=False)
    parser.add_argument('-l', '--level', dest='level', help='output only for this node', required=False, default=None)
    parser.add_argument('-k', '--nokids', dest='nokids', help='with --nodes exclude kids', action='store_true', required=False, default=False)
    parser.add_argument('-R', '--root', dest='root', help='extract the root', action='store_true', required=False, default=False)
    parser.add_argument('-u', '--unique', dest='unique', help='extract unique leaves for this root (-r) ; only with -L', action='store_true', required=False, default=False)
    parser.add_argument('-r', '--setroot', dest='setroot', help='set the root and output everything below', required=False, default=None)

    parser.add_argument('-n', '--nthreads', dest='nthreads', help='nr of threads to print (compatible with orthologer args)', required=False, default=None)
    parser.add_argument('-b', '--rebase', dest='rebase', help='change base _1 to _0', action='store_true', required=False, default=False)


    options = parser.parse_args()


    output = process(options)
    for N in output:
        print(N)
