#!/bin/python
'''
Insert a sub-tree into a given tree
'''
import os
import sys
import argparse
import newick as nw
import copy

def read_tree(ftree):
    tree = None
    with open(ftree,'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)
    return tree

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

def ascend_branch_to_next( node, processed, allnodes ):
    if node is None:
        return node,node,False
    if node.name=="":
        return node,node,False
    if node.ancestor is None:
        return node,node,False
    
    nextnode = node.ancestor
    lastnode = nextnode
    update = False
    while nextnode.name not in allnodes:
        print('node {}'.format(nextnode.name))
        # if node has already been processed, no need to continue with branch
        if nextnode.name in processed:
            update = False
            break
        processed.add(nextnode.name)
        lastnode = nextnode
        nextnode = nextnode.ancestor
        update = True
        if nextnode is None:
            break
    
    updnode = None
    if update:
        if nextnode is None:
            print('next is None! prev {}'.format(lastnode.name))
        else:
            updnode = allnodes[nextnode.name]
    return nextnode,updnode,copy.deepcopy(lastnode),update

def process(options):
    # read input trees
    tree = read_tree(options.tree)
    allnodes = {}
    print_node = options.node is None
    for node in tree[0].walk():
        # skip leafs
        if node.is_leaf:
            continue

        # check if node is to be printed
        if not print_node:
            if options.node is not None:
                if options.node != node.name:
                    continue
                print_node = True

        if print_node:
            print(node.name)
    
parser = argparse.ArgumentParser(description='Will print all sub-nodes (including root) starting from a given node.')

parser.add_argument('-t', '--tree', dest='tree', help='tree in newick format', metavar='FILE', required=True)
parser.add_argument('-n', '--node', dest='node', help='root node', metavar='string', default = None, required=False)

options = parser.parse_args()


process(options)
