import os
import math
import numpy
import time
import heapq
import codecs
from scipy import sparse

# find indices of top n largest elements in lis
def maxNum(lis, n):
    data = heapq.nlargest(n, enumerate(lis), key=lambda x:x[1])
    indices, vals = zip(*data)
    return list(indices)


def KG(KGraph, Ut, Ct, topnum, total_id2word, total_word2id, current_id2word, current_word2id, eta):

    UtUtT = Ut.dot(Ut.transpose())
    for i in range(UtUtT.shape[0]):
        UtUtT[i][i] = 0

    total_num_word = len(total_id2word)
    total_list1 = sparse.lil_matrix((total_num_word, total_num_word), dtype='int32')

    for i in range(UtUtT.shape[0]):
        word_list = maxNum(UtUtT[i], topnum)
        for j in range(topnum):
            word_i = current_id2word[i]
            t_i = total_word2id[word_i]
            word_j = current_id2word[word_list[j]]
            t_j = total_word2id[word_j]

            total_list1[t_j, t_i] = 1
            total_list1[t_i, t_j] = 1

    print("len(total_list1): ", total_list1.nnz)
    


    CtCtT = Ct.dot(Ct.transpose())
    for i in range(CtCtT.shape[0]):
        CtCtT[i][i] = 0

    total_list2 = sparse.lil_matrix((total_num_word, total_num_word), dtype='int32')

    for i in range(CtCtT.shape[0]):
        word_list = maxNum(CtCtT[i], topnum)
        for j in range(topnum):
            word_i = current_id2word[i]
            t_i = total_word2id[word_i]
            word_j = current_id2word[word_list[j]]
            t_j = total_word2id[word_j]

            total_list2[t_j, t_i] = 1*eta
            total_list2[t_i, t_j] = 1*eta

    print("len(total_list2): ", total_list2.nnz)

    KGraph += total_list1 # update KG using global info
    KGraph += total_list2 # update KG using local info

    print("KG edges size:", KGraph.nnz)

    '''
    index = total_list1[0]
    index1 = index[0]
    index2 = index[1]
    print(total_id2word[index1], total_id2word[index2])
    index = total_list2[0]
    index1 = index[0]
    index2 = index[1]
    print(total_id2word[index1], total_id2word[index2])
    '''

    return KGraph
