###############################################################################
# Assignment 2 - WSD    decision-list                                         #
# AIT-590, Summer 2021                                                        #
# Group 3: Fernando, Melissa, Archer                                          #
# July 4, 2021                                                                #
#                                                                             #
# Assuming you have python 3                                                  #
# TO EXECUTE RUN:                                                             #
#  python decision-list.py line-train.xml line-test.xml my-decision-list.txt  #
###############################################################################
#
# This is an application that will work to do Word Sense Desambiguation for a specific word,
# on this case specifically the word "line".
# The application will receive training data from the file line-train.xml which contains multiple instances of sentences
# each of these sentences in the context are attributed to a specific sense "phone" or "product"
#
# We will use the Input data to create a decision List based on Yarowsky Decision list presented in class.
# The decision list will work based on the collocation of words [-2, -1, 1] of the Ambiguous word Line.
# after calculating the likelihood of the collocated words we will use that to try to predict the correct sense of the
# line-test.xml data which does not have the sense attributed to it.
# If no sense is found on Decision List, we will use the sense with higherst probability from the training data
# based on counts.
#
#
# After we are able to find the sense classifications for the line-test.xml we will run our scorer.py application to
# evaludate the results.
#
# The scorer.py application will take the answers from the main application and the line-answers.txt file and compare
# how good were our predictions and output the accuracy of our model.
# then it will output a confusion matrix to illustrate the performance.
#
# The expected output of current application is as follows:
#
# First use the inputs to create the decision list
# decision list likelihood is created based on the ConditionalFreqDistribution
# which is used in the ConditionalProbDist to get the likelihood of each collocated item in the collocation table.
# Decision LIst:
# {'position': '-1W telephone line', 'likelihood': 6.149747119504682, 'classification': 'phone'}
# {'position': '-1W car line', 'likelihood': 4.392317422778761, 'classification': 'product'}
# {'position': '1W line rose', 'likelihood': 2.3219280948873626, 'classification': 'product'}
# {'position': '-2W recently introduced line', 'likelihood': 2.8073549220576046, 'classification': 'product'}
#
#
# After we have the decison list we lookup the sense probability based on the max likelihood of collocations.
# and output in the same format as the line-answers.txt for easier scoring later on.
# Predicted answers:
# <answer instance="line-n.w8_059:8174:" senseid="phone"/>
# <answer instance="line-n.w7_098:12684:" senseid="phone"/>
# <answer instance="line-n.w8_119:16927:" senseid="product"/>
# <answer instance="line-n.w8_008:13756:" senseid="product"/>
#
# We read both line-answers.txt and my-line-answers.txt as input on the scorer to compare our results and find
# the accuracy of our predictions.
#
# Output from Scorer.py:
#
# Our Decision List Accuracy: 76.98 %
#
#  ----- Confusion Matrix -----
#         |     p |
#         |     r |
#         |  p  o |
#         |  h  d |
#         |  o  u |
#         |  n  c |
#         |  e  t |
# --------+-------+
#   phone |<49>23 |
# product |  6<48>|
# --------+-------+

import math
import sys
import scorer

from bs4 import BeautifulSoup
from nltk import RegexpTokenizer, ConditionalFreqDist, ConditionalProbDist, ELEProbDist
from nltk.corpus import stopwords

# Word to disabiguate
a_word = "line"

# Collocation range used for future on decision List creation as well as testing data.
collocation_windows = [-2, -1, 1, 2]


# Read Arguments from the input on the required order
def read_arguments():
    """
    This function will read the first 3 arguments of the application
    and return as 3 variables with the file paths.
    :return: line-train.xml line-test.xml my-decision-list.txt
    """
    return sys.argv[1], sys.argv[2], sys.argv[3]


# Read XML files
def read_xml(xml_file):
    """
    This function will receive a file path as input
    read as xml file with BeautifulSoup and return as an object
    :param xml_file:
    :return: BeautifulSoup Object.
    """
    return BeautifulSoup(open(xml_file, 'r'), 'xml')


# normalize_and_tokenize function tokenize text and set them in lowercase
def normalize_and_tokenize(text):
    """
    This function will perform a few steps:
    - Lower the text and call the word_tokenizer(text) to retrieve the tokenized lower case text
    - Replace all occurences of "lines" to "line"
    - Tokenize (The tokenization process uses a regex tokenized to retrieve only words.
    - Remove stop words
    :param text:
    :return: lowered tokenized text.
    """
    text = text.lower()
    text = text.replace("lines", "line")
    tokens = word_tokenize(text)
    tokens = [token for token in tokens if token not in stopwords.words('english')]
    return tokens


def word_tokenize(text):
    """Work Tokenizer
        This Tokenizer will user a RegexpTokenizer \w+
        to identify words from the sentences, this will exclude punctuations and special characters while tokenizing words.

    :param text:
    :return: tokenized text
    """
    tokenizer = RegexpTokenizer(r'\w+')
    return tokenizer.tokenize(text)


def get_word_index(word_list, word):
    """
    This function will receive a list of words and a word to look up
    it will find the index where the word is present and if found will return the index on the list.
    if not found it returns none.
    :param word_list:
    :param word:
    :return: word index in the list.
    """
    try:
        index = word_list.index(word)
    except ValueError:
        index = None
        print(f'Warning from get_word_index: Did not find word_list.index(word) for {word}')
    return index


def create_training_dict_from_xml(xml):
    """
    This function creates a dictionary based on a XML input with the BeautifulSoup Object
    it will parse all elements from the input XML object into a List of dictionaries
    of the following format:
    [{
            'id': 'String',
            'sense': 'String'
            'tokens': []
        }]
    :param xml:
    :return: list of dict
    """
    training_dict = list()
    for instance in xml.find_all("instance"):
        tokens = list()

        for sentence in instance.context.find_all('s'):
            tokens.extend(normalize_and_tokenize(sentence.get_text()))

        line = {
            'id': instance['id'],
            'sense': instance.answer['senseid'] if instance.answer else 'to_predict',
            'tokens': tokens
        }

        training_dict.append(line)
    return training_dict


def create_collocation_distribution(training_dict):
    """
    using the training dictionary create a collocation distribution table for using later to create decision list.
    this is referent on the collocation distribution from yarowski article.
    :param training_dict:
    :return:
    """
    # Set the collocation range
    # collocation_windows = [-2, -1, 1, 2]

    # create a ConditionalFreqDist to Measure Collocational Distribution
    collocation_distribution = ConditionalFreqDist()

    # Starting Collocational distribution counts
    # for each instance of the training data
    for instance in training_dict:
        sense, tokens = instance['sense'], instance['tokens']
        # find the index of the ambiguous word "a_word"
        index = get_word_index(tokens, a_word)

        # if word exists in list of tokens
        if index:
            # start processing the collocation collocation_windows
            for window in collocation_windows:
                # for each window (-2, -1, +1, +2) define a range
                w_range = range(window, 0) if window < 0 else range(1, window + 1)

                # for each window(w) in the window range find the context and append to the collocation.
                # range is used just to make it easier to travel the collocation_windows
                collocation = ''
                # +/- k references on Step 3 of Yarowski article.
                for k in w_range:
                    context_index = index + k
                    if 0 <= context_index <= len(tokens) - 1:
                        collocation = collocation + tokens[context_index] + ' '
                    else:
                        pass
                        # print('notInList')

                # decide to append real a_word depending on the window
                if window < 0:
                    collocation = collocation + str(tokens[index])
                else:
                    collocation = str(tokens[index]) + ' ' + collocation

                # add to the freq_distribution
                # NOT SURE OF THIS IF WE NEED TO HAVE BOTH +/- W with collocation
                # collocation_distribution[str(window) + 'W ' + collocation][sense] += 1
                collocation_distribution[str(window) + 'W ' + collocation.rstrip()][sense] += 1
    return collocation_distribution


def create_decision_list_from_conditinal_prob(cpdist):
    """
    Step 4 of the Yarowski Article, we use the conditional probability of the collocation table
    to create the decision list.
    :param cpdist:
    :return: decision list
    """
    decision_list = []
    # calculate log likelihood
    for position in cpdist.conditions():
        # print(position)
        ratio = cpdist[position].prob('phone') / cpdist[position].prob('product')
        if ratio == 0:
            likelihood = 0
        else:
            likelihood = math.log(ratio, 2)
        classification = 'phone' if ratio >= 1 else 'product'
        decision_list.append({'position': position, 'likelihood': abs(likelihood), 'classification': classification})

    return sorted(decision_list, key=lambda i: i['likelihood'], reverse=True)


def lookup_decision_list(decision_list, value):
    """
    this function is used to lookup the collocation in the decision list to identify the likelihood
    of the collocated term.
    :param decision_list:
    :param value:
    :return:
    """
    return list(filter(lambda dl: dl['position'] in [value], decision_list))


#####################################################
#    Main Execution for Decision List stats here    #
#####################################################
def main():
    print("Starting Decision List")
    line_train, line_test, decision_list_output = read_arguments()

    # Read the XML input line_train.xml file
    xml = read_xml(line_train)

    # Create training dict from training data xml
    training_dict = create_training_dict_from_xml(xml)

    # create a ConditionalFreqDist to Measure Collocational Distribution
    collocation_distribution = create_collocation_distribution(training_dict)

    cpdist = ConditionalProbDist(collocation_distribution, ELEProbDist, 10)

    # create decision list using ConditionalProbDist as on the hints and tips.
    decision_list = create_decision_list_from_conditinal_prob(cpdist)

    # write decision list to file
    # type(decision_list) is a list, each element containing a dict like this:
    # {'position': '-2W man end line', 'likelihood': 2.321928094887362, 'classification': 'phone'}
    with open(decision_list_output, 'w') as f:
        for eachline in decision_list:
            f.write(str(eachline))
            f.write('\n')

    ##########################
    # START TESTING DATA     #
    ##########################
    xml_test = read_xml(line_test)
    testing_dict = create_training_dict_from_xml(xml_test)  # a list

    # for each instance of the training data
    unknown_sense = 'product'
    count_phone = 0
    count_product = 0
    # Find the most common sense to be used in case we dont find anything on decision list.
    for instance in training_dict:
        if instance['sense'] == 'phone':
            count_phone += 1
        else:
            count_product += 1

    # unknown_sense will be default sense, when sense is ambiguous
    if count_phone > count_product:
        unknown_sense = 'phone'

    # initiating list of results
    testing_results = []

    # for each instance in the testing dataset.
    for instance in testing_dict:
        sense_lookups = []
        max_likelihood = {}
        instance_id, tokens = instance['id'], instance['tokens']
        index = get_word_index(tokens, a_word)
        # if word exists in list of tokens
        if index:
            # start processing the collocation collocation_windows
            for window in collocation_windows:
                # for each window (-2, -1, +1, +2) define a range
                w_range = range(window, 0) if window < 0 else range(1, window + 1)

                # for each window(w) in the window range find the context and append to the collocation.
                # range is used just to make it easier to travel the collocation_windows
                collocation = ''
                # +/- k references on Step 3 of Yarowski article.
                for k in w_range:
                    context_index = index + k
                    if 0 <= context_index <= len(tokens) - 1:
                        collocation = collocation + tokens[context_index] + ' '
                    else:
                        pass
                        # print('notInList')

                # decide to append real a_word depending on the window
                if window < 0:
                    collocation = collocation + str(tokens[index])
                else:
                    collocation = str(tokens[index]) + ' ' + collocation
                position = str(window) + 'W ' + collocation.rstrip()

                # Lookup the likelihood of the item based on the window positon for the collocation.
                likelihood = lookup_decision_list(decision_list, position)

                if max_likelihood == {} and len(likelihood) > 0:
                    # When likelihood is found at the first time set as the max_likelihood
                    max_likelihood = likelihood[0]
                elif max_likelihood == {} and len(likelihood) == 0:
                    # When likelihood is not found and it is the first time, set as unknown.
                    max_likelihood = {'position': position, 'likelihood': 0, 'classification': unknown_sense}
                elif len(likelihood) > 0:
                    # When likelihood is found, check if it is greater than the previous found then set as the max.
                    for lkl in likelihood:
                        max_likelihood = lkl if lkl['likelihood'] > max_likelihood['likelihood'] else max_likelihood

            # for each instance on the test data, lookup the likelihood from the decision table for collocation +/- k
            # at the end of the collection, choose the item with the highest likelihood.
            # if no item is found in decision table, use the sense with highest probability.

            testing_results.append({'id': instance_id, 'sense': max_likelihood['classification']})

    # validate our test output:
    warn_str = f"WARNING: testing_results has {len(testing_results)} elements, expected {len(testing_dict)}"
    try:
        # moved this into a try/except so we don't die in some special case during grading
        assert len(testing_results) == len(testing_dict), warn_str
    except:
        print(warn_str)

    # In order to use the expected "script" command from the Assignment instruction we dont run things automatically
    # run scorer.py
    # display confusion matrix

    for item in testing_results:
        result = f'<answer instance="{item["id"]}" senseid="{item["sense"]}"/>'
        print(result)


if __name__ == "__main__":
    main()
