#!/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)
    subtree = read_tree(options.subtree)
    
    # find root of subtree in tree
    # create a map between name and Node
    # find startnode
    subroot = subtree[0].name
    startnode = None
    allnodes = {}
    for node in tree[0].walk():
        if node.name == subroot:
            startnode = node
        allnodes[node.name] = node

    startIsRoot = subroot == tree[0].name
    
    # fail if not found
    if startnode is None:
        print('root node of sub-tree not found in tree, name {}'.format(subroot))
        return
    
    # set of all node names in subtree already in tree
    allsubs = set()
    for node in subtree[0].walk('skip_root'):
        if node.name in allnodes:
            allsubs.add(node.name)

    # loop over all leaves (leaf == nw.Node())
    # first expand all leaves
    subnodes = subtree[0].get_leaves()
    for leaf in subnodes:
        if leaf.name in allnodes:
            desc = allnodes[leaf.name].descendants
            if len(desc)>0:
                for d in desc:
                    leaf.add_descendant(copy.deepcopy(d))
                    allsubs.add(d.name)

    # look at startnode, find all branches NOT in subtree and add them
    for node in startnode.descendants:
        if node.name not in allsubs:
            subtree[0].add_descendant(copy.deepcopy(node))
            allsubs.add(node.name)

    if startIsRoot:
        print('{};'.format(subtree[0].newick))
    else:
        parent = startnode.ancestor
        tree[0].prune([startnode])
        parent.add_descendant(subtree[0])
        print('{};'.format(tree[0].newick))

    
parser = argparse.ArgumentParser(description='Insert a sub-tree into a tree.')

parser.add_argument('-t', '--tree', dest='tree', help='tree in newick format', metavar='FILE', required=True)
parser.add_argument('-s', '--subtree', dest='subtree', help='sub-tree to be inserted, in newick format', metavar='FILE', required=True)

options = parser.parse_args()


process(options)
