"""
data_padding.py
description: pre-process data, pads sentances to the samle fixed length, converts
words to vocabulary ID representations, saves the training data in a numpy array
format for later use.
"""
import numpy as np
import pickle

test_data = True;

"""
pre-process training data
combines positive and negative training sempales in one file and generate labels
pad each sentance to the same specified fixed length with padding word: <PAD/>
convert words to vocabulary ID representation
"""
def preprocessData(seq_length, vocab ,padding_word="<PAD/>"):
    combinedData = [];
    labels = [];
    label = [0, 1];
    if padding_word == '<PAD/>':
        padding_word_id = vocab.get(padding_word, -1);
    for fn in ['twitter-datasets/train_pos_full.txt', 'twitter-datasets/train_neg_full.txt']:
        with open(fn) as f:
            for line in f:
                lineData = [];
                sline = line.strip().split(); # split sentance into words
                for word in sline:
                    word_id = vocab.get(word, -1); # get a vocabulary ID representation of a word
                    if word_id != -1:
                        lineData.append(word_id); # append the word IDs
                num_padding = seq_length - len(lineData) # calculate the number of words to pad
                if num_padding>0:
                    new_line = lineData + [padding_word_id] * num_padding # pad the words
                else:
                    new_line = lineData[:seq_length]; # if length of sentence is more then fixed number of words cut off at fixed length
                combinedData.append(new_line); # append processed sentence to the document
                labels.append(label); # append labels

            label = [1, 0];
    return [combinedData, labels];

"""
Same functionality as preprocessData function, but deals with evaluation data
"""
def preprocessTestData(seq_length, vocab ,padding_word="<PAD/>"):
    combinedData = [];
    padding_word_id = vocab.get(padding_word, -1);
    for fn in ['twitter-datasets/test_data.txt']:
        with open(fn) as f:
            for line in f:
                lineData = [];
                lline = line.strip().split(",", 1);
                sline = lline[1].strip().split();
                for word in sline:
                    word_id = vocab.get(word, -1);
                    if word_id != -1:
                        lineData.append(word_id);
                num_padding = seq_length - len(lineData);
                if num_padding>0:
                    new_line = lineData + [padding_word_id] * num_padding
                else:
                    new_line = lineData[:seq_length];
                combinedData.append(new_line);
    return [combinedData];

def main():
    numWords = 31; # fixed number of words per sentance.
    # load vocabulary file generated by the script provided
    with open('modelData/vocab.pkl', 'rb') as f:
        vocab = pickle.load(f)

    if test_data == False:
        data = preprocessData(numWords, vocab);
        X = np.array(data[0]);
        Y = np.array(data[1]);

        print("Shape of combined data: ", X.shape);
        print("Label shape: ", Y.shape)
        # save the training data as a numpy array
        np.save('twitter-datasets/trainX_combined', X);
        np.save('twitter-datasets/labels', Y);
    else:
        data = preprocessTestData(numWords, vocab);
        X = np.array(data[0]);
        np.save("twitter-datasets/testX", X);


if __name__ == '__main__':
    main();
