import tensorflow as tf
import numpy as np

batch_size = 100
embedding_size = 200
vocabulary_size = 10000
generations = 100000
print_loss_every = 2000

num_sampled = int(batch_size / 2)  # Number of negative examples to sample.
window_size = 2  # skip-gram 模型中的 window

# %%
# 读入数据
'http://mattmahoney.net/dc/' + 'text8.zip'
import zipfile


def read_data(filename):
    """Extract the first file enclosed in a zip file as a list of words"""
    with zipfile.ZipFile(filename) as f:
        data = tf.compat.as_str(f.read(f.namelist()[0]))  # .split()
    return data


filename = 'text8.zip'
words = read_data(filename)
texts = [words]

import re
regex=re.compile('[a-zA-Z]{1,}') # 1次以上字母组合，去除标点符号和数字
texts=[regex.findall(sentence.lower()) for sentence in texts]
texts=[' '.join(words) for words in texts] # 重新整合成句子

import collections


class VocabularyProcessor:
    '''
    author: guofei
    site: www.guofei.site
    输入多个句子组成的列表，输出用序号替代单词，例如
    ```
    corpus = ['this is the first document',
              'this is the second second document',
              'and the third one',
              'is this the first document']
    wp = VocabularyProcessor(max_vocabulary_size=5)
    wp.fit(corpus)
    X = wp.transform(corpus)
    ```
    剔除低频词，并且把所有的低频词标记为 'RARE'，其序号固定为 0
    '''

    def __init__(self, max_vocabulary_size=10000):
        self.max_vocabulary_size = max_vocabulary_size
        self.word_count = None
        self.dictionary = None
        self.reverse_dictionary = None

    def fit(self, sentences):
        words = [word for sentence in sentences for word in sentence.split(' ')]
        word_count_ = collections.Counter(words)  # 词频
        word_count = word_count_.most_common(self.max_vocabulary_size - 1)  # 最多词频

        words_most = [word for word, num in word_count]  # 频率最高的词
        len_words_most = len(words_most)
        dictionary = dict(zip(words_most, range(1, len_words_most)))
        reverse_dictionary = dict(zip(range(1, len_words_most), words_most))

        # 添加 'RARE' ,也就是低频词语
        rare_count = sum([word_count_[word] for word in word_count_ if word not in words_most])
        word_count = [('RARE', rare_count)] + word_count
        dictionary['RARE'] = 0
        reverse_dictionary[0] = 'RARE'

        self.word_count = word_count
        self.dictionary = dictionary
        self.reverse_dictionary = reverse_dictionary

    def transform(self, sentences):
        return [[self.dictionary.get(word, 0) for word in sentence.split(' ')] for sentence in sentences]

    def fit_transform(self, sentences):
        self.fit(sentences)
        return self.transform(sentences)


vp = VocabularyProcessor(max_vocabulary_size=vocabulary_size)
vp.fit(texts)
text_data = vp.transform(texts)
word_dictionary = vp.dictionary
word_dictionary_rev = vp.reverse_dictionary

def generate_batch_data(sentences, batch_size, window_size):
    batch_data, label_data, batch_size_current = [], [], 0
    while True:
        rand_sentence = sentences[np.random.randint(len(sentences))]
        len_sentence = len(rand_sentence)
        for ind, word in enumerate(rand_sentence):
            x = word
            y_list = rand_sentence[max(ind - window_size, 0):ind] \
                     + rand_sentence[ind + 1:min(ind + window_size + 1, len_sentence)]
            for y in y_list:
                batch_data.append(x)
                label_data.append([y])
                batch_size_current += 1
                if batch_size_current == batch_size:
                    yield batch_data, label_data
                    batch_data, label_data, batch_size_current = [], [], 0



batch_data_generater = generate_batch_data(sentences=text_data, batch_size=batch_size, window_size=2)
# 示例：
# batch_inputs, batch_labels = batch_data_generater.next()
# 将来的预测输入：feed_dict = {x_inputs: batch_inputs, y_target: batch_labels}

# 每5000次打印一次最近的单词
print_valid_every = 5000
valid_words = ['into', 'on', 'five', 'have', 'as']
valid_examples = [word_dictionary[x] for x in valid_words]  # 待评估的单词在vocabulary中的index


# %%


x_inputs = tf.placeholder(tf.int32, shape=[batch_size])
y_target = tf.placeholder(tf.int32, shape=[batch_size, 1])
valid_dataset = tf.constant(valid_examples, dtype=tf.int32)


embeddings = tf.Variable(tf.random_uniform([vocabulary_size, embedding_size], -1.0, 1.0))
embed = tf.nn.embedding_lookup(embeddings, x_inputs)

# NCE loss parameters
nce_weights = tf.Variable(tf.truncated_normal([vocabulary_size, embedding_size], stddev=1.0 / np.sqrt(embedding_size)))
nce_biases = tf.Variable(tf.zeros([vocabulary_size]))


# Get loss from prediction
loss = tf.reduce_mean(
    tf.nn.nce_loss(weights=nce_weights,
                   biases=nce_biases,
                   labels=y_target,
                   inputs=embed,
                   num_sampled=num_sampled,
                   num_classes=vocabulary_size))

# Create optimizer
optimizer = tf.train.GradientDescentOptimizer(learning_rate=1.0).minimize(loss)

# Cosine similarity between words
norm = tf.sqrt(tf.reduce_sum(tf.square(embeddings), 1, keepdims=True))
normalized_embeddings = embeddings / norm
valid_embeddings = tf.nn.embedding_lookup(normalized_embeddings, valid_dataset)
similarity = tf.matmul(valid_embeddings, normalized_embeddings, transpose_b=True)

# %%
sess = tf.Session()
init = tf.global_variables_initializer()
sess.run(init)

# Run the skip gram model.
loss_vec = []
for i in range(generations):
    batch_inputs, batch_labels = batch_data_generater.next()
    feed_dict = {x_inputs: batch_inputs, y_target: batch_labels}
    sess.run(optimizer, feed_dict=feed_dict)

    # Return the loss
    if (i + 1) % print_loss_every == 0:
        loss_val = sess.run(loss, feed_dict=feed_dict)
        loss_vec.append(loss_val)
        print('Loss at step {} : {}'.format(i + 1, loss_val))

    # Validation: Print some words and top 5 related words
    if (i + 1) % print_valid_every == 0:
        sim = sess.run(similarity, feed_dict=feed_dict)
        for j in range(len(valid_words)):
            valid_word = valid_words[j]
            top_k = 5
            nearest = (-sim[j, :]).argsort()[1:top_k + 1]
            log_str = "Nearest to {}:".format(valid_word)+ \
                    ','.join([word_dictionary_rev[nearest_k] for nearest_k in nearest])
            print(log_str)

final_embeddings=sess.run(embeddings)
def plot_with_labels(low_dim_embs, labels, filename='tsne.png'):
    assert low_dim_embs.shape[0] >= len(labels), "More labels than embeddings"
    plt.figure(figsize=(18, 18))  #in inches
    for i, label in enumerate(labels):
        x, y = low_dim_embs[i,:]
        plt.scatter(x, y)
        plt.annotate(label,xy=(x, y),xytext=(5, 2),textcoords='offset points',
                     ha='right',va='bottom')

    plt.show()

from sklearn.manifold import TSNE
import matplotlib.pyplot as plt

tsne = TSNE(perplexity=30, n_components=2, init='pca', n_iter=5000)
plot_only = 200
low_dim_embs = tsne.fit_transform(final_embeddings[:plot_only,:])
labels = [word_dictionary_rev[i] for i in range(plot_only)]
plot_with_labels(low_dim_embs, labels)

import pandas as pd

url_neg = 'http://www.guofei.site/datasets_for_ml/rt-polaritydata/rt-polaritydata/rt-polarity.neg'
url_pos = 'http://www.guofei.site/datasets_for_ml/rt-polaritydata/rt-polaritydata/rt-polarity.pos'
df_neg = pd.read_csv(url_neg, sep='\t', header=None, names=['text'])
df_pos = pd.read_csv(url_pos, sep='\t', header=None, names=['text'])
texts = list(df_pos.text) + list(df_neg.text)
target = [1] * df_pos.shape[0] + [0] * df_neg.shape[0]


valid_words = ['cliche', 'love', 'hate', 'silly', 'sad']