import pandas as pd
import numpy as np
from collections import defaultdict, Counter
import random
from tqdm.notebook import tqdm
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_distances, euclidean_distances
from sklearn.decomposition import LatentDirichletAllocation
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.utils import to_categorical
from tensorflow.keras.layers import Input, Bidirectional, LSTM, Dense, MaxPooling1D, AveragePooling1D, Conv1D
from tensorflow.keras.layers import Flatten, Reshape, Concatenate, Permute, Activation, Dropout, multiply
from tensorflow.keras.layers import BatchNormalization
from tensorflow.keras.models import Model
from tensorflow.keras.losses import cosine_similarity
from tensorflow.keras import regularizers
import tensorflow.keras.backend as K
import tensorflow as tf
import pymorphy2
import nltk
nltk.download('stopwords')
from nltk.corpus import stopwords
import numpy as np
import matplotlib.pyplot as plt
import pickle
import os
import re
from conllu import parse_incr

files = {'train': 'ru_syntagrus-ud-train.conllu',
         'test':  'ru_syntagrus-ud-test.conllu',
         'dev':   'ru_syntagrus-ud-dev.conllu'}
database = {}
for data_type in files:
    filename = files[data_type]
    database = {}
    with open(os.path.join('UD_Russian-SynTagRus-master', filename), encoding='utf-8') as f:
        parsed = parse_incr(f)
        for token_list in parsed:
            topic_name = token_list.metadata['sent_id'].split('.')[0]
            # убрём цифры из названий темы
            topic_name = re.sub(r'\d+', '', topic_name)
            if topic_name not in database:
                database[topic_name] = []
            sentence = ' '.join([token['form'] for token in token_list]).lower()
            database[topic_name].append(sentence)

choosen_for_evaluation = ['I_slepye_prozreyut',
                          'Interviyu_Mariny_Astvatsaturyan',
                          'Byudzhet']
texts_for_evaluation = {}
texts_for_training = {}
for topic in database:
    if topic in choosen_for_evaluation:
        texts_for_evaluation[topic] = database[topic]
    else:
        texts_for_training[topic] = database[topic]
TEXTS_CORPUS = [sentence for topic in texts_for_training for sentence in texts_for_training[topic]]
# напечатаем примеры
for topic in texts_for_evaluation:
    print(topic, len(texts_for_evaluation[topic]))
    for index, sentence in enumerate(texts_for_evaluation[topic]):
        print('\t', sentence[:100])
        if index > 5:
            break
    print('\n')

Byudzhet 70
     кандидат исторических наук н. митрофанов .
     все медные и серебряные на императорском кону .
     екатерина ii и бюджет .
     приходит день , и на все лады отовсюду звучит слово " бюджет " .
     газеты расшифровывают смысл его основополагающих статей , сравнивая прошлогодние наметки с нынешними
     многие возмущаются по поводу малого внимания к культуре , обороне , армии … - продолжайте перечень в
     конформисты принимаются разоблачать " выскочек " , уличая их в передергивании процентных долей , в б
Interviyu_Mariny_Astvatsaturyan 72
     " моя научная журналистика стоит на царских костях " .
     - как вы , биолог , пришли в научную журналистику ? с чего всё начиналось ?
     - началось по воле случая , в 1993-м .
     одному из тогдашних музыкальных ведущих " эха москвы " понадобились записи ива монтана , и он обрати
     у нас дома , сколько я себя помню , лежали виниловые пластинки - концертные записи ива монтана , при
     я предоставила их для подготовки радиопередачи об иве монтане .
     так я познакомилась с тогдашней редакцией " эха " , которая сидела на никольской , напротив гума .
I_slepye_prozreyut 72
     и слепые прозреют …
     опыты , проведенные специалистами института мозга человека ран , подтвердили существование у homo sa
     три года назад московский ученый вячеслав бронников начал учить слепых видеть .
     он разработал оригинальную методику , позволяющую резко активизировать деятельность правого полушари
     в результате за десять дней занятий бронников развивал у своих подопечных навыки так называемого пря
     дети с пороками зрения после специального обучения могли кататься на велосипедах , играть в шахматы
     слепой человек , как утверждают медики , видит перед собой пелену .

np.random.seed(42)
random.seed(777)
index2topic = {}
index2text = {}
index = 0
for topic in texts_for_evaluation:
    for sentence in texts_for_evaluation[topic]:
        index2topic[index] = topic
        index2text[index] = sentence
        index += 1
def get_similarity_values(sentences):
    return np.random.rand(len(sentences), len(sentences))
chart_methods = {}
bottom_minimum = -7626.2336448598135
def evaluate(get_similarity_values, method_name=None, add_to_chart=True):
    test_messages = [index2text[index] for index in range(len(index2text))]
    distances_each_to_each = get_similarity_values(test_messages)
    evaluations = []
    for target_index in index2topic:
        distances = distances_each_to_each[target_index]
        distances_indexes = sorted(zip(distances, range(len(index2topic))), key=lambda x: x[0])
        evaluation_result = 0
        for i, (distance, index) in enumerate(distances_indexes):
            if index2topic[index] == index2topic[target_index]:
                evaluation_result += len(test_messages) - i
            else:
                evaluation_result -= len(test_messages) - i
        evaluations.append(evaluation_result)
    # сделаем случайное распределние искусственным нулём (baseline)
    result = round(np.mean(evaluations) - bottom_minimum, 1)
    if add_to_chart:
        #добавляем на график, только лучший результат по всем прогонам
        if method_name not in chart_methods or result > chart_methods[method_name][0]:
            chart_methods[method_name] = (result, np.std(evaluations))
    return f'{method_name}: {str(result)}'
def parse_result(result):
    return float(new_result.split(': ')[1])
evaluate(get_similarity_values, 'random arrange')

count_vectorizer = CountVectorizer()
corpus = TEXTS_CORPUS
count_vectorizer.fit(corpus)
def get_similarity_values(sentences):
    sentences_bow = count_vectorizer.transform(sentences)
    distances = cosine_distances(sentences_bow, sentences_bow)
    return distances
evaluate(get_similarity_values, 'BOW')

morph = pymorphy2.MorphAnalyzer()
def lemmatize(corpus, verbose=False):
    clear_corpus = []
    if verbose:
        iterator = tqdm(corpus, leave=False)
    else:
        iterator = corpus
    for sentence in iterator:
        tokens = sentence.split() # разбиваем текст на слова
        res = []
        for token in tokens:
            p = morph.parse(token)[0]
            res.append(p.normal_form)
        clear_corpus.append(' '.join(res))
    return clear_corpus
count_vectorizer = CountVectorizer()
corpus = lemmatize(TEXTS_CORPUS, True)
count_vectorizer.fit(corpus)
def get_similarity_values(sentences):
    sentences_bow = count_vectorizer.transform(lemmatize(sentences))
    distances = cosine_distances(sentences_bow, sentences_bow)
    return distances
evaluate(get_similarity_values, 'BOW с леммами слов', False)

ru_stopwords = stopwords.words('russian')
ru_stopwords += ['.', ',', '"', '!',
                 '?','(', ')', '-',
                 ':', ';', '_', '\\']
def delete_stopwords(corpus, verbose=False):
    clear_corpus = []
    if verbose:
        iterator = tqdm(corpus, leave=False)
    else:
        iterator = corpus
    for sentence in iterator:
        tokens = sentence.split() # разбиваем текст на слова
        res = []
        without_stopwords = [token for token in tokens if token not in ru_stopwords]
        clear_corpus.append(' '.join(without_stopwords))
    return clear_corpus
count_vectorizer = CountVectorizer()
corpus = lemmatize(delete_stopwords(TEXTS_CORPUS), True)
count_vectorizer.fit(corpus)
def get_similarity_values(sentences):
    sentences_bow = count_vectorizer.transform(lemmatize(delete_stopwords(sentences)))
    distances = cosine_distances(sentences_bow, sentences_bow)
    return distances
evaluate(get_similarity_values, 'BOW с леммами и без стоп слов')

def similarity_values_wrapper(lda, count_vectorizer, do_lemmatize=False, do_delete_stopwords=False):
    def get_similarity_values(sentences):
        if do_delete_stopwords:
            sentences = delete_stopwords(sentences)
        if do_lemmatize:
            sentences = lemmatize(sentences)
        sent_vector = count_vectorizer.transform(sentences)
        sent_vector = lda.transform(sent_vector)
        distances = cosine_distances(sent_vector, sent_vector)
        return distances
    return get_similarity_values
lda = LatentDirichletAllocation(n_components=300)
corpus = TEXTS_CORPUS
count_vectorizer = CountVectorizer().fit(corpus)
corpus = count_vectorizer.transform(corpus)
lda.fit(corpus)
get_similarity_values = similarity_values_wrapper(lda, count_vectorizer)
print(evaluate(get_similarity_values, 'LDA', False))
lda = LatentDirichletAllocation(n_components=300)
corpus = lemmatize(TEXTS_CORPUS, True)
count_vectorizer = CountVectorizer().fit(corpus)
corpus = count_vectorizer.transform(corpus)
lda.fit(corpus)
get_similarity_values = similarity_values_wrapper(lda, count_vectorizer, do_lemmatize=True)
print(evaluate(get_similarity_values, 'LDA с леммами', True))
lda = LatentDirichletAllocation(n_components=300)
corpus = lemmatize(delete_stopwords(TEXTS_CORPUS), True)
count_vectorizer = CountVectorizer().fit(corpus)
corpus = count_vectorizer.transform(corpus)
lda.fit(corpus)
get_similarity_values = similarity_values_wrapper(lda, count_vectorizer, do_lemmatize=True, do_delete_stopwords=True)
print(evaluate(get_similarity_values, 'LDA с леммами и без стоп слов', False))

%matplotlib inline
def plot_results():
    methods = sorted(chart_methods.items(), key=lambda x: x[1][0])
    labels = [m[0] for m in methods]
    x_pos = np.arange(len(labels))
    mean = [m[1][0] for m in methods]
    std = [m[1][1] for m in methods]
    # Build the plot
    fig, ax = plt.subplots(figsize=(12,8))
    ax.bar(x_pos,
           mean,
           yerr=std,
           align='center',
           alpha=0.5,
           ecolor='black',
           capsize=10)
    ax.set_ylabel('Баллы и стандартное отклонение')
    ax.set_xticks(x_pos)
    ax.set_xticklabels(labels, rotation=20, ha='right')
    ax.set_title('Сранительный график методов')
    ax.yaxis.grid(True)
    plt.show()
plot_results()

# Скачивание предобученных эмбедингов
import fasttext.util
from wikipedia2vec import Wikipedia2Vec
fasttext.util.download_model('ru', if_exists='ignore')
wiki2vec = Wikipedia2Vec.load('ruwiki_20180420_300d.pkl')
ft = fasttext.load_model('cc.ru.300.bin')

def vectorize(token, use_word2vec=True, use_fasttext=True):
    assert use_word2vec or use_fasttext
    if use_fasttext:
        try:
            fast_text_vector = ft.get_word_vector(token)
        except KeyError:
            fast_text_vector = np.zeros((ft.get_dimension()))
    if use_word2vec:
        try:
            word2vec_vector = wiki2vec.get_word_vector(token)
        except KeyError:
            word2vec_vector = np.zeros((len(wiki2vec.get_word_vector('the'))))
    if use_fasttext and use_word2vec:
        return np.concatenate([word2vec_vector, fast_text_vector])
    elif use_fasttext:
        return np.array(fast_text_vector)
    elif use_word2vec:
        return np.array(word2vec_vector)
    else:
        return 'something went wrong on vectorisation'
print(np.shape(vectorize('any_token')))

def similarity_values_wrapper(use_word2vec=True, use_fasttext=True, distance_function=cosine_distances):
    def get_similarity_values(sentences):
        sent_vector = []
        for sentence in sentences:
            sentence_vector = []
            for token in sentence.split():
                sentence_vector.append(vectorize(token, use_word2vec, use_fasttext))
            sent_vector.append(np.mean(sentence_vector, axis=0))
        distances = distance_function(sent_vector, sent_vector)
        return distances
    return get_similarity_values
get_similarity_values = similarity_values_wrapper(use_word2vec=True, use_fasttext=True, distance_function=euclidean_distances)
print(evaluate(get_similarity_values, 'среднее по embedings с euclidean_distances с word2vec + fast_text', add_to_chart=False))
get_similarity_values = similarity_values_wrapper(use_word2vec=False, use_fasttext=True, distance_function=euclidean_distances)
print(evaluate(get_similarity_values, 'среднее по embedings с euclidean_distances с fast_text', add_to_chart=False))
get_similarity_values = similarity_values_wrapper(use_word2vec=True, use_fasttext=False, distance_function=euclidean_distances)
print(evaluate(get_similarity_values, 'среднее по embedings с euclidean_distances с word2vec'))
get_similarity_values = similarity_values_wrapper(use_word2vec=True, use_fasttext=True, distance_function=cosine_distances)
print(evaluate(get_similarity_values, 'среднее по embedings с cosine_distance с word2vec + fast_text', add_to_chart=False))
get_similarity_values = similarity_values_wrapper(use_word2vec=False, use_fasttext=True, distance_function=cosine_distances)
print(evaluate(get_similarity_values, 'среднее по embedings с cosine_distance с fast_text', add_to_chart=False))
get_similarity_values = similarity_values_wrapper(use_word2vec=True, use_fasttext=False, distance_function=cosine_distances)
print(evaluate(get_similarity_values, 'среднее по embedings с cosine_distance с word2vec'))

def similarity_values_wrapper(use_word2vec=True, use_fasttext=True, distance_function=cosine_distances):
    def get_similarity_values(sentences):
        sentences = delete_stopwords(sentences)
        sent_vector = []
        for sentence in sentences:
            sentence_vector = []
            for token in sentence.split():
                sentence_vector.append(vectorize(token, use_word2vec, use_fasttext))
            sent_vector.append(np.mean(sentence_vector, axis=0))
        distances = distance_function(sent_vector, sent_vector)
        return distances
    return get_similarity_values
get_similarity_values = similarity_values_wrapper(use_word2vec=True, use_fasttext=True, distance_function=euclidean_distances)
print(evaluate(get_similarity_values, 'среднее по embedings без стопслов с euclidean_distances с word2vec + fast_text', add_to_chart=False))
get_similarity_values = similarity_values_wrapper(use_word2vec=False, use_fasttext=True, distance_function=euclidean_distances)
print(evaluate(get_similarity_values, 'среднее по embedings без стопслов с euclidean_distances с fast_text', add_to_chart=False))
get_similarity_values = similarity_values_wrapper(use_word2vec=True, use_fasttext=False, distance_function=euclidean_distances)
print(evaluate(get_similarity_values, 'среднее по embedings без стопслов с euclidean_distances с word2vec'))
get_similarity_values = similarity_values_wrapper(use_word2vec=True, use_fasttext=True)
print(evaluate(get_similarity_values, 'среднее по embedings без стопслов с cosine_distance с word2vec + fast_text', add_to_chart=False))
get_similarity_values = similarity_values_wrapper(use_word2vec=False, use_fasttext=True)
print(evaluate(get_similarity_values, 'среднее по embedings без стопслов с cosine_distance с fast_text', add_to_chart=False))
get_similarity_values = similarity_values_wrapper(use_word2vec=True, use_fasttext=False)
print(evaluate(get_similarity_values, 'среднее по embedings без стопслов с cosine_distance с word2vec'))

tf_idf_vectorizer = TfidfVectorizer()
tf_idf_vectorizer.fit(TEXTS_CORPUS)
vocab = tf_idf_vectorizer.get_feature_names()

def similarity_values_wrapper(use_word2vec=True, use_fasttext=True, distance_function=cosine_distances):
    def get_similarity_values(sentences):
        sent_vector = [[]]*len(sentences)
        weights_data = tf_idf_vectorizer.transform(sentences).tocoo()
        for row, col, weight in zip(weights_data.row, weights_data.col, weights_data.data):
            sent_vector[row].append(weight*vectorize(vocab[col], use_word2vec, use_fasttext))
        for row in range(len(sent_vector)):
            if not sent_vector[row]:
                sent_vector.append((len(vectorize('zoros_vector'))))
        sent_vector = np.sum(sent_vector, axis=1)
        distances = distance_function(sent_vector, sent_vector)
        return distances
    return get_similarity_values
get_similarity_values = similarity_values_wrapper(use_word2vec=True, use_fasttext=True)
print(evaluate(get_similarity_values,'среднее по embedings с tf-idf с cosine_distance с word2vec + fast_text', add_to_chart=False))
get_similarity_values = similarity_values_wrapper(use_word2vec=False, use_fasttext=True)
print(evaluate(get_similarity_values,'среднее по embedings с tf-idf с cosine_distance с fast_text', add_to_chart=False))
get_similarity_values = similarity_values_wrapper(use_word2vec=True, use_fasttext=False)
print(evaluate(get_similarity_values,'среднее по embedings с tf-idf с cosine_distance с word2vec', add_to_chart=False))
get_similarity_values = similarity_values_wrapper(use_word2vec=True, use_fasttext=True, distance_function=euclidean_distances)
print(evaluate(get_similarity_values,'среднее по embedings с tf-idf с euclidian_distance с word2vec + fast_text', add_to_chart=True))
get_similarity_values = similarity_values_wrapper(use_word2vec=False, use_fasttext=True, distance_function=euclidean_distances)
print(evaluate(get_similarity_values,'среднее по embedings с tf-idf с euclidian_distance с fast_text', add_to_chart=False))
get_similarity_values = similarity_values_wrapper(use_word2vec=True, use_fasttext=False, distance_function=euclidean_distances)
print(evaluate(get_similarity_values,'среднее по embedings с tf-idf с euclidian_distance с word2vec', add_to_chart=False))

plot_results()

max_len = 20
min_len = 5
embedding_size = len(vectorize('any token'))
class EmbedingsDataGenerator():
    def __init__(self, texts_corpus=TEXTS_CORPUS, min_len=5, max_len=20, batch_size=32, batches_per_epoch=100, use_word2vec=True, use_fasttext=True):
        self.texts = texts_corpus
        self.min_len = min_len
        self.max_len = max_len
        self.batch_size = batch_size
        self.batches_per_epoch = batches_per_epoch
        self.use_word2vec = use_word2vec
        self.use_fasttext = use_fasttext
        self.embedding_size = len(vectorize('token', use_word2vec=self.use_word2vec, use_fasttext=self.use_fasttext))
    def vectorize(self, sentences):
        vectorized_sentences = []
        for text in sentences:
            text_vec = []
            tokens = str(text).split()
            for token in tokens:
                text_vec.append(vectorize(token, use_word2vec=self.use_word2vec, use_fasttext=self.use_fasttext))
            vectorized_sentences.append(text_vec)
        vectorized_sentences = pad_sequences(vectorized_sentences, maxlen=self.max_len, dtype='float32')
        return vectorized_sentences
    def __iter__(self):
        for _ in tqdm(range(self.batches_per_epoch), leave=False):
            X_batch = []
            y_batch = []
            finished_batch = False
            while not finished_batch:
                text = random.choice(self.texts)
                tokens = str(text).split()
                if len(tokens) < self.min_len:
                    continue
                x_vec = []
                for token in tokens:
                    token_vec = vectorize(token, use_word2vec=self.use_word2vec, use_fasttext=self.use_fasttext)
                    if len(x_vec) >= self.min_len:
                        X_batch.append(x_vec)
                        y_batch.append(token_vec)
                        if len(X_batch) == self.batch_size:
                            X_batch = pad_sequences(X_batch, maxlen=self.max_len, dtype='float32')
                            yield np.array(X_batch), np.array(y_batch)
                            finished_batch = True
                            break
                    x_vec.append(token_vec)
class IndexesDataGenerator(EmbedingsDataGenerator):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.token2index = {}
        index = 0
        for text in self.texts:
            tokens = str(text).split()
            for token in tokens:
                if token not in self.token2index:
                    self.token2index[token] = index
                    index += 1
    def __iter__(self):
        for _ in tqdm(range(self.batches_per_epoch), leave=False):
            X_batch = []
            X_batch_indexes = []
            y_batch = []
            finished_batch = False
            while not finished_batch:
                text = random.choice(self.texts)
                tokens = str(text).split()
                if len(tokens) < self.min_len:
                    continue
                x_vec = []
                x_tokens = []
                for token in tokens:
                    token_vec = vectorize(token, use_word2vec=self.use_word2vec, use_fasttext=self.use_fasttext)
                    if len(x_vec) >= self.min_len:
                        X_batch.append(x_vec)
                        X_batch_indexes.append(to_categorical(x_tokens, num_classes=len(self.token2index)))
                        y_batch.append(self.token2index[token])
                        if len(X_batch) == self.batch_size:
                            X_batch = pad_sequences(X_batch, maxlen=self.max_len, dtype='float32')
                            X_batch_indexes = pad_sequences(X_batch_indexes, maxlen=self.max_len, dtype='int32')
                            y_batch = to_categorical(y_batch, num_classes=len(self.token2index))
                            yield np.array(X_batch), np.array(X_batch_indexes), np.array(y_batch)
                            finished_batch = True
                            break
                    x_vec.append(token_vec)
                    x_tokens.append(self.token2index[token])

data_generator = EmbedingsDataGenerator()

%%timeit
for x, y in data_generator:
    pass

data_generator = IndexesDataGenerator()

%%timeit
for x_e, x_i, y_i in data_generator:
    pass

def similarity_values_wrapper(embedder, vectorizer, distance_function=cosine_distances):
    def get_similarity_values(sentences):
        sent_vec = vectorizer(sentences)
        sent_embedings = embedder(sent_vec)
        distances = distance_function(sent_embedings, sent_embedings)
        return distances
    return get_similarity_values

def model_builder(data_generator):
    complexity = 500
    inp = Input(shape=(data_generator.max_len, data_generator.embedding_size))
    X = inp
    X = LSTM(complexity, return_sequences=True)(X)
    X = LSTM(complexity)(X)
    X = Dense(complexity, activation='elu')(X)
    X = Dense(complexity, activation='elu')(X)
    X = Dense(data_generator.embedding_size, activation='linear')(X)
    model = Model(inputs=inp, outputs=X)
    model.compile(loss=cosine_similarity, optimizer='adam')
    model.summary()
    return model
data_generator = EmbedingsDataGenerator(use_fasttext=False)
next_word_model = model_builder(data_generator)
get_similarity_values = similarity_values_wrapper(next_word_model.predict, data_generator.vectorize)

new_result = -10e5
for i in tqdm(range(1000)):
    if i%3==0:
        previous_result = new_result
        new_result = evaluate(get_similarity_values, 'Language Model on embedings')
        new_result = parse_result(new_result)
        print(i, new_result)
        # stopping condition
        if new_result < previous_result and i > 20:
            break
    for x, y in data_generator:
        next_word_model.train_on_batch(x, y)

def model_builder(data_generator):
    complexity = 200
    inp = Input(shape=(data_generator.max_len, data_generator.embedding_size))
    X = inp
    X = LSTM(complexity, return_sequences=True)(X)
    X = LSTM(complexity)(X)
    X = Dense(complexity, activation='linear', name='embedding_output')(X)
    X = Dense(complexity, activation='elu')(X)
    X = Dense(len(data_generator.token2index), activation='softmax')(X)
    model = Model(inputs=inp, outputs=X)
    model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['acc'])
    model.summary()
    embedder = Model(inputs=inp, outputs=model.get_layer('embedding_output').output)
    return model, embedder
data_generator = IndexesDataGenerator()
next_word_model, embedder = model_builder(data_generator)
get_similarity_values = similarity_values_wrapper(embedder.predict, data_generator.vectorize)

new_result = -10e5
for i in tqdm(range(1000)):
    if i%3==0:
        previous_result = new_result
        new_result = evaluate(get_similarity_values, 'Language Model on token index')
        new_result = parse_result(new_result)
        print(i, new_result)
        if new_result < previous_result and i > 20:
            break
    for x_e, x_i, y in data_generator:
        next_word_model.train_on_batch(x_e, y)

plot_results()