# coding=utf8
"""
@author: <NAME>
@date: 09/26/2019
@code description: It is a Python3 file to implement cosine similarity with TF-IDF and Word Embedding methods.
"""

from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import statistics

from nltk.stem import WordNetLemmatizer
from nltk.corpus import stopwords


glove_path = "../../data/glove.6B.50d.txt"

punctuation_list = ['.', ',', '!', '?', '@', '#', '$', '%', '^', '&', '*', '(', ')']


def tf_idf(tfidf_vectorizer, corpus):
    """
    It is a function to use TF-IDF to vectorize corpus.
    :param corpus: corpus to fit
    :return: vector
    """
    X = tfidf_vectorizer.fit_transform(corpus)
    return X.toarray()


def calculate_cosine_similarity(vec, *args):
    """
    It is a function to calculate cosine similarity.
    :param vec: vector
    :return: cosine similarity result
    """
    return cosine_similarity(vec, args)


def load_GloVe_model(path):
    """
    It is a function to load GloVe model
    :param path: model path
    :return: model array
    """
    print("Load GloVe Model.")
    with open(path, 'r') as f:
        content = f.readlines()
    model = {}
    for line in content:
        splitLine = line.split()
        word = splitLine[0]
        embedding = np.array((splitLine[1:]))
        model[word] = embedding
    print("Done.", len(model), " words loaded!\n")
    return model


def process(sentence, lemmatizer):
    """
    It is a function to
    :param sentence:
    :return:
    """
    res = []
    # 1. Split
    splitSentence = sentence.split()

    # 2. To lower case
    for word in splitSentence:
        word = word.lower()

        # 3. Lemmatize
        word = lemmatizer.lemmatize(word)

        # 4. Remove stop words
        stopword_set = set(stopwords.words("english"))
        if word in stopword_set:
            continue

        # 5. Remove punctuation
        if word in punctuation_list:
            continue

        res.append(word)
    return res


def get_glove_vec(sentence, lemmatizer, glove_model):
    """
    It is a function to get glove vector for cosine similarity calculation.
    :param process_sen: processed sentence
    :param glove_model: GloVe model
    :return: np.mean(process_sen)
    """
    res = []
    for word in process(sentence, lemmatizer):
        try:
            vec = glove_model.get(word).astype(float)
            res.append(np.mean(vec))
        except Exception:
            continue

    return res


def calculate_glove_cosine_similarity(s1, s2, lemmatizer, glove_model):
    """
    It is a function to calculate GloVe embedding cosine similarity.
    :param glove_model: GloVe model
    :return: GloVe cosine similarity
    """
    # 1. Get GloVe Vector
    s1_vec = get_glove_vec(s1, lemmatizer, glove_model)     # <List> object
    s2_vec = get_glove_vec(s2, lemmatizer, glove_model)

    # 2. Measure the length of vector
    try:
        if len(s1_vec) == len(s2_vec):
            s1_array = np.array((s1_vec)).reshape(1, -1)
            s2_array = np.array((s2_vec)).reshape(1, -1)

        elif len(s1_vec) > len(s2_vec):
            s1_array = np.array((s1_vec)).reshape(1, -1)
            s2_array = np.zeros(shape=(1, len(s1_vec)))
            s2_array[0, :len(s2_vec)] = s2_vec
        else:
            s2_array = np.array((s2_vec)).reshape(1, -1)
            s1_array = np.zeros(shape=(1, len(s2_vec)))
            s1_array[0, :len(s1_vec)] = s1_vec

        assert s1_array.shape == s2_array.shape
        s1_mean = np.mean(s1_array, axis=0).reshape(1, -1)
        s2_mean = np.mean(s2_array, axis=0).reshape(1, -1)
        return cosine_similarity(s1_mean, s2_mean)[0][0]

    except Exception as e:
        print(e)


def main():

    corpus = ['The president greets the press in Chicago',
              'Obama speaks to the media in Illinois']
    s1 = 'The president greets the press in Chicago'
    s2 = 'Obama speaks to the media in Illinois'
    s3 = 'I love you'
    s4 = 'We went to Starbucks to buy hazelnut lattee yesterday'
    s5 = 'We often go to Starbucks to buy coffee and chat with each other.!!!!!!'


    ############## 1. TF-IDF ###############
    tf_idf_vectorizer = TfidfVectorizer()
    tf_idf_vec = tf_idf(tfidf_vectorizer=tf_idf_vectorizer,
                        corpus=corpus)
    print("tf_idf_vec = ", tf_idf_vec)
    print("tf_idf_vec.shape = ", tf_idf_vec.shape)
    print("tf_idf_vectorizer.get_feature_names() = ", tf_idf_vectorizer.get_feature_names())


    ##### 2. TF-IDF Cosine Similarity ######
    tfidf_cosine_res = cosine_similarity(tf_idf_vec)[0][1]
    print("tfidf_cosine_res = ", tfidf_cosine_res)
    print("\n")


    ########### 3. Lemmatization ###########
    lemmatizer = WordNetLemmatizer()


    ########## 4. Load GloVe Model #########
    glove_model = load_GloVe_model(glove_path)      # len(glove_model) = 400000


    ###### 5. GloVe Cosine Similarity ######
    res = calculate_glove_cosine_similarity(s1, s3, lemmatizer, glove_model)
    print("res = ", res)

    res1 = calculate_glove_cosine_similarity(s1, s2, lemmatizer, glove_model)
    print("res1 = ", res1)

    res2 = calculate_glove_cosine_similarity(s2, s3, lemmatizer, glove_model)
    print("res2 = ", res2)

    res3 = calculate_glove_cosine_similarity(s5, s4, lemmatizer, glove_model)
    print("res3 = ", res3)


if __name__ == '__main__':

    main()