import pandas as pd
import copy
from collections import Counter
from tqdm import tqdm
import sys
import random
sys.path.append('GraphEmbedding')
#from GraphEmbedding.examples.struc2vec_flight import *
import time
import numpy as np
import matplotlib.pyplot as plt
import networkx as nx
from sklearn.manifold import TSNE
import pandas as pd
from collections import Counter


data_list = [set(['A','B']),set(['D','C']),set(['E','A']),set(['E','D']),set(['A','D']),set(['B','D'])]


item_id = set()
data_list_3 = {}
total_freq = len(data_list)

for dl in tqdm(data_list):
    # 
    for d in dl:
        item_id.add(str(d).replace(' ',''))
    # 
    dl = copy.deepcopy(tuple(str(d).replace(' ','') for d in dl))
    if dl not in data_list_3.keys():
        data_list_3[dl] = 0
    data_list_3[dl] += 1

data_label = {str(ii).replace(' ',''):str(n) for n,ii in enumerate(item_id)} # 节点：节点编号
data_list_3 = {k:v for k,v in data_list_3.items()}   # [A,B]出现频率

# 数据导出 - 节点-节点
edgelist_file = open('GraphEmbedding/test/data.edgelist','w+',encoding = 'utf-8') 

for k,v in tqdm(data_list_3.items()):
    k = [str(_k) for _k in k]
    text = ' '.join([data_label[_k] for _k in k]+[str(v)]) + '
'
    # text = ' '.join(list(k)) + '
'
    edgelist_file.write(text)

edgelist_file.close()

# 数据导出 - 节点-解释
edgelist_label = open('GraphEmbedding/test/data.label','w+',encoding = 'utf-8') 

for k,v in tqdm(data_label.items()):
    text = ' '.join([k,v]) + '
'
    # text = ' '.join(list(k)) + '
'
    edgelist_label.write(text)

edgelist_label.close()



#----
def bit_product_sum(x, y):
    return sum([item[0] * item[1] for item in zip(x, y)])


def cosine_similarity(x, y, norm=False):
    """ 计算两个向量x和y的余弦相似度 """
    assert len(x) == len(y), "len(x) != len(y)"
    zero_list = [0] * len(x)
    if x == zero_list or y == zero_list:
        return float(1) if x == y else float(0)

    # method 1
    res = np.array([[x[i] * y[i], x[i] * x[i], y[i] * y[i]] for i in range(len(x))])
    cos = sum(res[:, 0]) / (np.sqrt(sum(res[:, 1])) * np.sqrt(sum(res[:, 2])))

    # method 2
    # cos = bit_product_sum(x, y) / (np.sqrt(bit_product_sum(x, x)) * np.sqrt(bit_product_sum(y, y)))

    # method 3
    # dot_product, square_sum_x, square_sum_y = 0, 0, 0
    # for i in range(len(x)):
    #     dot_product += x[i] * y[i]
    #     square_sum_x += x[i] * x[i]
    #     square_sum_y += y[i] * y[i]
    # cos = dot_product / (np.sqrt(square_sum_x) * np.sqrt(square_sum_y))

    return 0.5 * cos + 0.5 if norm else cos  # 归一化到[0, 1]区间内


def plot_embeddings(emb_list,color_random = []):
    #X, Y = read_node_label('../data/wiki/wiki_labels.txt')
    X = emb_list
    if len(color_random) == 0:
        color_random = [random.sample(range(20), 1)[0] for _ in range(len(embeddings))]

    model = TSNE(n_components=2)
    node_pos = model.fit_transform(emb_list)
    
    color_idx = {}
    for i in range(len(X)):
        color_idx.setdefault(color_random[i], [])
        color_idx[color_random[i]].append(i)

    for c, idx in color_idx.items():
        plt.scatter(node_pos[idx, 0], node_pos[idx, 1], label=c)
    plt.legend()
    plt.show()


G = nx.read_edgelist('GraphEmbedding/test/data.edgelist', create_using=nx.DiGraph(), nodetype=None,
                     data=[('weight', int)])
# struc2vec -> 最有规律
from ge import Struc2Vec
s2v_model = Struc2Vec(G, 10, 80, workers=4, verbose=40, )
s2v_model.train()
embeddings = s2v_model.get_embeddings()



# from ge import SDNE
# sdne_model = SDNE(G, hidden_size=[1024, 512],)
# sdne_model.train(batch_size=20000, epochs=50, verbose=2)
# embeddings = sdne_model.get_embeddings()

# node2vec -> 有一小簇有规律
from ge import Node2Vec
n2v_model = Node2Vec(G, walk_length=10, num_walks=80,
                 p=0.25, q=4, workers=1, use_rejection_sampling=0)
n2v_model.train(window_size = 5, iter = 3)
embeddings=n2v_model.get_embeddings()

# DeepWalk -> 最散
from ge import DeepWalk
dw_model = DeepWalk(G, walk_length=10, num_walks=80, workers=1)
dw_model.train(window_size=5, iter=3)
embeddings = dw_model.get_embeddings()

# 画图
emb_list = np.array(list(embeddings.values()))
#evaluate_embeddings(embeddings)
plot_embeddings(emb_list)


# id中文名
label = pd.read_table('GraphEmbedding/test/data.label',header = None)
id_to_label = {l.split(' ')[1]:l.split(' ')[0]  for l in label[0]}
label_to_id = {l.split(' ')[0]:l.split(' ')[1]  for l in label[0]}


# 求相似产品
simi = n2v_model.w2v_model.wv.most_similar(label_to_id['D'],topn = 20)
simi
n2v_model.sentences