#! /usr/bin/env python3
# vRhyme
# Author: Kristopher Kieft
# University of Wisconsin-Madison

import subprocess
import pickle
import sys
import networkx as nx
from networkx.algorithms.community import label_propagation_communities


def edge_check(x,y,max_edges):
    '''
    Check if one of the nodes passed max edges threshold
    '''
    check = False
    if x <= max_edges or y <= max_edges:
        check = True

    return check

def network_stuff(net_data, folder, bin_size, refine_m, max_edges, iteration):
    '''
    Main wrapper for generating clusters/networks of bins.
    Clusters are generated by connected nodes.
    Automatic refinement with KMeans clustering.
    '''
    if len(net_data) == 0:
        return
    
    net_data.sort(key = lambda x: x[2])

    counter = {}
    data = []
    bins = {}
    duos = {}
    b = 1

    for item in net_data:
        qry = item[0]
        sbj = item[1]

        try:
            q = counter[qry]
            q += 1
            counter[qry] = q
        except Exception:
            counter[qry] = 1
            q = 1
        try:
            s = counter[sbj]
            s += 1
            counter[sbj] = s
        except Exception:
            counter[sbj] = 1
            s = 1

        check = edge_check(q,s,max_edges)
        if check == True:
            added = False
            for key,value in bins.items():
                if qry in value or sbj in value:
                    value.update([qry,sbj])
                    bins[key] = value
                    duos[key].append((qry,sbj))
                    added = True
                    break

            if added == False:
                bins[b] = set([qry,sbj])
                duos[b] = [(qry,sbj)]
                b += 1

    counter = None
    net_data = None

    while True:
        names = list(bins.keys())
        l = len(names)
        alter = False
        for i in range(0,l-1):
            for j in range(i+1,l):
                x = names[i]
                y = names[j]

                try:
                    x_values = bins[x]
                    y_values = bins[y]

                    if x_values & y_values != set():
                        x_values = x_values.union(y_values)
                        bins[x] = x_values
                        duos[x].extend(duos[y])

                        del bins[y]
                        del duos[y]
                        alter = True

                except Exception:
                    # y was deleted
                    pass

        if alter == False: # nothing was combined
            break

    names = None
    final_bins = {}
    b = 1
    for key,data in duos.items():
        names = bins[key]
        ln = len(names)

        if ln >= refine_m:
            G = nx.Graph(data)
            comm = label_propagation_communities(G)

            for new_bin in comm:
                if len(new_bin) >= bin_size:
                    final_bins[b] = set(new_bin)
                    b += 1
        else:
            if ln >= bin_size:
                final_bins[b] = names
                b += 1

    with open(folder + iteration + ".vRhyme-bins.tsv", "w") as net_out:
        s = 0
        for val in final_bins.values():
            names = list(val)
            s += len(names)

            names = [str(n) for n in names]
            names = '\t'.join(names)
            net_out.write(f'{names}\n')

    b = len(final_bins)
    bins = None

if __name__ == '__main__':
    folder = sys.argv[2]
    bin_size = int(sys.argv[3])
    refine_m = int(sys.argv[4])
    max_edges = int(sys.argv[5])
    iteration = sys.argv[6]

    with open(sys.argv[1], 'rb') as read_net:
        net_data = pickle.load(read_net)

    network_stuff(net_data, folder, bin_size, refine_m, max_edges, iteration)

    subprocess.run('rm ' + sys.argv[1], shell=True)

#
#
#
