__author__ = 'Hao'


from keras.datasets import imdb
import numpy as np
import re
import itertools
from collections import Counter
import cPickle
import numpy as np


def clean_str(string):
    """
    Tokenization/string cleaning for all datasets except for SST.
    Original taken from https://github.com/yoonkim/CNN_sentence/blob/master/process_data.py
    """
    string = re.sub(r"[^A-Za-z0-9(),!?\'\`]", " ", string)
    string = re.sub(r"\'s", " \'s", string)
    string = re.sub(r"\'ve", " \'ve", string)
    string = re.sub(r"n\'t", " n\'t", string)
    string = re.sub(r"\'re", " \'re", string)
    string = re.sub(r"\'d", " \'d", string)
    string = re.sub(r"\'ll", " \'ll", string)
    string = re.sub(r",", " , ", string)
    string = re.sub(r"!", " ! ", string)
    string = re.sub(r"\(", " \( ", string)
    string = re.sub(r"\)", " \) ", string)
    string = re.sub(r"\?", " \? ", string)
    string = re.sub(r"\s{2,}", " ", string)
    return string.strip().lower()


def load_data_and_labels(folder_name):
    """
    Loads MR polarity data from files, splits the data into words and generates labels.
    Returns split sentences and labels.
    """
    # Load data from files
    positive_examples = list(open(folder_name+"rt-polarity.pos").readlines())
    positive_examples = [s.strip() for s in positive_examples]
    negative_examples = list(open(folder_name+"rt-polarity.neg").readlines())
    negative_examples = [s.strip() for s in negative_examples]
    # Split by words
    x_text = positive_examples + negative_examples
    x_text = [clean_str(sent) for sent in x_text]
    x_text = [s.split(" ") for s in x_text]
    # Generate labels
    positive_labels = [1 for _ in positive_examples]
    negative_labels = [0 for _ in negative_examples]
    y = np.concatenate([positive_labels, negative_labels], 0)
    return [x_text, y]

def build_vocab(sentences):
    """
    Builds a vocabulary mapping from word to index based on the sentences.
    Returns vocabulary mapping and inverse vocabulary mapping.
    """
    # Build vocabulary
    word_counts = Counter(itertools.chain(*sentences))
    # Mapping from index to word
    vocabulary_inv = [x[0] for x in word_counts.most_common()]
    # Mapping from word to index
    vocabulary = {x: i+1 for i, x in enumerate(vocabulary_inv)}
    return [vocabulary, vocabulary_inv]

def vocab_to_word2vec(fname, vocab, k=300):
    """
    Load word2vec from Mikolov
    """
    word_vecs = {}
    with open(fname, "rb") as f:
        header = f.readline()
        vocab_size, layer1_size = map(int, header.split())
        binary_len = np.dtype('float32').itemsize * layer1_size
        for line in xrange(vocab_size):
            word = []
            while True:
                ch = f.read(1)
                if ch == ' ':
                    word = ''.join(word)
                    break
                if ch != '\n':
                    word.append(ch)
            if word in vocab:
               word_vecs[word] = np.fromstring(f.read(binary_len), dtype='float32')
            else:
                f.read(binary_len)
    print str(len(word_vecs))+" words found in word2vec."

    #add unknown words by generating random word vectors
    count_missing = 0
    for word in vocab:
        if word not in word_vecs:
            word_vecs[word] = np.random.uniform(-0.25, 0.25, k)
            count_missing+=1
    print str(count_missing)+" words not found, generated by random."
    return word_vecs

def build_word_embedding_mat(word_vecs, vocabulary_inv, k=300):
    """
    Get the word embedding matrix, of size(vocabulary_size, word_vector_size)
    ith row is the embedding of ith word in vocabulary
    """
    vocab_size = len(vocabulary_inv)
    embedding_mat = np.zeros(shape=(vocab_size+1, k), dtype='float32')
    for idx in range(len(vocabulary_inv)):
        embedding_mat[idx+1] = word_vecs[vocabulary_inv[idx]]
    print "Embedding matrix of size "+str(np.shape(embedding_mat))
    #initialize the first row,
    embedding_mat[0]=np.random.uniform(-0.25, 0.25, k)
    return embedding_mat


def build_input_data(sentences, labels, vocabulary):
    """
    Maps sentencs and labels to vectors based on a vocabulary.
    """
    x = [[vocabulary[word] for word in sentence] for sentence in sentences]
    y = np.array(labels)
    return [x, y]


def transform_text(text, vocabulary):
    """
    Transform text to list index that can be fed into CNN and LSTM
    :param text:
    :param vocabulary: {word: index} mapping
    :return:
    """
    words = clean_str(text).split(' ')
    print words
    words = [vocabulary[i] for i in words if i in vocabulary]
    return words




if __name__ == "__main__":
    sentences, labels = load_data_and_labels()
    print str(len(sentences)) + " sentences read"
    vocabulary, vocabulary_inv = build_vocab(sentences)
    print "Vocabulary size: "+str(len(vocabulary))


    print transform_text("I Like you, girl! hello...", vocabulary)

    word2vec = vocab_to_word2vec("../data/GoogleNews-vectors-negative300.bin", vocabulary)
    embedding_mat = build_word_embedding_mat(word2vec, vocabulary_inv)
    x, y = build_input_data(sentences, labels, vocabulary)
    cPickle.dump([x, y, embedding_mat], open('../data/train_mat.pkl', 'wb'))
    # cPickle.dump(word2vec, open('../data/word2vec.pkl', 'wb'))
    cPickle.dump(vocabulary, open('../data/vocab.pkl', 'wb'))
    print "Data created"

