from pygoose import *

import nltk

"""
Explanation: Preprocessing: Clean Up & Tokenize Questions
Break question titles into tokens, and perform token-level normalization: expand shortened words, correct spelling, etc.
Imports
This utility package imports numpy, pandas, matplotlib and a helper kg module into the root namespace.
End of explanation
"""


project = kg.Project.discover()

"""
Explanation: Config
Automatically discover the paths to various data folders and compose the project structure.
End of explanation
"""


df_train = pd.read_csv(project.data_dir + 'train.csv').fillna('none')
df_test = pd.read_csv(project.data_dir + 'test.csv').fillna('none')

df_all = pd.concat([df_train, df_test])

"""
Explanation: Load Data
Original question datasets.
End of explanation
"""


stopwords = set(kg.io.load_lines(project.aux_dir + 'stopwords.vocab'))

"""
Explanation: Stopwords customized for Quora dataset.
End of explanation
"""


spelling_corrections = kg.io.load_json(project.aux_dir + 'spelling_corrections.json')

"""
Explanation: Pre-composed spelling correction dictionary.
End of explanation
"""


tokenizer = nltk.tokenize.RegexpTokenizer(r'\w+')

"""
Explanation: Load Tools
End of explanation
"""


def translate(text, translation):
    for token, replacement in translation.items():
        text = text.replace(token, ' ' + replacement + ' ')
    text = text.replace('  ', ' ')
    return text

def spell_digits(text):
    translation = {
        '0': 'zero',
        '1': 'one',
        '2': 'two',
        '3': 'three',
        '4': 'four',
        '5': 'five',
        '6': 'six',
        '7': 'seven',
        '8': 'eight',
        '9': 'nine',
    }
    return translate(text, translation)

def expand_negations(text):
    translation = {
        "can't": 'can not',
        "won't": 'would not',
        "shan't": 'shall not',
    }
    text = translate(text, translation)
    return text.replace("n't", " not")

def correct_spelling(text):
    return ' '.join(
        spelling_corrections.get(token, token)
        for token in tokenizer.tokenize(text)
    )

def get_question_tokens(question, lowercase=True, spellcheck=True, remove_stopwords=True):
    if lowercase:
        question = question.lower()
    
    if spellcheck:
        question = correct_spelling(question)
    
    question = spell_digits(question)
    question = expand_negations(question)

    tokens = [token for token in tokenizer.tokenize(question.lower() if lowercase else question)]    
    if remove_stopwords:
        tokens = [token for token in tokens if token not in stopwords]
    
    tokens.append('.')
    return tokens

def get_question_pair_tokens_spellcheck(pair):
    return [
        get_question_tokens(pair[0], lowercase=False, spellcheck=True, remove_stopwords=False),
        get_question_tokens(pair[1], lowercase=False, spellcheck=True, remove_stopwords=False),
    ]

def get_question_pair_tokens_lowercase_spellcheck(pair):
    return [
        get_question_tokens(pair[0], lowercase=True, spellcheck=True, remove_stopwords=False),
        get_question_tokens(pair[1], lowercase=True, spellcheck=True, remove_stopwords=False),
    ]

def get_question_pair_tokens_lowercase_spellcheck_remove_stopwords(pair):
    return [
        get_question_tokens(pair[0], lowercase=True, spellcheck=True, remove_stopwords=True),
        get_question_tokens(pair[1], lowercase=True, spellcheck=True, remove_stopwords=True),
    ]

"""
Explanation: Preprocess and tokenize questions
End of explanation
"""


tokens_spellcheck = kg.jobs.map_batch_parallel(
    df_all.as_matrix(columns=['question1', 'question2']),
    item_mapper=get_question_pair_tokens_spellcheck,
    batch_size=1000,
)

"""
Explanation: Tokenize the questions, correct spelling, but keep the upper/lower case.
End of explanation
"""


tokens_lowercase_spellcheck = kg.jobs.map_batch_parallel(
    df_all.as_matrix(columns=['question1', 'question2']),
    item_mapper=get_question_pair_tokens_lowercase_spellcheck,
    batch_size=1000,
)

"""
Explanation: Tokenize the questions, convert to lowercase and correct spelling, keep the stopwords (useful for neural models).
End of explanation
"""


tokens_lowercase_spellcheck_no_stopwords = kg.jobs.map_batch_parallel(
    df_all.as_matrix(columns=['question1', 'question2']),
    item_mapper=get_question_pair_tokens_lowercase_spellcheck_remove_stopwords,
    batch_size=1000,
)

"""
Explanation: Just as before, but also with stopwords removed.
End of explanation
"""


vocab = set()
for question in progressbar(np.array(tokens_lowercase_spellcheck).ravel()):
    for token in question:
        vocab.add(token)

vocab_no_stopwords = vocab - stopwords

"""
Explanation: Extract question vocabulary
End of explanation
"""


kg.io.save(
    tokens_spellcheck[:len(df_train)],
    project.preprocessed_data_dir + 'tokens_spellcheck_train.pickle'
)
kg.io.save(
    tokens_spellcheck[len(df_train):],
    project.preprocessed_data_dir + 'tokens_spellcheck_test.pickle'
)

kg.io.save(
    tokens_lowercase_spellcheck[:len(df_train)],
    project.preprocessed_data_dir + 'tokens_lowercase_spellcheck_train.pickle'
)
kg.io.save(
    tokens_lowercase_spellcheck[len(df_train):],
    project.preprocessed_data_dir + 'tokens_lowercase_spellcheck_test.pickle'
)

kg.io.save(
    tokens_lowercase_spellcheck_no_stopwords[:len(df_train)],
    project.preprocessed_data_dir + 'tokens_lowercase_spellcheck_no_stopwords_train.pickle'
)
kg.io.save(
    tokens_lowercase_spellcheck_no_stopwords[len(df_train):],
    project.preprocessed_data_dir + 'tokens_lowercase_spellcheck_no_stopwords_test.pickle'
)

"""
Explanation: Save preprocessed data
Tokenized questions.
End of explanation
"""


kg.io.save_lines(
    sorted(list(vocab)),
    project.preprocessed_data_dir + 'tokens_lowercase_spellcheck.vocab'
)

kg.io.save_lines(
    sorted(list(vocab_no_stopwords)),
    project.preprocessed_data_dir + 'tokens_lowercase_spellcheck_no_stopwords.vocab'
)

"""
Explanation: Question vocabulary.
End of explanation
"""


kg.io.save(df_train['is_duplicate'].values, project.features_dir + 'y_train.pickle')

"""
Explanation: Ground truth.
End of explanation
"""


