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

__all__ = ['process_sentence', 'process_doc', 'decontracted', 'get_stop_words', 'token_filter', 'get_tokens',
           '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 process_sentence(doc_id:int, sent_id:int, sent_tokens:list, sent_pos:list, raw_sent:str):
    "Pre-process a single sentence"
    row_id = f"doc.{doc_id}.sent.{sent_id}"
    sent_text = raw_sent[sent_id]
    tokens = sent_tokens[sent_id]
    pos_tags = sent_pos[sent_id]
    data_row = [row_id, doc_id, sent_id, sent_text, tokens, pos_tags]
    return data_row

def process_doc(doc_id:int, text:str, stop_words:list):
    "Pre-process a single document"

    # Pre-process document into cleaned sentences
    sent_tokens, sent_pos, raw_sent = tokenize_and_stem(text, stop_words)

    # Populate a row in data for each sentence in the document
    data_rows = [process_sentence(doc_id, sent_id, sent_tokens, sent_pos, raw_sent) for sent_id in range(len(sent_tokens))]

    # Populate a row in doc_df for each file loaded
    doc_tokens = [token for sent in sent_tokens for token in sent]
    doc_row = [doc_id, None, text, doc_tokens]

    return doc_row, data_rows

# 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"),
        (r"/", " ")
    ]

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

    return text

# Cell
def get_stop_words(stop_words_file:str=None):
    "Gets a list of all stop words"
    # 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')

    stop_words = set(stopwords.words('english')) | set(custom_stop_words)
    return stop_words

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

def get_tokens(text:string):
    """
    Gets list of tokens from string

    Returns list
    """
    table = str.maketrans(' ', ' ', string.punctuation+"“"+"”")
    return WhitespaceTokenizer().tokenize(decontracted(text).translate(table))

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

    Returns (list, list, list)
    """

    lemmatizer = nltk.WordNetLemmatizer()

    # Split into sentences and condense whitespace
    sent_raw = [re.sub("\s+", " " ,s).strip() for s in nltk.sent_tokenize(text)]
    sent_tokens = []
    sent_pos = []

    # Tokenize sentences
    tokenized = np.array([get_tokens(s) for s in sent_raw])

    # Part of Speech tags
    pos_tags = [[p[1] for p in nltk.pos_tag(sent)] for sent in tokenized]

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

    # Create filter to remove invalid tokens
    filter_mask = np.array([[token_filter(token, stop_words) for token in sent] for sent in lemma])

    # Remove invalid tokens
    sent_tokens = [np.array(lemma[i])[filter_mask[i]] for i in range(len(lemma))]
    sent_pos = [np.array(pos_tags[i])[filter_mask[i]] for i in range(len(pos_tags))]

    return sent_tokens, sent_pos, sent_raw