import re
import pickle
import numpy as np
import pandas as pd
from pprint import pprint

# Gensim
import gensim
import gensim.corpora as corpora
from gensim.utils import simple_preprocess
from gensim.models import CoherenceModel

import nltk
from nltk.corpus import stopwords 
from nltk.stem.wordnet import WordNetLemmatizer
# nltk.download('wordnet')
# nltk.download('stopwords')
lemma = WordNetLemmatizer()

# Plotting tools
import pyLDAvis
import pyLDAvis.gensim  # don't skip this
import matplotlib.pyplot as plt

# Enable logging for gensim - optional
import logging
logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s', level=logging.ERROR)

import warnings
warnings.filterwarnings("ignore",category=DeprecationWarning)

import nltk
nltk.download('wordnet')


documents = ['На втором этаже не работает кондиционер, очень душно, невозможно работать.'
    ,'Сделайте кондиционер на втором этаже, жарко'
    ,'Нужен кондиционер на втором этаже, чтобы не было душно'
         ,'Хотелось бы кондиционер на втором этаже, но нужно будет поставить его туда, где не сидят сотрудники, чтобы не дуло сильно'
        ,'Очень жарко в офисе, особенно на втором этаже, нужен кондиционер'
        ,'В кофейне на первом этаже кофе очень плохого качества'
        ,'Поменяйте нашу кофейню, кофе ужасный и невкусный'
        ,'Кофе на первом этаже плохой, пусть будет дорогая кофейня, но с хорошим кофе, а эту, думаю, нужно убрать'
        ,'Кофейня делает ужасный кофе'
        ,'Coffee boom нужно заменить, кофе некачественный'
        ,'Лифт уже две недели не работает, сложно подниматься на четвёртый этаж, скорей почините, пожалуйста'
        ,'У меня проблемы с ногами, мне сложно подниматься и спускаться с третьего этажа, поскорей почините лифт'
        ,'Почему лифт уже две недели не работает? Сотрудникам приходится спускаться и подниматься по лестнице по несколько раз в день'
        ,'Сделайте лифт от другой компании, чтобы не ломались без конца'
        ,'Установите камеры на парковке, мою машину кто-то поцарапал'
        ,'Мою машину стукнули на парковке, почему там нет камер, чтобы посмотреть?!'
        ,'Машину моего коллеги поцарапали на нашей парковке и внаглую уехали, почему там нет камер?!']



def clean(text):
    # Remove new line characters
    text = re.sub('\s+', ' ', text)
    # Remove distracting single quotes
    text = re.sub("\'", "", text)
    return text


def doc_to_words(doc, stop_words, lemma):
    stop_words = set(stop_words)
    # remove stop words and punctuation
    words = [w for w in gensim.utils.simple_preprocess(str(doc), deacc=True) if w not in stop_words]
    
    # make lemmatization
    words = [lemma.lemmatize(w) for w in words]
    
    return words


def main():
	stop_words = stopwords.words('russian')
	stop_words = set(stop_words)
	if prepare:
	    # Convert to list
	    data = documents
	    data = [clean(t) for t in data]
	    data_words = [doc_to_words(t, stop_words, lemma) for t in data]
	    
	    with open('data.pkl', 'wb') as f:
	        pickle.dump({'data': data, 'data_words': data_words}, f)
	else:
	    with open('data.pkl', 'rb') as f:
	        d = pickle.load(f)
	        data = d['data']
	        data_words = d['data_words']

	if prepare:
		# Build the bigram and trigram models
	    bigram = gensim.models.Phrases(data_words, min_count=5, threshold=100) # higher threshold fewer phrases.

	    # Faster way to get a sentence clubbed as a trigram/bigram
	    bigram_mod = gensim.models.phrases.Phraser(bigram)
	    bigram_mod.save('bigram_mod.pkl')
	    
	    data_words_bigrams = [bigram_mod[w] for w in data_words]
	    with open('bigrams.pkl', 'wb') as f:
	        pickle.dump(data_words_bigrams, f)

	else:
	    bigram_mod = gensim.models.Phrases.load('bigram_mod.pkl')
	    
	    with open('bigrams.pkl', 'rb') as f:
	        data_words_bigrams = pickle.load(f)
	
	if prepare:
		id2word = corpora.Dictionary(data_words_bigrams)
		id2word.save('id2word.pkl')
	else:
	    id2word = corpora.Dictionary.load('id2word.pkl')

	corpus = [id2word.doc2bow(text) for text in data_words_bigrams]
	
	if prepare:
		lda_model = gensim.models.ldamodel.LdaModel(
			    corpus=corpus,
			    id2word=id2word,
			    num_topics=4, 
			    random_state=100,
			    update_every=1,
			    chunksize=100,
			    passes=10,
			    alpha='auto',
			    per_word_topics=True)
		lda_model.save('lda_default.pkl')
	else:
		lda_model = gensim.models.ldamodel.LdaModel.load('lda_default.pkl')
	print(1245)
	vis = pyLDAvis.gensim.prepare(lda_model, corpus, id2word)
	pyLDAvis.save_html(vis, 'lda.html')

if __name__== "__main__":
	prepare = True
	main()