#!/bin/python3

import os
import argparse
import sys
from ete3 import NCBITaxa

def clean_taxid(T):
    elems = T.split('_')
    if len(elems) == 2:
        return elems
    return elem[0],None

def process( orthodb_nodes, taxids, outdir ):
    '''
    orthodb_nodes : a set of nodes defined in OrthoDB
    taxids        : a set of OrthoDB taxids

    Each taxid is of the form <ncbi taxid>_<version> where the NCBI taxid version
    are both numerical.
    '''
    
    # instance of NCBI query engine
    ncbi = NCBITaxa()

    # create todo DB - maps OrthoDB to set of full taxids
    todoDB = {}
    todoDBleaf = {}
    print('adding {} OrthoDB nodes'.format(len(orthodb_nodes)))
    for node in orthodb_nodes:
        #print('add OrthoDB node: {}'.format(node))
        todoDB[node] = set()
        todoDBleaf[node] = set()

    print('get lineages for {} taxids'.format(len(taxids)))
    # loop over all given taxids
    for T in taxids:
        # split taxid into NCBI id and version
        cleanT,versionT = clean_taxid(T)

        # get lineage - save full taxid only in the OrthoDB nodes
        lineage = ncbi.get_lineage(cleanT)
        Nend = None
        for node in lineage:
            N = int(node)
            if N in orthodb_nodes:
                Nend = N
                todoDB[N].add(T)

        # Nend now contains the last node found in OrthoDB
        if Nend is not None:
            todoDBleaf[Nend].add(T)

    # output full todo files
    for item in todoDB.items():
        node = item[0]
        todo_file = os.path.join( outdir, 'map{}.todo'.format(item[0]) )
        if os.path.isfile(todo_file):
            print('file exists - skip {}'.format(todo_file))
            continue
        if len(item[1]) == 0:
            print('node is empty - skip file {}'.format(todo_file))
            continue
        
        print('writing file {}'.format(todo_file))
        with open( todo_file, 'w' ) as fh:
            sT = sorted(item[1])
            for taxid in sT:
                fh.write('{}\n'.format(taxid))

    # output leaf todo files
    for item in todoDBleaf.items():
        node = item[0]
        todo_file = os.path.join( outdir, 'leaf{}.todo'.format(item[0]) )
        if os.path.isfile(todo_file):
            print('file exists - skip {}'.format(todo_file))
            continue
        if len(item[1]) == 0:
            print('node is empty - skip file {}'.format(todo_file))
            continue
        
        print('writing file {}'.format(todo_file))
        with open( todo_file, 'w' ) as fh:
            sT = sorted(item[1])
            for taxid in sT:
                fh.write('{}\n'.format(taxid))

def load_taxids( options ):
    taxids = set()
    with open(options.taxids,'r') as fh:
        for line in fh:
            line = line.strip().rstrip('\n')
            if len(line)==0:
                continue
            if line[0]=='#':
                continue
            taxids.add(line)
            
    if len(taxids) == 0:
        print('warning: no taxids found in file {}'.format(options.taxids))
        return None
    
    return taxids

def load_odbnodes( options ):
    odbnodes = set()
    with open( options.odbnodes ) as fh:
        for line in fh:
            line = line.strip().rstrip('\n')
            if len(line)==0:
                continue
            if line[0]=='#':
                continue
            odbnodes.add(int(line))
    
    if len(odbnodes) == 0:
        print('warning: no OrthoDB node ids found in file {}'.format(options.odbnodes))
        return None

    return odbnodes
    
#----------------------------------------------------------------------------
if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Create todo files per node using NCBI lineage')

    parser.add_argument('-v', '--odbver',    dest='odbver',   help='OrthoDB version - used if --odb is not given', metavar='ODBVER',  required=False)
    parser.add_argument('-d', '--odb',    dest='odbnodes',   help='file with valid orthodb nodes, overrides --odbver', metavar='FILE',  required=False)
    parser.add_argument('-t', '--taxids', dest='taxids',     help='file with taxids',              metavar='FILE',  required=True)
    parser.add_argument('-o', '--outdir', dest='outdir',     help='output directory',              metavar='FILE',  required=True)

    options = parser.parse_args()

    # set of OrthoDB species taxids as strings e.g '9606_0'
    taxids = load_taxids( options )
    if taxids is None:
        sys.exit(1)

    # set OrthoDB nodes (levels) as NCBI taxids (integers)
    odbnodes = load_odbnodes( options )
    if odbnodes is None:
        sys.exit(1)
    
    process( odbnodes, taxids, options.outdir )
