#!/bin/python
'''
Propagate new leaf nodes upwards in a tree.

Input:
  srctree    - source tree in newick format
  nodes      - file containing node to todo-file map

The nodes file has two columns, space separated <node id> <todo file with labels> .
Only node id's that are in the original source tree are considered.

'''
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 read_todo( ftodo ):
    labels = None
    with open(ftodo,'r') as fh:
        labels = []
        for lbl in fh:
            if len(lbl)==0:
                continue
            if lbl[0]=='#':
                continue
            labels.append(lbl.strip())
    return labels
            
def read_nodes( fnodes ):
    nodes = None
    nmissing = 0
    ndup = 0
    nempty = 0
    with open(fnodes,'r') as fh:
        # read list of todo files
        nodes = None
        for line in fh:
            line = line.strip().rstrip('\n')
            if len(line)==0:
                continue
            if line[0]=='#':
                continue

            elems = line.split()
            ftodo = elems[1]
            node  = elems[0]
            
            # make sure todo file exists
            if not os.path.isfile(ftodo):
                nmissing += 1
            else:
                # todo file format: <path>/<node>.todo
                if nodes is None:
                    nodes = {}
                
                # warn duplicates
                if node in nodes:
                    ndup += 1

                # init for given node, add each label in file
                nodes[node] = read_todo( ftodo )
                if len(nodes[node]) == 0:
                    nempty += 1
                
                
    return nodes,nmissing,ndup,nempty

def process_tree( inptree, nodes, options ):
    tree = copy.deepcopy( inptree )


    leaves = tree[0].get_leaves()

    # get all leaf nodes (one up wrt the actual leaves)
    leafnodes = set()
    for L in leaves:
        leafnodes.add( L.ancestor )

    # loop over all leaf nodes and process
    for L in leafnodes:
        # clear descendants if keepsrc is false
        if not options.keepsrc:
            #newdesc = []
            L.descendants[:] = [ D for D in L.descendants if not D.is_leaf ]
            #for D in L.descendants:
            #    if not D.is_leaf:
            #        newdesc.append(D)
            #L.descendants = newdesc
                    
        #    L.descendants = [ N for N in L.descendants if not N.is_leaf ]

        if L.name in nodes:
            # leaf node in set of input nodes

            # loop over new leaves and add them to descendants
            for newleaf in nodes[L.name]:
                L.add_descendant( nw.Node(newleaf) )


    # remove duplicates
    for N in tree[0].walk():
        if N.is_leaf:
            continue
        uniq = set()
        common = set()
        for D in N.descendants:
            # first add leaves - set to be filtered
            if D.is_leaf:
                uniq.add(D.name)
            else:
                # it's not a leaf, add its descendants that are leaves
                # by construiction above it contains ALL leaves at this node
                for DD in D.descendants:
                    if DD.is_leaf:
                        common.add(DD.name)
        
        # remove the leaves of descendants, thus creating a list of unique leaves
        uniq = uniq - common
        newdes = []
        index = 0
        nleafstot=0
        nleafskeep=0
        nbef = len( N.descendants )
        N.descendants[:] = [D for D in N.descendants if (D.is_leaf and (D.name in uniq)) or (not D.is_leaf)]
        #for D in N.descendants:
        #    if D.is_leaf:
        #        nleafstot += 1
        #        if D.name not in uniq:
        #            del N.descendants[index]
        #        else:
        #            nleafskeep += 1
        #    index += 1
        if (len(N.descendants) == 0):
            print('NODE:{} tot={} keep={}, bef = {}, after={}'.format(N.name,nleafstot,nleafskeep,nbef,len(N.descendants)))
    return tree

def process(options):
    '''
    Process input tree and nodes.
    Add new leaves to each leaf node and possiby clear old leaves.
    '''
    # read input trees
    srctree = read_tree( options.srctree )
    if srctree is None:
        print("error: failed reading source tree {}".format(options.srctree))
        return 255

    print(f"# tree source: {options.srctree}")
    # returns a dict where key is node and value is a set of leaf names
    nodes,nmissing,ndup,nempty = read_nodes( options.nodes )
    if nodes is None:
        print("error: failed reading nodes {}".format(options.nodes))
        return 255
    print(f"# read {len(nodes)} nodes from {options.srctree}")
    
    # process the tree
    newtree = process_tree( srctree, nodes, options )

    if options.output is None:
        nw.dumps( newtree )
    else:
        print("# writing updated tree to file {}".format(options.output))
        nw.write( newtree, options.output )
    
    return 0
    
parser = argparse.ArgumentParser(description='Propagate new leaves upwards a tree.')

parser.add_argument('-t', '--srctree', dest='srctree', help='source tree in newick format', metavar='FILE', required=True)
parser.add_argument('-n', '--nodes', dest='nodes', help='file containing nodes', metavar='FILE', required=True)
parser.add_argument('-o', '--output', dest='output', help='output tree in newick format', metavar='FILE', default=None, required=False)
parser.add_argument('-k', '--keep', dest='keepsrc', help='keep source tree nodes',  action='store_true', default=False, required=False )

###parser.add_argument('-d', '--dest', dest='dest', help='destination directory', metavar='DIR', required=False)

options = parser.parse_args()


process(options)
