# AUTOGENERATED! DO NOT EDIT! File to edit: preprocessing.ipynb (unless otherwise specified).

__all__ = ['decontracted', 'token_filter', 'tokenize_and_stem']

# Cell
import nltk
from nltk.tokenize import WhitespaceTokenizer
from nltk.corpus import stopwords
import numpy as np
import string
import re
import os

# Cell
def decontracted(text:str):
    """
    Removes contractions from input `text`.

    Credit: https://stackoverflow.com/questions/19790188/expanding-english-language-contractions-in-python

    Returns str
    """
    contractions = [
        (r"won\'t", "will not"),
        (r"can\'t", "can not"),
        (r"hadn\'t", "had not"),
        (r"doesnt", "does not"),
        (r"youre", "you are"),
        (r"dont", "do not"),
        (r"im\s", "i am"),
        (r"ive\s", "i have"),
        (r"won\'t", "will not"),
        (r"can\'t", "can not"),
        (r"hadn\'t", "had not"),
        (r"dont\s", "do not"),
        (r"n\'t", " not"),
        (r"\'re", " are"),
        (r"\'s", " is"),
        (r"\'d", " would"),
        (r"\'ll", " will"),
        (r"\'t", " not"),
        (r"\'ve", " have"),
        (r"\'m", " am")
    ]

    # Replace any contractions with their decontracted equivalent
    for contraction, decontracted in contractions:
        text = re.sub(contraction, decontracted, text)

    text = re.sub(r"/", " ", text) ### I added this line (Amy)
    return text

# Cell
def token_filter(token:str, stop_words:list):
    """
    Checks if a token should be kept (True) or removed during pre-processing (i.e stop words, non-alpha tokens)

    Returns bool
    """
    match = re.match('[A-z]', token)
    return match is not None and token not in stop_words

# Cell
def tokenize_and_stem(text:str, stop_words_file:str=None):
    """
    Parse out sentences, remove contractions, tokenize by white space, remove all punctuation,
    lemmatize tokens, and get parts of speech.

    Returns (list, list, list)
    """

    # Load custom stop words from file if given
    custom_stop_words = []
    if stop_words_file is not None:
        if os.path.isfile(stop_words_file):
            with open(stop_words_file, encoding="utf-8") as file:
                custom_stop_words = file.read().strip().split('\n')
        else:
            print(f'The stop words file {stop_words_file} does not exist. Ignoring and using defaults.')

    lemmatizer = nltk.WordNetLemmatizer()
    stop_words = set(stopwords.words('english')) | set(custom_stop_words)
    table  = str.maketrans(' ', ' ', string.punctuation+"“"+"”")
    sent = nltk.sent_tokenize(text)
    sent_tokens = []
    sent_pos = []
    sent_raw = []

    #For each sentence in document get back the list of tokenized words with contractions normalized and punctuation removed
    for s in sent:
        # Condense whitespace
        s = re.sub("\s+", " " ,s).strip()

        sent_raw.append(s)
        tokenized = WhitespaceTokenizer().tokenize(decontracted(s).translate(table))

        # Part of Speech tags
        pos_tags = nltk.pos_tag(tokenized)

        # Lemmatize and convert to lowercase
        lemma = [lemmatizer.lemmatize(t).lower() for t in tokenized]

        # Create a filter to remove stopwords (including some punctuation) and tokens with no alpha characters
        filter_mask = [token_filter(token, stop_words) for token in lemma]

        # Add filtered list of tokens and PoS tags
        filtered_tokens = list(np.array(lemma)[filter_mask])
        filtered_pos = list(np.array(pos_tags)[filter_mask])
        sent_tokens.append(filtered_tokens)
        sent_pos.append(filtered_pos)

    return sent_tokens, sent_pos, sent_raw