import nltk
# nltk.download('wordnet') # 第一次運行需運行此命令，安裝wordnet數(shù)據(jù)集
from nltk.corpus import wordnet as wn
from math import *
import random
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import networkx as nx

network = {} # 構(gòu)建層級網(wǎng)絡(luò)
last_level = 8 # 最深的層設(shè)置為8層
levelOfNode = {} # 數(shù)據(jù)的層級信息，0為哺乳動物（根節(jié)點），1為哺乳動物下一結(jié)構(gòu)
# 遞歸構(gòu)建network
def get_hyponyms(synset, level):
    if (level == last_level):
        levelOfNode[str(synset)] = level
        return
    if not str(synset) in network:
        network[str(synset)] = [str(s) for s in synset.hyponyms()]
        levelOfNode[str(synset)] = level
    for hyponym in synset.hyponyms():
        get_hyponyms(hyponym, level + 1)
# 構(gòu)建以哺乳動物為根節(jié)點的層次結(jié)構(gòu)數(shù)據(jù)集
mammal = wn.synset('mammal.n.01')
get_hyponyms(mammal, 0)
levelOfNode[str(mammal)] = 0
# 將終端葉子節(jié)點補到network字典中
for a in levelOfNode:
    if not a in network:
        network[a] = []

def norm(x):
    return np.dot(x, x)

def traverse(graph, start, node):
    node_name = node.name().split(".")[0]
    graph.depth[node_name] = node.shortest_path_distance(start)
    for child in node.hyponyms():
        child_name = child.name().split(".")[0]
        graph.add_edge(node_name, child_name) # 添加邊
        traverse(graph, start, child) # 遞歸構(gòu)建

def hyponym_graph(start):
    G = nx.Graph() # 定義一個圖
    G.depth = {}
    traverse(G, start, start)
    return G

def graph_draw(graph):
    plt.figure(figsize=(10, 10)) # 展示整體的網(wǎng)絡(luò)
    # plt.figure(figsize=(3, 3)) # 展示大象網(wǎng)絡(luò)
    nx.draw(graph,
         node_size = [10 * graph.degree(n) for n in graph],
         node_color = [graph.depth[n] for n in graph],
         alpha = 0.8,
         font_size = 4,
         width = 0.5,
         with_labels = True)
    
def get_keys(d, value):
        return [k for k,v in d.items() if v == value]
    root_name = get_keys(graph.depth, 0)[0]
    plt.savefig("~/hyperE/fig/" + root_name + ".png", dpi = 300)
    
graph = hyponym_graph(mammal)
graph_draw(graph)

elephant = wn.synset('elephant.n.01')
graph = hyponym_graph(elephant)
graph_draw(graph)