from __future__ import print_function

from functools import reduce

import tarfile

import numpy as np

import re

import pdb

import pickle

import h5py

from keras.preprocessing.sequence import pad_sequences

from glove import Glove

def check_repeated(name,repeated_list):

    name = name.lower().strip()

    return name if not (name in repeated_list) else repeated_list[name]

def process_title(word):

    return re.sub(r'\W+', ' ', word).strip().lower()

def tokenize(sent):

    '''Return the tokens of a sentence including punctuation.

    >>> tokenize('Bob dropped the apple. Where is the apple?')

    ['Bob', 'dropped', 'the', 'apple', '.', 'Where', 'is', 'the', 'apple', '?']

    '''

    return [x.strip().lower() for x in re.split('(\W+)?', sent) if x.strip()]

def parse_stories(lines, only_supporting=False,repeated_list=None):

    '''Parse stories

    If only_supporting is true, only the sentences that support the answer are kept.

    '''

    data = []

    story = []

    for line in lines:

        line = line.decode('utf-8').strip()

        spl = line.split(' ', 1)

        if len(spl) > 1:

            nid, line = spl

        else:

            continue

        try:

            nid = int(nid)

        except ValueError:

            pdb.set_trace()

        if nid == 0:

            story = []

        if '\t' in line:

            supporting, a = line.split('\t')

            a = map(process_title,a.split(','))

            options = [] if len(a) == 1 else list(set(a[1:]))

            a = a[0]    

            substory = None

            # Provide all the substories

            if supporting:

                story.append([tokenize(supporting) + [u'.']])

            substory = [x for x in story if x]

            # TODO: I should have done the lower in previous processing steps

            if not substory:

                continue

            data.append((substory, a.lower(), map(lambda x:x.lower(),options)))

        else:

            sent = tokenize(line)

            story.append([sent + [u'.']])

    return data

def get_stories(f, only_supporting=False, max_length=None,repeated_list=None, min_length=None):

    '''Given a file name, read the file, retrieve the stories, and then convert the sentences into a single story.

    If max_length is supplied, any stories longer than max_length tokens will be discarded.

    '''

    data = parse_stories(f.readlines(), only_supporting=only_supporting,repeated_list=repeated_list)

    flatten = lambda data: reduce(lambda x, y: x + y, data)

    data = [[flatten(reversed(story)), answer, options] for story,answer,options in data if not max_length or len(flatten(story)) < max_length]

    # At least two facts

    print(len(data))

    if min_length: 

        data = filter(lambda x: len(x[0]) > min_length, data)

    print(len(data))

    return data

def vectorize_stories(data, word_idx, story_maxlen, query_maxlen):

    X = []

    Xq = []

    Y = []

    for story, query, answer in data:

        x = [word_idx[w] for w in story]

        xq = [word_idx[w] for w in query]

        y = np.zeros(len(word_idx))  # let's not forget that index 0 is reserved

        y[word_idx[answer]] = 1

        X.append(x)

        Xq.append(xq)

        Y.append(y)

    return (pad_sequences(X, maxlen=story_maxlen),

            pad_sequences(Xq, maxlen=query_maxlen), np.array(Y))

def get_spacy_vectors(data, answer_dict, story_maxlen, model):

    X = []

    Y = []

    for story,answer in data:

        story = story[:story_maxlen] if len(story) > story_maxlen else story

        x = [model(unicode(w)).vector for w in story]

        X.append(x)

        if not answer_dict is None:

            y = np.zeros(len(answer_dict))

            y[answer_dict[answer]] = 1

            Y.append(y)

    return (pad_sequences(X, maxlen=story_maxlen,dtype='float32'),

            np.array(Y))

def get_word_vectors(data, answer_dict, story_maxlen, model):

    X = []

    Y = []

    for story,answer in data:

        story = story[:story_maxlen] if len(story) > story_maxlen else story

        x = [model.word_vectors[model.dictionary[w]] for w in story if 

             not model.dictionary.get(w) is None]

        X.append(x)

        if not answer_dict is None:

            y = np.zeros(len(answer_dict))

            y[answer_dict[answer]] = 1

            Y.append(y)

    return (pad_sequences(X, maxlen=story_maxlen,dtype='float32'),

            np.array(Y))

def create_vectors_dataset(input_files, vector_files, max_len=500):

    print('Creating word vectors file')

    training_set_file, test_set_file = input_files

    train_word_file, test_word_file = vector_files

    train_stories = pickle.load(open(training_set_file,'r'))

    test_stories = pickle.load(open(test_set_file,'r'))

    train_stories = [(reduce(lambda x,y: x + y, map(list,fact)),q) for fact,q in train_stories]

    test_stories = [(reduce(lambda x,y: x + y, map(list,fact)),q) for fact,q in test_stories]

    vocab = sorted(reduce(lambda x, y: x | y, (set(story + [answer]) for story, answer in train_stories + test_stories)))

    # Reserve 0 for masking via pad_sequences

    vocab_size = len(vocab) + 1

    story_maxlen = max(map(len, (x for x, _ in train_stories + test_stories)))

    print('-')

    print('Vocab size:', vocab_size, 'unique words')

    print('Story max length:', story_maxlen, 'words')

    print('Number of training stories:', len(train_stories))

    print('Number of test stories:', len(test_stories))

    print('-')

    print('Here\'s what a "story" tuple looks like (input, query, answer):')

    print(train_stories[0])

    print('-')

    print('Vectorizing the word sequences...')

    word_idx = dict((c, i + 1) for i, c in enumerate(vocab))

    answer_vocab = sorted(reduce(lambda x, y: x | y, (set([answer]) for _, answer in train_stories + test_stories)))

    # Reserve 0 for masking via pad_sequences

    answer_dict = dict((word, i) for i, word in enumerate(answer_vocab))

    print('Answers dict len: {0}'.format(len(answer_dict)))

    # I need to check also if this exist

    word_vectors_dir = 'word_vectors/glove.42B.300d.txt'

    word_vectors_model = Glove.load_stanford(word_vectors_dir)

    inputs_train, answers_train = get_word_vectors(train_stories, answer_dict, 

                                                   max_len, word_vectors_model)

    inputs_test, answers_test = get_word_vectors(test_stories, answer_dict, max_len,

                                                 word_vectors_model)

    with h5py.File(train_word_file,'w') as train_f:

        _ = train_f.create_dataset('inputs',data=inputs_train)

        _ = train_f.create_dataset('answers',data=answers_train)

    with h5py.File(test_word_file,'w') as test_f:

        _ = test_f.create_dataset('inputs',data=inputs_test)

        _ = test_f.create_dataset('answers',data=answers_test)

    return (inputs_train, answers_train),(inputs_test, answers_test)

def save_vectors_dict(input_files):

    # I need to check also if this exist

    filename = 'word_vectors/glove.42B.300d.txt'

    word_vectors_dict = 'word_vectors/glove_dict.hdf5'

    dct = {}

    vectors = array.array('d')

    # Read in the data.

    with io.open(filename, 'r', encoding='utf-8') as savefile:

        for i, line in enumerate(savefile):

            tokens = line.split(' ')

            word = tokens[0]

            entries = tokens[1:]

            dct[word] = i

            vectors.extend(float(x) for x in entries)

    print('Saving to hf5 file')

    with h5py.File(word_vectors_dict,'w') as vector_f:

        _ = vector_f.create_dataset('vectors',data=dct)