#!/bin/python

import os
import sys
import argparse
import urllib.request
import json
import newick

def load_json(url,fjson):
    jdata = None
    
    # first try file
    if fjson is not None:
        with open(fjson,'r') as fh:
            jdata = json.load(fh)

    # not found... try URL
    if jdata is None:
        try:
            data = urllib.request.urlopen(url)
            jdata = json.loads(data.readline())
        except:
            jdata=None
    
    return jdata


def find_root( node, root ):
    if root is None:
        return node
    
    if node is None:
        return None

    if str(node['key'])==str(root):
        return node

    if 'children' not in node:
        return None

    for child in node['children']:
        found = find_root( child, root )
        if found is not None:
            return found

def json2newick( jdata, nwstr, doleaves ):
    keystr = jdata['key']
    if 'children' not in jdata:
        if doleaves:
            if (options.rebase):
                elems = keystr.split('_')
                if len(elems)==2:
                    count = int(elems[1])
                    if count==0:
                        raise Exception('already rebased')
                    count=count-1
                    keystr='{}_{}'.format(elems[0],count)

            nwstr = keystr+nwstr;
            
        return nwstr
    
    nwstr = jdata['key']+nwstr;

    delim=''
    nwchild=''
    for child in jdata['children']:
        if not doleaves:
            if 'children' not in child:
                continue

        nwadd = ''
        nwadd = json2newick(child,nwadd,doleaves)
        nwchild = nwadd+delim+nwchild
        delim = ','

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

    return nwstr

def getLeaves( nwroot ):
    leaves = nwroot[0].get_leaves()
    allnodes = [ x.name for x in leaves ]
    allnodes.sort()
    return '\n'.join(allnodes)

def getNodes( nwroot, includeRoot ):
    '''
    retrieve all internal nodes, ignoring root and leafs
    '''
    if includeRoot:
        # all non-leaves plus root
        allnodes = [n.name for n in nwroot[0].walk() if not n.is_leaf]
    else:
        # all internal (no leafs, no root)
        allnodes = [n.name for n in nwroot[0].walk() if n.is_internal]
    
    allnodes.sort()
    return '\n'.join(allnodes)

    
#jroot = json.loads('{"key":"7147","children":[{"key":"L11"},{"key":"L12"},{"key":"13","children":[{"key":"21","children":[{"key":"L31"}]},{"key":"L22"}]}]}')
#nw=''
#nw = json2newick(jroot,nw,True)
#print(nw)
#sys.exit(0)

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

parser.add_argument('-v', '--odbver', dest='odbver',help='OrthoDB API version', required=False, default='current')
parser.add_argument('-u', '--baseurl', dest='baseurl',help='OrthoDB API base URL', required=False, default='https://data.orthodb.org')
parser.add_argument('-j', '--json', dest='odbjson',help='OrthoDB JSON, previously downloaded using --odburl', required=False, default=None)
parser.add_argument('-r', '--root', dest='root',help='root node to extract (taxid)', required=False,default=None)
parser.add_argument('-L', '--leaves', dest='doleaves',help='add leaves to tree', action='store_true',required=False,default=False)
parser.add_argument('-R', '--listroot', dest='listroot',help='list all root nodes', action='store_true',required=False,default=False)
parser.add_argument('-T', '--listleaves', dest='listleaves',help='extract only leaves', action='store_true',required=False,default=False)
parser.add_argument('-N', '--listnodes', dest='listnodes',help='extract only internal nodes', 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)

parser.add_argument( '-b', '--rebase', dest='rebase', help='change base _1 to _0', action='store_true', required=False, default=False)

options = parser.parse_args()

treeurl = os.path.join( options.baseurl, options.odbver, 'tree' )

dotree = True
if options.listleaves or options.listnodes:
    dotree = False

# retrieve OrthoDB tree
jdata = load_json(treeurl,options.odbjson)
if jdata is None:
    print('error: unable to retrieve OrthoDB tree',file=sys.stderr)
    if treeurl is not None:
        print('error: URL  - {}'.format(treeurl), file=sys.stderr)
    if options.odbjson is not None:
        print('error: JSON - {}'.format(options.odbjson),file=sys.stderr)
    sys.exit(255)


# process input
rval=''

if options.listroot:
    print('%-40s %s'%('#Name','Taxid'))
    for jstart in jdata['data']:
        print('%-40s %s'%(jstart['name'],jstart['key']))

#elif options.gennodes is not None:
    
else:
    if options.root is None:
        # if no root is given, dump the full JSON
        rval = json.dumps(jdata)
    else:
        # extract the newick tree for the given root
        for jstart in jdata['data']:
            jroot = find_root( jstart, options.root )
            if jroot is not None:
                print('# URL    : {}'.format(treeurl))
                print('# parent : {}'.format(jstart['name']))
                print('# name   : {}'.format(jroot['name']))
                print('# taxid  : {}'.format(jroot['key']))
                rval=''
                nwroot=None
                if options.listleaves:
                    nwstr = json2newick(jroot,rval,True)
                    nwroot = newick.loads(nwstr)
                    rval = getLeaves(nwroot)
                if options.listnodes:
                    rval=''
                    nwstr = json2newick(jroot,rval,True)
                    nwroot = newick.loads(nwstr)
                    # get all nodes excluding root
                    rval = getNodes(nwroot,not options.skiproot)
                if dotree:
                    rval = json2newick(jroot,rval, options.doleaves)
                    rval += ';'
                break

if len(rval)>0:
    print(rval)
else:
    print('error: nothing found!',file=sys.stderr)
    sys.exit(255)
