#!/bin/python
'''
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 os
import sys
import argparse
import newick as nw
import copy

def find_node( start,name ):
    for node in start.walk():
        if node.name==name:
            return node
    return None

def get_labels(options,name,ext):
    
    fname = os.path.join( options.lbldir,'%s%s'%(name,ext) )
    lblset = set()
    with open(fname,'r') as fhnode:
        for elem in fhnode:
            elem = elem.strip().rstrip('\n')
            if len(elem)==0:
                continue
            if elem[0]=='#':
                continue
            # add label
            lblset.add(elem)

    return lblset

'''                
def process(options):

    tree=None
    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 = nw.loads(nwstr)
    
    if tree is None:
        return

    keymap={}
    for node in tree[0].walk():
'''        


def process(options):

    newtree = None
    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 = nw.loads(nwstr)
        if tree is None:
            return

        # make sure extension starts with '.'
        ext = options.lblext
        ext = ext if ext[0]=='.' else '.'+ext

        if options.rmleaves:
            newtree = tree[0]
            newtree.prune([l for l in newtree.walk() if l.is_leaf])
        else:
            # get all labels per node
            node2keys = {}
            for node in tree[0].walk():
                node2keys[node.name] = set()
                node2keys[node.name] = get_labels(options,node.name,ext)

            # walk through the tree again and add unique descendants
            newtree = copy.deepcopy(tree[0])
            newnode = newtree
            for node in tree[0].walk():
                # find node in newtree
                newnode = find_node( newtree, node.name )
                if newnode is None:
                    print('bug: node not found in new tree {}'.format(node.name))
                    continue
                
                # extract unique keys by removing all descendants
                uniqKeys = node2keys[node.name]
                for d in node.descendants:
                    uniqKeys = uniqKeys-node2keys[d.name]
                    
                # no unique keys - skip to next node
                if len(uniqKeys)==0:
                    continue
                
                # add unique keys - they are by construct leafs
                for u in uniqKeys:
                    newnode.add_descendant(nw.Node(u))
                    
        # print new tree
        print(newtree.newick)


parser = argparse.ArgumentParser(description='Either remove (-r) or expand leaf nodes. If expanding, the labels are obtained from deduced files (<lbldir>/<leaf>.<lblext>)')

parser.add_argument('-t', '--nwtree', dest='nwtree', help='tree in newick format', metavar='FILE', required=True)
parser.add_argument('-l', '--lbldir', dest='lbldir', help='directory with node label files', metavar='DIR', default="./todo", required=False)
parser.add_argument('-e', '--lblext', dest='lblext', help='node label file extension', metavar='string', default=".todo", required=False)
parser.add_argument('-r', '--rmleaves', dest='rmleaves', help='this will remove all leaves, ignoring all other options (except -t)', action='store_true',default=False, required=False)

options = parser.parse_args()


process(options)
