__author__ = 'mjohnpayne'

# 1 - generate dict of clusters with clusterid as key and gene ids as data
# 2 - generate dict of interpro ID:description objects with gene id as key and id-descriptions as data
# 2 - generate dict of GOTermobjects with gene id as key and id-descriptions as data
# 3 - merge by having clusterid link to collapsed list of interpro ids (set)
# 4 - print out

import sys
from time import sleep as sl

clusterinf = open('/Volumes/MP_HD/Pm_Ts_Tf_Pf_comparison/orthomcl_run/orthomcl_data/orthomcl_ortho_groups.txt','r')

gafs = ['/Volumes/MP_HD/Blast2Go/Pm/exported_data/pm_b2g.gaf.txt','/Volumes/MP_HD/Blast2Go/Tf/exported data/blast2go_annotgaf_20141204_2159.txt','/Volumes/MP_HD/Blast2Go/Pf/exported_data/blast2go_annotgaf_20141204_2344.txt','/Volumes/MP_HD/Blast2Go/Ts/exported data/blast2go_annotgaf_20141204_2348.txt']


def rename(gene):
    if 'TF' in gene:
        no = 5-len(gene[3:])
        gene = 'TFLA_' + no*'0' + gene[3:] + '0'
    elif 'Pf' in gene:
        no = 5-len(gene[3:])
        gene = 'PFUN_' + no*'0' + gene[3:] + '0'
    return gene


def ortho_clust_to_dict(inf):
    clust = {}
    for line in inf:
        col = line.strip('\n').split(' ')
        clust[col[0][:-1]] = [rename(x[4:]) for x in col[1:]]
    return clust



def desc_to_dict(intproin):
    gos = {}
    for f in intproin:
        tmp = open(f,'r').readlines()
        for line in tmp[4:]:
            col = line.strip('\r\n').split('\t')
            gene = col[1]#rename(col[0][4:])
            if gene not in gos:
                gos[gene] = ['\t'.join(col[4:])]
            else:
                gos[gene].append('\t'.join(col[4:]))
    ogos = {}
    for i in gos:
        ogos[i] = list(set(gos[i]))
    return ogos

def merge(orthos,gos):
    annots = {}
    for i in orthos:
        go = []
        for j in orthos[i]:
            if j in gos:
                go += gos[j]
        gol = list(set(go))
        annots[i] = gol
    return annots



def main(cl,int):
    clu = ortho_clust_to_dict(cl)
    gos = desc_to_dict(int)
    annots = merge(clu,gos)
    return annots

annotations = main(clusterinf,gafs)



outfile = open('/Volumes/MP_HD/Pm_Ts_Tf_Pf_comparison/orthomcl_run/ortho_groups_by_species/4_species_ortho_group_b2g_annots.gaf','w')

outfile.write('!Generated by Blast2GO\n!Note: This is a beta function of Blast2GO\n!Date: 20141204\n!gaf-version: 2.0\n')

for i in annotations:
    for j in annotations[i]:
        outfile.write('B2G\t' + i + '\t' + i + '\t\t' + j + '\n')

outfile.close()