import networkx as nx
from numpy import *
import math
import random
import time

def Compute_Degree(G):
    #将G.degree()的返回值变为字典
    node_to_degree = {}
    for item in G.degree():
        node_to_degree[item[0]]=item[1]
    return node_to_degree

def k_shell(G):
    #返回一个key为核数，value为节点的dict
    newG=G.copy()
    core_to_node={}
    now_ks=1
    while len(newG.degree):
        core_to_node[now_ks]=[]
        while True:
            level_node_list=[]
            for item in newG.degree:
                if item[1]<=now_ks:
                    level_node_list.append(item[0])
            newG.remove_nodes_from(level_node_list)
            core_to_node[now_ks].extend(level_node_list)
            if not len(newG.degree):
                return core_to_node
            if min(newG.degree, key=lambda x:x[1])[1]>now_ks:
                break
        now_ks=min(newG.degree, key=lambda x:x[1])[1]
    return core_to_node

def sumDegree(G):
    """
    计算G中度的和
    """
    G_degrees = Compute_Degree(G)
    sum = 0
    for v in G_degrees.values():
        sum += v
    return sum


def degree_norm(G):
    """
    归一化
    """
    sum = sumDegree(G)
    I = {}
    G_degrees = Compute_Degree(G)
    for k, v in G_degrees.items():
        I[k] = v / sum
    return I


def Entropy(G):
    """
    Entropy(G) 计算出G中所有节点的熵
    I 为重要性
    e 为节点的熵sum += I[i]*math.log(I[i])
    """
    I = degree_norm(G)
    ent_of_G = {}
    for k, v in I.items():
        sum = 0
        for i in G.neighbors(k):
            sum += I[i] * math.log(I[i])
        sum = -sum
        ent_of_G[k] = sum
    return ent_of_G


def kshellEntropy(G):

    # 计算所有壳层下，所有节点的熵值，返回字典，key是核数，value是一个字典，代表节点及其对应熵值

    ks = k_shell(G)
    ent_of_G = Entropy(G)
    kshell_ent = {}
    new_ks = sorted(ks.keys(), reverse=True)
    for ksI in new_ks:
        ksE = {}
        for i in ks[ksI]:
            ksE[i] = ent_of_G[i]
        kshell_ent[ksI] = ksE
    return kshell_ent


def kshellEntropySort(G):
    #排序
    k_shell_Ent = kshellEntropy(G)
    k_shell_ans = []
    sorted_ksEnt = sorted(k_shell_Ent.keys(), reverse=True)
    for item in sorted_ksEnt:
        t = sorted([(v, k) for k, v in k_shell_Ent[item].items()], reverse=True)

        # 熵值相同的节点放在一个set里面
        t_new = {}
        for item in t:
            t_new.setdefault(item[0], list()).append(item[1])
        # 按熵值排序变成列表
        t = sorted([(k, v) for k, v in t_new.items()], reverse=True)

        # 把相同熵值的节点列表打乱顺序随机选择
        sub_ksES = []
        for item in t:
            if len(item[1]) == 1:
                sub_ksES += item[1]
            else:
                random.shuffle(item[1])
                sub_ksES += item[1]

        k_shell_ans.append(sub_ksES)
    return k_shell_ans


def get_top_nodes(G):
    starttime = time.time()
    top_nodes = []
    node_rank = kshellEntropySort(G)
    while (len(node_rank) != 0):
        for i in range(len(node_rank)):
            top_nodes.append(node_rank[i].pop(0))
        while True:
            if [] in node_rank:
                node_rank.remove([])
            else:
                break
    endtime = time.time()
    print('改进的基于信息熵的k_shell算法： %f second' % (endtime - starttime))
    return top_nodes