#!/bin/python
#
# This retrieves OrthoDB tree for a given API version.
# In addition it may do some additional post-processing.
#
# get_orthodb_tree.py [-nuvbJrlRLN]
#
#  -i           input: tree (JSON from OrthoDB API or newick)
#                      overrides -u,-v and -b
#  -u           input: full URL where to retrieve tree in JSON format
#                      overrides -v and -b
#  -v           input: OrthoDB API version (default 'current')
#  -b           input: OrthoDB API base URL (default https://data.orthodb.org)
#
#  -J           dump full tree as JSON (by default it dumps a newick file)
#
#  -r           set root node
#  -k           flag - if set, keep all leaves
#  -R           print a list of root nodes
#  -L           print a list of leaves
#  -N           print a list of nodes
#  -C           print one line per node containing the node plus all children (no leaves)
#               CSV: <root>,<child1>,<child2>,...
#  -I           with -C, print only branches with at least one child
#  -o           with -C, print only root
#  -n           do not print a header
#  
import os
import sys
import argparse
import urllib.request
import json
import newick

def load_json( fjson ):
    '''
    Load JSON from file.
    File is allowed to have comment lines (starting with '#').
    '''
    jdata = None
    
    # first try file
    if fjson is not None:
        with open(fjson,'r') as fh:
            jstr=''
            for line in fh:
                if line.startswith('#'):
                    continue
                jstr += line
                
            jdata = json.loads(jstr)
            return jdata
    
    return jdata

def load_json_url(url):
    '''
    Load JSON from url.
    '''
    jdata = None
    
    try:
        data = urllib.request.urlopen(url)
        jdata = json.loads(data.readline())
        return jdata['data']
    except:
        jdata=None
    
    return jdata

def load_newick(fname):
    '''
    Load newick from file.
    File is allowed to have comment lines (starting with '#').
    '''
    rval = None
    
    if fname is not None:
        with open(fname,'r') as fh:
            nstr=''
            for line in fh:
                if line.startswith('#'):
                    continue
                nstr += line

            try:
                nw = newick.loads(nstr)
                if isinstance(nw,list):
                    rval = nw
                else:
                    rval = [ nw ]
            except:
                rval = None

    return rval

def _json2newick( jdata, nwstr ):
    keystr = jdata['key']
    if 'children' not in jdata:
        return keystr+nwstr;
    
    nwstr = jdata['key']+nwstr;

    delim=''
    nwchild=''
    for child in jdata['children']:
        nwadd = ''
        nwadd = _json2newick(child,nwadd)
        nwchild = nwadd+delim+nwchild
        delim = ','

    if len(nwchild)>0:
        nwstr = '('+nwchild+')'+nwstr

    return nwstr

def json2newick( jdata ):
    '''
    Convert JSON to newick format.
    '''
    if jdata is None:
        return None
    start=''
    rval = []
    if isinstance(jdata,list):        
        for J in jdata:
            start = ''
            rval += newick.loads( _json2newick( J, start ))
    else:
        start = ''
        rval += newick.loads( _json2newick( jdata, start ))

    return rval

def get_leaves( nwroot ):
    leaves = None
    allnodes = []
    root = nwroot
    if not isinstance(nwroot,list):
        root = [ nwroot ]

    for N in nwroot:
        leaves = nwroot[0].get_leaves()
        allnodes += [ x.name for x in leaves ]
    
    allnodes.sort()
    return allnodes

def get_nodes( nwroot, includeRoot ):
    '''
    retrieve all internal nodes, ignoring root and leafs
    '''
    root = nwroot
    if not isinstance(nwroot,list):
        root = [ nwroot ]

    for N in nwroot:
        if includeRoot:
            # all non-leaves plus root
            allnodes = [n.name for n in N.walk() if not n.is_leaf]
        else:
            # all internal (no leafs, no root)
            allnodes = [n.name for n in N.walk() if n.is_internal]
    
    allnodes.sort()
    return allnodes

def clean_odbver( odbver ):
    if odbver is None:
        return None
    odbver = odbver.replace('.','-')
    return odbver.replace('-0','')

def get_orthodb_version( options ):
    ver = None
    url = None
    if options.fullurl is None:
        if (options.baseurl is not None) and (options.odbver is not None):
            url = os.path.join( options.baseurl, options.odbver, 'orthodb_release_id' )
    else:
        url = options.fullurl
        url = url.replace('tree','orthodb_release_id')
    if url is None:
        if options.odbinp is not None:
            with open( options.odbinp, 'r' ) as fh:
                line = fh.readline()
                elems = line.split(':')
                if len(elems)>1:
                    ver = elems[1].strip()
    else:
        data = urllib.request.urlopen(url)
        ver = json.loads(data.readline())
    return ver
        
def build_treeurl( options ):
    treeurl = None
    if options.fullurl is None:
        if (options.baseurl is not None) and (options.odbver is not None):
            treeurl = os.path.join( options.baseurl, options.odbver, 'tree' )
    else:
        treeurl = options.fullurl
    return treeurl

def find_nw_root( tree, root ):
    # assumes tree is a list of newick Node
    if root is None:
        return tree
    rval = None
    for T in tree:
        for n in T.walk():
            if n.name == root:
                rval = [ n ]
                break
        if rval is not None:
            break
    
    return rval

def load_tree( options ):
    '''
    Loads tree data.
    1. from --input - if set and not found, return None
    2. from URL info

    Returns in newick format (not string).
    
    '''
    data = None
    msg = None
    source=None
    if options.odbinp is not None:
        source = options.odbinp
        # test first newick format
        data = load_newick( options.odbinp )
        if data is None:
            # test JSON and convert to newick
            jdata = load_json( options.odbinp )
            data = json2newick( jdata )
        if data is None:
            # error - given input file gave nothing
            msg='input file {}'.format(options.odbinp)
        
    else:
        # build tree URL
        url = build_treeurl( options )
        if url is None:
            msg='failed to construct tree URL using given input'
            return None
        
        source = url
        
        # load JSON from url
        jdata = load_json_url( url )
        if options.dumpjson:
            return jdata,source,msg
        else:
            data = json2newick( jdata )

        if data is None:
            msg='failed constructing tree from url {}'.format(url)

    if data is not None:
        # if options.listnodes then the original leafs should always be removed
        if options.rmleaves or options.listnodes:
            for D in data:
                # the following will set all leaf names to None
                D.remove_leaf_names()
                for N in D.walk():
                    # remove those with None name - note slice notation '[:]' which is required
                    N.descendants[:] = [ ch for ch in N.descendants if ch.name is not None ]

        # if options.listnodes AND rmleaves are set, remove the next level of leaf nodes
        if options.rmleaves and options.listnodes:
            for D in data:
                # the following will set all leaf names to None
                D.remove_leaf_names()
                for N in D.walk():
                    # remove those with None name - note slice notation '[:]' which is required
                    N.descendants[:] = [ ch for ch in N.descendants if ch.name is not None ]
            
            
    return data,source,msg

def process( tree, options ):
    if tree is None:
        return 1

    odbver = get_orthodb_version( options )
    if odbver is None:
        odbver = 'unknown'

    result = 'newick tree'
    if options.dumpjson:
        result = 'raw OrthoDB json'
    elif options.listroot:
        result = 'list of root nodes'
    elif options.listleaves:
        result = 'list of all leaf nodes'
        if not options.rmleaves:
            result = 'list of all leaves'
    elif options.listnodes:
        result = 'list of all internal nodes (including root)'
    elif options.listchildren:
        result = 'list of all nodes plus their children (CSV)'

    if not options.noheader:
        print('# OrthoDB version   : {}'.format(odbver))
        print('# Data              : {}'.format(result))
    
    if options.dumpjson:
        print(json.dumps( tree ))
        return 0
    
    # tree is now a list of newick Node (root node(s))
    newtree = find_nw_root( tree, options.root )

    if newtree is None:
        print('error: unable to find root {} in tree'.format(options.root),file=sys.stderr)
        return 255

    if not options.noheader:
        print('# Root level(s)     : {}'.format(','.join([T.name for T in newtree])))

    # loop over list and apply options
    for T in newtree:
        # handle list root option
        if options.listroot:
            # print just the root
            print(T.name)
            
        elif options.listleaves:
            # print all leaves
            print('\n'.join(L for L in T.get_leaf_names()))
            
        elif options.listnodes:
            # print all nodes
            print('\n'.join(N.name for N in T.walk()))
            
        elif options.listchildren:
            if options.onlyroot:
                cstr = ','.join(C.name for C in T.descendants)
                if len(cstr)==0:
                    print(T.name)
                else:
                    print(T.name+','+cstr)
            else:
                lines=[]
                for N in T.walk():
                    cstr = ','.join(C.name for C in N.descendants)
                    if (len(cstr) == 0):
                        if not options.rmleaves:
                            continue
                    else:
                        cstr = ','+cstr
                    if options.withchild and (',' not in cstr):
                        continue
                    lines.append(N.name+cstr)
                for L in lines[::-1]:
                    print(L)
        else:
            # print full tree as a newick string
            print(newick.dumps(T))

#-----------------------------------------------------------------------------
parser = argparse.ArgumentParser(description='Extract taxids in BRH file')

parser.add_argument('-i', '--input',      dest='odbinp', help='input: OrthoDB newick or JSON file, supersedes all other input options', required=False, default=None)
parser.add_argument('-u', '--fullurl',    dest='fullurl',help='input: OrthoDB API full URL, supersedes --odbver + --baseurl', required=False, default=None)
parser.add_argument('-v', '--odbver',     dest='odbver', help='input: OrthoDB API version, construct full URL together with --baseurl', required=False, default='current')
parser.add_argument('-b', '--baseurl',    dest='baseurl',help='input: OrthoDB API base URL', required=False, default='https://data.orthodb.org')

parser.add_argument('-J', '--dumpjson',   dest='dumpjson',help='dump raw JSON (does not support -k and -r or any of the --listX options)', action='store_true',required=False,default=False)

parser.add_argument('-r', '--setroot',    dest='root',      help='root node to extract (taxid)', required=False,default=None)
parser.add_argument('-k', '--rmleaves',   dest='rmleaves',  help='remove source leaves in tree', action='store_true',required=False,default=False)
parser.add_argument('-R', '--listroot',   dest='listroot',  help='extract all root nodes', action='store_true',required=False,default=False)
parser.add_argument('-L', '--listleaves', dest='listleaves',help='extract only leaves', action='store_true',required=False,default=False)
parser.add_argument('-N', '--listnodes',  dest='listnodes', help='extract non-leaf nodes', action='store_true',required=False,default=False)
parser.add_argument('-C', '--listchildren',dest='listchildren', help='list node plus its children as a CSV: <node>,<child1>,<child2>,...', action='store_true',required=False,default=False)
parser.add_argument('-I', '--only_with_children',dest='withchild', help='together with -C, list only nodes with atleast ne child', action='store_true',required=False,default=False)
parser.add_argument('-o', '--onlyroot',    dest='onlyroot',help='with -C -> print only root with children', action='store_true',required=False,default=False)
parser.add_argument('-n', '--noheader',    dest='noheader',help='do not print heder lines', action='store_true',required=False,default=False)
###parser.add_argument('-S', '--skiproot', dest='skiproot',help='in conjunction with -N, skip the root node', action='store_true',required=False,default=False)

# special argument for a fix...
# parser.add_argument( '-b', '--rebase', dest='rebase', help='change base _1 to _0', action='store_true', required=False, default=False)

options = parser.parse_args()
if options.odbinp is not None:
    options.odbver = None
    options.baseurl = None
    options.fullurl = None
else:
    if options.dumpjson and options.root is not None:
        print('error: option -J does not support --setroot',file=sys.stderr)
        sys.exit(255)
        
options.odbver = clean_odbver( options.odbver )

#--------------------------------------------------------
# load the tree based on given options
tree,source,msg = load_tree( options )

if tree is None:
    print('error: failed loading tree data -> {}'.format(msg),file=sys.stderr)
    sys.exit(255)

# process tree according to options
ec = process( tree, options )
sys.exit(ec)

