# reader.py
# ---------------
# Licensing Information:  You are free to use or extend this project for
# educational purposes provided that (1) you do not distribute or publish
# solutions, (2) you retain this notice, and (3) you provide clear
# attribution to the University of Illinois at Urbana-Champaign
"""
This file is responsible for providing functions for reading the files
"""
from nltk.stem.porter import PorterStemmer
from nltk.tokenize import RegexpTokenizer
import pandas as pd
import tqdm

porter_stemmer = PorterStemmer()
tokenizer = RegexpTokenizer(r'\w+')
bad_words = {'aed','oed','eed'} # these words fail in nltk stemmer algorithm


def loadDataset(filename, subsets, stemming=False):
    '''
    Load a file, and returns a list of texts.

    Parameters:
    filename (str): the .parquet file containing the data
    subsets (list of str): List of subsets
                           (in 'acad', 'blog', 'news', 'fic', 'spok', 'tvm') to load
    stemming (bool, optional): if True, use NLTK's stemmer to remove suffixes
    
    Output:
    texts (list of lists): texts[m][n] is the n'th word in the m'th text in the dataset
    count (int): number of files loaded
    '''
    texts = []
    data = pd.read_parquet(filename)
    for itm in subsets:
        if itm not in ['acad', 'blog', 'news', 'fic', 'spok', 'tvm']:
            raise RuntimeException(f'{itm} is not a valid subset')
        data_filtered = data[data['doc_id'].str.split('_').str[0] == itm]    
        for _, _, text in data_filtered.itertuples():
            texts.append(tokenizer.tokenize(text.lower()))
    if stemming:
        for text in tqdm.tqdm(texts):
            for i in range(len(text)):
                if text[i] in bad_words:
                    continue
                text[i] = porter_stemmer.stem(text[i])
    return texts
    
def val_set_helper(dataset):
    """
    Splits a dataset into a labels list and a texts list

    Parameters:
    dataset (dict of string->list of lists): dictionary mapping label to list of texts (each text is a list)
    
    Output:
    labels (list): labels[m] is the label of the m'th text in the dataset
    texts (list of lists): texts[m][n] is the n'th word in the m'th text in the dataset
    """
    labels = []
    texts = []
    for label, text_list in dataset.items():
        labels.extend([label] * len(text_list))
        texts.extend(text_list)
    return labels, texts


