#!/usr/bin/python3

import argparse
import xml.etree.ElementTree as ET

def tag( label ):
    namespace = '{http://orthoXML.org/2011/}'
    return '{}{}'.format(namespace,label)

def geneRefs( groups ):
    if groups is None:
        return None
    if type(groups) != list:
        listGroups = [ groups ]
    else:
        listGroups = groups

    rval = []
    for group in listGroups:
        gall = group.findall(tag('geneRef'))
        if gall is None:
            continue
        for G in gall:
            rval.append(int(G.attrib['id']))
    return rval

def dumpGenes( clusterId, grefs, geneMap ):
    if grefs is None:
        return None
    for g in grefs:
        geneId = geneMap.get(g)
        print('{} {}'.format(clusterId,geneId))


#------------------------------------------------------------------------------
parser = argparse.ArgumentParser(description='Simple plot tool')

parser.add_argument('-f', '--file', dest='finput',
                    help='OrthoXML file', metavar='FILE', required=True)
options = parser.parse_args()

tree = ET.parse(options.finput)
if tree is None:
    print('error: failed reading tree from input file {}'.format(options.finput))
    sys.exit(255)


root = tree.getroot()
groups = root.findall(tag('groups'))
if groups is None:
    print('error: cannot find tag \'groups\'')
    sys.exit(255)
    
species = root.findall(tag('species'))
if species is None:
    print('error: cannot find tag \'species\'')
    sys.exit(255)

# first read the full gene map
geneMap = {}
for S in species:
    DB = S.find(tag('database'))
    genes = DB.find(tag('genes'))
    for G in genes:
        numId = int(G.attrib['id'])
        protId = G.attrib['protId']
        if numId in geneMap:
            print('duplicate id {}'.format(numid))
        geneMap[numId] = protId

print('# found {} genes'.format(len(geneMap)))
print('# found {} group tags'.format(len(groups)))

# there should be ony one groups but, not sure
clusterId = 0
for G in groups:
    mainOrthologs = G.findall(tag('orthologGroup'))
    for og in mainOrthologs:
        paralogs  = geneRefs( og.findall(tag('paralogGroup')) )
        orthologs = geneRefs( og.findall(tag('orthologGroup')) )
        genes     = geneRefs( og )
        dumpGenes(clusterId, paralogs,  geneMap)
        dumpGenes(clusterId, orthologs, geneMap)
        dumpGenes(clusterId, genes,     geneMap)
        clusterId += 1

    
