#!/usr/bin/python3
'''
parse orthoxml file
'''

import os
import argparse
import sys

from collections import OrderedDict

import orthoxml

numid2genid = {}
def fillGeneMap( species ):
    global numid2genid
    # species is a list of orthoxml.species class
    for sp in species:
        for db in sp.database:
            for gene in db.genes.gene:
                # gene is a orthoxml.gene class
                numid = gene.id
                genid = gene.protId
                if genid is None:
                    genid = gene.geneId
                if genid is None:
                    genid = gene.transcriptId
                if numid in numid2genid:
                    print('warning: id {} found more than once! Skip!'.format(numid))
                    continue
                if genid is None:
                    numid2genid[numid] = 'UNKNOWN:{}'.format(numid)
                else:
                    numid2genid[numid] = genid


def getId( numid ):
    global numid2genid
    if not isinstance(numid,int):
        numid = int(numid)
    if numid in numid2genid:
        genid = numid2genid[numid]
    else:
        genid = 'NOTFOUND:{}'.format(numid)
    return genid

def dumpPairs(G,geneRefs):
    global numid2genid
    isLeaf = len(G.orthologGroup)==0
    Ngenes = len(G.geneRef)
    if geneRefs is None:
        geneRefs = []
        
    if isLeaf:
        Nsrc = len(geneRefs)
        print('#ID {}'.format(G.id))
        # first print previous geneRefs in pair with all in the leaf
        for i in range(0,Nsrc):
            idA = getId(geneRefs[i].get_id())
            for j in range(0,Ngenes):
                idB = getId(G.geneRef[j].get_id())
                print('{} {} N'.format(idA,idB))

        # print all pairs in the leaf
        for i in range(0,Ngenes):
            idA = getId(G.geneRef[i].get_id())
            for j in range(i+1,Ngenes):
                idB = getId(G.geneRef[j].get_id())
                print('{} {} L'.format(idA,idB))
        return True

    geneRefs += G.geneRef
    for og in G.orthologGroup:
        dumpPairs(og,geneRefs)
    
    return True

#-----------------------------------------------------------------
# MAIN


parser = argparse.ArgumentParser(description='generate pairs from orthoxml file')

parser.add_argument('-f', '--file',    dest='fname',     help='input cluster file',         metavar='FILE',   default = None, required=True)

options = parser.parse_args()

rootobj = orthoxml.parse(options.fname)

fillGeneMap( rootobj.species)

indent=' '
for G in rootobj.groups.orthologGroup:
    dumpPairs(G,None)

    
