#!/usr/bin/env python
# coding=utf-8

from multiprocessing import cpu_count

import numpy as np

from gensim.models.word2vec import Word2Vec
from gensim.models.doc2vec import Doc2Vec, TaggedDocument
from keras.preprocessing.text import Tokenizer, text_to_word_sequence
from keras.utils import to_categorical

from nltk.downloader import download
from nltk.corpus import stopwords
from nltk.tokenize import RegexpTokenizer, sent_tokenize
from nltk.stem import WordNetLemmatizer
import pymorphy2 as pm
from alphabet_detector import AlphabetDetector
from sklearn.preprocessing import MultiLabelBinarizer

RUSSIAN_REGEX = u'[А-Яа-яёЁa-zA-Z^,!.+-\/\']+'

def load(data_folder):
    X_train = np.load(data_folder + "x_train.npy")
    X_test = np.load(data_folder + "x_test.npy")
    Y_train = np.load(data_folder + "y_train.npy")
    Y_test = np.load(data_folder + "y_test.npy")
    idx_train = np.load(data_folder + "idx_train.npy")
    idx_test = np.load(data_folder + "idx_test.npy")
    document_X_title = np.load(data_folder + "x_title.npy")

    return X_train, X_test, Y_train, Y_test, idx_train, idx_test, document_X_title

def dump(data_folder, X_train=None, X_test=None, Y_train=None, Y_test=None, idx_train=None, idx_test=None, document_X_title=None):
    if X_train is not None:
        np.save(data_folder + "x_train.npy", X_train)
    if X_test is not None:
        np.save(data_folder + "x_test.npy", X_test)
    if Y_train is not None:
        np.save(data_folder + "y_train.npy", Y_train)
    if Y_test is not None:
        np.save(data_folder + "y_test.npy", Y_test)
    if idx_train is not None:
        np.save(data_folder + "idx_train.npy", idx_train)
    if idx_test is not None:
        np.save(data_folder + "idx_test.npy", idx_test)
    if document_X_title is not None:
        np.save(data_folder + "x_title.npy", document_X_title)

def print_labels(pr):
    return " ".join(str(x) for x in pr)


def doc_to_vec(newsline_documents, number_of_documents, document_Y, selected_categories, data_folder, model_name=None,
               num_features=500, document_max_num_words=100, load=False):
    if model_name is None:
        model_name = 'reuters.doc2vec'

    if load is True:
        # Load an existing Word2Vec model
        d2v_model = Doc2Vec.load(data_folder + model_name)
    else:
        corpus = []
        for i, words in enumerate(newsline_documents):
            corpus.append(TaggedDocument(words, [i]))

        d2v_model = Doc2Vec(corpus, size=num_features, min_count=5, window=10, workers=cpu_count())
        d2v_model.init_sims(replace=True)
        d2v_model.save(data_folder + model_name)

    num_categories = len(selected_categories)
    X = np.zeros(shape=(number_of_documents, document_max_num_words, num_features)).astype(np.float32)
    Y = np.zeros(shape=(number_of_documents, num_categories)).astype(np.float32)

    for idx, text in enumerate(newsline_documents):
        inferred_vec = d2v_model.infer_vector(text)
        X[idx, :] = inferred_vec

    for idx, key in enumerate(document_Y.keys()):
        Y[idx, :] = document_Y[key]

    return X, Y, num_categories


def word_to_vec(newsline_documents, number_of_documents, document_Y, selected_categories, data_folder, model_name=None,
                num_features=500, document_max_num_words=100, load=False, sg=None):
    if model_name is None:
        model_name = 'reuters.word2vec'

    if load is True:
        # Load an existing Word2Vec model
        w2v_model = Word2Vec.load(data_folder + model_name)
    else:
        w2v_model = Word2Vec(newsline_documents, size=num_features, sg=sg, min_count=1, window=10, workers=cpu_count())
        w2v_model.init_sims(replace=True)
        w2v_model.save(data_folder + model_name)

    num_categories = len(selected_categories)
    X = np.zeros(shape=(number_of_documents, document_max_num_words, num_features)).astype(np.float32)
    Y = np.zeros(shape=(number_of_documents, num_categories)).astype(np.float32)

    empty_word = np.zeros(num_features).astype(np.float32)

    for idx, document in enumerate(newsline_documents):
        for jdx, word in enumerate(document):
            if jdx == document_max_num_words:
                break

            else:
                if word in w2v_model:
                    X[idx, jdx, :] = w2v_model[word]
                else:
                    X[idx, jdx, :] = empty_word

    for idx, key in enumerate(document_Y.keys()):
        Y[idx, :] = document_Y[key]

    return X, Y, num_categories


def tokenize_prepare():
    download('stopwords')
    download('wordnet')
    download('punkt')


def tokenize_documents(document_X, document_Y, lang=None, regex=None, decode=False):
    tokenize_prepare()

    if lang is None:
        lang = 'english'

    if regex is None:
        regex = '[\'a-zA-Z]+'

    # Load stop-words
    stop_words_en = set(stopwords.words('english'))
    stop_words_any = set(stopwords.words(lang))
    stop_words = stop_words_en.union(stop_words_any)

    # Initialize tokenizer
    # It's also possible to try with a stemmer or to mix a stemmer and a lemmatizer
    tokenizer = RegexpTokenizer(regex)

    pm_analyzer = pm.MorphAnalyzer()
    lemmatizer_ru = lambda (w): pm_analyzer.parse(w)[0].normal_form

    wn_analyzer = WordNetLemmatizer()
    lemmatizer_en = lambda (w): wn_analyzer.lemmatize(w)

    ad = AlphabetDetector()

    # Initialize lemmatizer
    def lemmatizer(w):
        if ad.is_cyrillic(w):
            return lemmatizer_ru(w)
        else:
            return lemmatizer_en(w)

    # Tokenized document collection
    newsline_documents = []

    def decode_document(document):
        if decode is True:
            return document.decode('utf-8')
        else:
            return document

    def tokenize(document):
        words = []

        for sentence in sent_tokenize(decode_document(document)):
            tokens = [lemmatizer(t.lower()) for t in tokenizer.tokenize(sentence) if
                      t.lower() not in stop_words]
            words += tokens

        return words

    # Tokenize
    for key in document_X.keys():
        newsline_documents.append(tokenize(document_X[key]))

    number_of_documents = len(document_X)

    return newsline_documents, number_of_documents


def print_predictions(predicted, x, classes, idx, y=None, with_keys=False, show_words=50, encode=True):
    with_labels = []

    for item in predicted:
        pr = zip(classes, item)
        p = sorted(pr, key=lambda t: t[1], reverse=True)
        zipped = list(map((lambda x:
                           (x[0], float("%.3f" % x[1]))),
                          p[:5])
                      )
        with_labels.append(zipped)

    def encoded(value):
        if encode is True:
            return value.encode('utf-8')
        else:
            return value

    if with_keys is True:
        x_getter = lambda (i): encoded(x[i.astype(str)][:show_words]) if i > 0 else None
    else:
        x_getter = lambda (i): encoded(x[i][:show_words]) if i > 0 else None

    if y is None:
        y_getter = lambda (i): None
    else:
        y_getter = lambda (i): encoded(print_labels(y[i]))

    for i, pr in zip(idx, with_labels):
        line = x_getter(i)
        y_line = y_getter(i)
        labels_string = encoded(print_labels(pr))
        # cats = y_train_text[]

        print('{0} => {1} => {2}'.format(line, labels_string, y_line))


def update_frequencies(news_categories, categories, column='Newslines'):
    for category in categories:
        idx = news_categories[news_categories.Name == category].index[0]
        f = news_categories.get_value(idx, column)
        news_categories.set_value(idx, column, f + 1)

    return news_categories


def to_category_vector(categories, target_categories):
    vector = np.zeros(len(target_categories)).astype(np.float32)

    for i in range(len(target_categories)):
        if target_categories[i] in categories:
            vector[i] = 1.0

    return vector


MAX_NB_WORDS = 20000


def keras_prepare_text(df, y=None, max_sent_length=500, max_sents=100):
    texts = []
    labels = []
    sources = []

    for idx in range(df.Text.shape[0]):
        text = df.Text[idx]
        sources.append(text)
        sentences = sent_tokenize(text)
        texts.append(sentences)

        labels.append(df.Category[idx])

    tokenizer = Tokenizer(num_words=MAX_NB_WORDS)
    tokenizer.fit_on_texts(sources)

    data = np.zeros((len(sources), max_sents, max_sent_length), dtype='int32')

    for i, sentences in enumerate(texts):
        for j, sent in enumerate(sentences):
            if j < max_sents:
                wordTokens = text_to_word_sequence(sent)
                k = 0
                for _, word in enumerate(wordTokens):
                    if k < max_sent_length and tokenizer.word_index[word] < MAX_NB_WORDS:
                        data[i, j, k] = tokenizer.word_index[word]
                        k = k + 1

    word_index = tokenizer.word_index
    print('Total %s unique tokens.' % len(word_index))

    encoder = MultiLabelBinarizer()
    encoded_Y = encoder.fit_transform(y)
    labels = to_categorical(encoded_Y, num_classes=len(encoded_Y[0]))
    # labels = to_categorical(np.asarray(labels))

    print('Shape of data tensor:', data.shape)
    print('Shape of label tensor:', labels.shape)

    indices = np.arange(data.shape[0])
    np.random.shuffle(indices)
    data = data[indices]
    labels = labels[indices]

    return data, labels
