## set of supporting methods for parsers and converters

from collections import defaultdict
import networkx as nx
import itertools
import multiprocessing as mp

def split_to_layers(input_network):

    layer_info = defaultdict(list)
    subgraph_dictionary = {}
    
    for node in input_network.nodes(data=True):
        try:
            layer_info[node[0][1]].append(node[0])
        except Exception as err:
            layer_info[node[1]['type']].append(node[0])
    
    for layer,nodes in layer_info.items():
        subnetwork = input_network.subgraph(nodes)
        subgraph_dictionary[layer] = subnetwork #nx.relabel_nodes(subnetwork,mapping)
    del layer_info
    
    return subgraph_dictionary

def add_mpx_edges(input_network):
    
    _layerwise_nodes = split_to_layers(input_network)
    
    min_node_layer = {}
    for layer,network in _layerwise_nodes.items():
        min_node_layer[layer] = set([n[0][0] for n in network.nodes(data=True)])
        
    for pair in itertools.combinations(list(min_node_layer.keys()),2):
        layer_first = pair[0]
        layer_second = pair[1]
        pair_intersection = set.intersection(min_node_layer[layer_first],min_node_layer[layer_second])

        for node in pair_intersection:
            n1 = (node,layer_first)
            n2 = (node,layer_second)
            input_network.add_edge(n1,n2,key="mpx",type="mpx")

    return input_network
        

def parse_gaf_to_uniprot_GO(gaf_mappings,filter_terms=None):
    uniGO = defaultdict(list)    
    with open(gaf_mappings) as im:
        for line in im:
            parts = line.split("\t")
            try:
                if "GO:" in parts[4]:
                    uniGO[parts[1]].append(parts[4]) ## GO and ref both added
                if "GO:" in parts[3]:
                    uniGO[parts[1]].append(parts[3])
            except:
                pass

    all_terms = list(itertools.chain(*uniGO.values()))
    if filter_terms is not None:
        sorted_d = sorted(Counter(all_terms).items(), key=operator.itemgetter(1),reverse=True)
        top_100 = [x[0] for x in sorted_d[0:filter_terms]]
        new_map = defaultdict(list)
        for k,v in uniGO.items():
            v = [x for x in v if x in top_100]
            new_map[k] = v        
        return new_map
    
    else:
        return uniGO
