"""
Author: Prakash Dhimal
George Mason University
CS 584 Theory and Applications of Data Mining
Homework 1

The objective of this assignment are the following:
 - Implement the Nearest Neighbor Classification Algorithm
 - Handle Text Data (Reviews of Amazon Baby Products)
 - Design and Engineer Features from Text Data.
 - Choose the Best Model i.e., Parameters of a Nearest Neighbor Selection, Features and Similarity Functions

 Detailed Description:
 A practical application in e-commerce applications is to infer sentiment (or polarity)
 from free form review text submitted for range of products. For the purposes of this assignment you have to implement
 a k-Nearest Neighbor Classifier to predict the sentiment for 18506 reviews for baby products provided in the test file (test.data).
 Positive sentiment is represented by a review rating and given by +1 and Negative Sentiment is represented by a review rating of -1.

In test.data you are only provided the reviews but no ground truth rating which will be used for comparing your predictions.

Training data consists of 18506 reviews as well and exists in file train_file.dat. Each row begins with the sentiment score followed with a text of the rating.

For Evaluation Purposes (Leaderboard Ranking) we will use the Accuracy Metric comparing the Predictions submitted by you on the test set with the ground truth.

Some things to note:
 - The public leaderboard shows results for 50% of randomly chosen test instances only. This is a standard practice in data mining challenge to avoid gaming of the system.
 - The private leaderboard will be released after the deadline evaluates all the entries in the test set.
 - In a 24-hour cycle you are allowed to submit a prediction file 5 times only.
 - The final ranking will always be based on the last submission.
 - format.dat shows an example file containing 18506 rows alternating with +1 and -1. Your test.dat should be similar to format.dat
   with the same number of rows i.e., 18506 but of course the sentiment score generated by your developed model.

Implement:
  - a k-nearest neighbor classifier to predict the sentiment for 1806 reviews for baby products provided in the test file
    - test file - test.data
    - positive sentiment is represented by a review rating and given by a +1
    - negative sentiment is represented by a review rating of -1
    - there is a train_file.dat with training data?

  - parameters:
    - the number of neighbors (k),
    - the bag-of-words representation (binary vs. raw frequency count vs. TF*IDF),
    - distance/similarity measure (at least two).
"""
import math
import multiprocessing
import string
import time
from collections import Counter

import nltk
import numpy as np
import pandas as pd
from nltk.corpus import stopwords

# from scipy.sparse import csr_matrix

nltk.download('punkt')
nltk.download('stopwords')
stemmer = nltk.PorterStemmer()
stop_words = stopwords.words('english')


def get_training_data():
    # training file
    train_file = "../data/1580449515_4035058_train_file.dat"
    train = pd.read_table(train_file, header=None, skip_blank_lines=False)
    train.columns = ["label", "data"]
    train = train.dropna()  # TODO - should we do this?
    train = train.reset_index(drop=True)
    data_processed = [preprocess(line) for index, line in enumerate(train.data)]
    return data_processed, train.label


def get_raw_test_data():
    # test file
    test_file = "../data/1580449515_4313154_test.dat"
    with open(test_file, "r") as file:
        data = file.readlines()
    return data[:8000]


# todo - trim the data down even more:
def preprocess(document):
    print("Length of the initial document: ", len(document))
    # preserve any stars ratings, people usualy say I give this product 1 star, 2 star, 3 star, 4 star, and 5 star
    if "1 star" in document:
        document = document.replace("1 star", "onestar")
    if "2 star" in document:
        document = document.replace("2 star", "twostars")
    if "3 star" in document:
        document = document.replace("3 star", "threestars")
    if "4 star" in document:
        document = document.replace("4 star", "fourstars")
    if "5 star" in document:
        document = document.replace("5 star", "fivestars")

    tokens = nltk.word_tokenize(document)
    # delete any digits
    tokens = [word for word in tokens if not word.isdigit()]
    # lower case
    tokens = [word.lower() for word in tokens if type(word) == str]
    # remove punctuations
    tokens = [word for word in tokens if word not in string.punctuation]
    # remove stopwords
    tokens = [word for word in tokens if not word in stop_words]
    # remove any words less and 3 in length # todo
    tokens = [word for word in tokens if not len(word) <= 4]
    # stemming
    tokens = [stemmer.stem(word) for word in tokens]
    print("Length of the document after pre-processing: ", len(tokens))

    return tokens


def df_list(document, index, DF_list):
    for token in document:
        try:
            DF_list[token].add(index)
        except:
            DF_list[token] = {index}


def update_DF_list(word, DF_list):
    DF_list[word] = len(DF_list[word])


def document_vector(
        document,
        total_vocabulary,
        N,
        DF_list):
    """
    Returns TF-IDF for each document
    """
    # each document need the length of total vocab
    doc_vector = np.zeros((len(total_vocabulary)))

    counter = Counter(document)
    words_count = len(document)

    for token in np.unique(document):
        tf = counter[token] / words_count
        df = doc_freq(DF_list, token)
        idf = math.log((N + 1) / (df + 1))  # why plus 1?
        try:
            ind = total_vocabulary.index(token)
            doc_vector[ind] = tf * idf
        except:
            pass
    return doc_vector


def doc_freq(DF_matrix, word):
    """
    This method is just for lookup
    """
    frequenxy = 0
    try:
        frequenxy = DF_matrix[word]
    except:
        pass
    return frequenxy


def process_tasks(
        task_queue,
        k,
        total_vocabulary,
        train_matrix,
        train_label,
        test_data_raw,
        DF_list,
        test_y):
    """
    This method is to support multi-processing
    """
    while not task_queue.empty():
        task = task_queue.get()
        label(
            vocabulary=total_vocabulary,
            train_data_matrix=train_matrix,
            train_label=train_label,
            test_document_index=task,  # should have the document index
            test_data_raw=test_data_raw,
            DF_list_train=DF_list,
            test_y=test_y,
            k=k)
    return True


def add_tasks(
        task_queue,
        test_data_list):
    """
    This method is to support multi-processing
    """
    for test_document_index in range(len(test_data_list)):
        task_queue.put(test_document_index)
        print(test_document_index)
    return task_queue


def cosine_simimarity(document_a, document_b):
    # normalized cosine similarity
    return np.dot(document_a, document_b) / (np.linalg.norm(document_a) * np.linalg.norm(document_b))


def label(
        vocabulary,
        train_data_matrix,  # This is a TF-IDF matrix
        train_label,
        test_document_index,
        test_data_raw,  # this is the raw test data, not pre-processed yet
        DF_list_train,
        test_y,
        k):
    test_raw = test_data_raw[test_document_index]

    test_processed = preprocess(test_raw)

    # Generate the TF_IDF for the test instance
    test_TF_IDF_vector = document_vector(document=test_processed,
                                         N=len(train_label) + 1,  # because we are adding 1 more?
                                         DF_list=DF_list_train,
                                         total_vocabulary=vocabulary)

    print("Finding label for test item ", test_document_index)

    # cosine similarity
    distances = [(index, (cosine_simimarity(test_TF_IDF_vector, doc_vector))) for index, doc_vector in
                 enumerate(train_data_matrix)]

    # get k-nearest neighbors
    # np.array(distances).sort()[::-1]
    # the closer the documents are by angle, the higher is the Cosine Similarity
    distances.sort(key=lambda x: x[1])
    neighbor_indices = distances[::-1]
    neighbor_indices = neighbor_indices[:k]

    label_list = []
    for value in neighbor_indices:
        label_list.append(train_label[value[0]])

    # get a list of labels from the neighbors and sum the list
    # labels_list = train_label[neighbor_indices].tolist()
    label_sum = sum(label_list)

    # classify based on label_sum
    if label_sum > 0:
        label_this = +1
    else:
        label_this = -1
        # print(test_data[test_document_index])
        # print("labeled as : ", label_this)

    test_y[test_document_index] = label_this


def k_nearest_neighbor(
        k,
        vocabulary,
        test_data_raw,
        DF_list,
        train_matrix,
        train_label):
    test_y = multiprocessing.Array('i', len(test_data_raw))

    num_cpus = multiprocessing.cpu_count() - 2
    queue = multiprocessing.Queue()
    full_task_queue = add_tasks(
        queue,
        test_data_raw)

    processes = []
    print(f'Running with {num_cpus} processes!')
    start = time.time()
    for n in range(num_cpus):
        process = multiprocessing.Process(
            target=process_tasks, args=(
                full_task_queue,
                k,
                vocabulary,
                train_matrix,
                train_label,
                test_data_raw,
                DF_list,
                test_y,))
        processes.append(process)
        process.start()
    for process in processes:
        process.join()
    print(f'Time taken = {time.time() - start:.10f}')

    review_ratings = []
    for rating in test_y:
        if rating is 1:
            review_ratings.append("+1")
        else:
            review_ratings.append("-1")
    return review_ratings


##########################################################################
print("INFO: reading, pre-processing training file...")
training_data, training_labels = get_training_data()
# Raw test data
print("INFO: reading the test file!")
raw_test_data = get_raw_test_data()

print("INFO: Length of training data: ", len(training_data))
print("INFO: Length of test data: ", len(raw_test_data))
############################################################################
# vocab_training = []
# [vocab(document, vocab_training) for document in training_data]
# bag_of_word_matrix_training = [bag_of_words(document, vocab_training) for document in training_data]

############################################################################

print("INFO: Generating DF-IDF matrix for the training data set...")
DF_list_training = {}
# this is going to update the DF_
[df_list(document, index, DF_list_training) for index, document in enumerate(training_data)]
# again updates the list?
[update_DF_list(word, DF_list_training) for word in DF_list_training]
vocab_training = [word for word in DF_list_training]

############################################################################


N = len(training_data)
tf_idf_matrix_training = [document_vector(
    document,
    vocab_training,
    N,
    DF_list_training) for document in training_data]

# mat_sparse = csr_matrix(tf_idf_matrix_training)
# mat_sparse.sort_indices()

print("INFO: Starting K-NN...")
k = 25
# returns a set of results
ratings = k_nearest_neighbor(
    k=k,
    vocabulary=vocab_training,
    test_data_raw=raw_test_data,
    DF_list=DF_list_training,
    train_matrix=tf_idf_matrix_training,
    train_label=np.array(training_labels))
print("INFO: K-nearest neighbor done!")
print("INFO: Writing results")

# writing to .txt
fileName = "../output/dhimalp-k" + str(k) + "-train-" + str(len(training_data)) + "-test-" + str(
    len(raw_test_data)) + "-" + str(int(time.time())) + ".dat"

np.savetxt(fileName, ratings, fmt='%s')
print("INFO: Done writing results!")
print("Check ", fileName)
