'''
Created on 2016-5-13

@author: psi
'''

import numpy as np
from keras.utils.np_utils import to_categorical

# TODO: function for converting one-hot text to n-gram data points
# TODO: function for converting one-hot text to other classes of LMs

def get_categorical_generator_train(X, y, max_index, batch_size, batch_per_epoch):
    '''
    Create a new generator that randomly select (batch_size * batch_per_epoch) examples
    and convert y vector to categorical batch by batch. The generator will loop over
    its data infinitely.
    Might be used to avoid memory issue when the range of y is large.
    '''
    n = len(X)
    if batch_size * batch_per_epoch > n:
        raise ValueError("Input data has less than batch_size * batch_per_epoch examples")
    X = np.array(X).astype('int32')
    y = np.array(y).astype('int32')
    while True:
        sample_idx = np.random.choice(np.arange(n), batch_per_epoch * batch_size, replace=False)
        for i in range(0, batch_per_epoch):
            begin_idx = i * batch_size
            end_idx = begin_idx + batch_size
            batchX = X[sample_idx[begin_idx:end_idx]]
            batchy = y[sample_idx[begin_idx:end_idx]]
            batchy = to_categorical(batchy, max_index)  
            yield (batchX, batchy)

def get_categorical_generator_eval(X, y, max_index, batch_size):
    '''
    Create a new generator that convert y vector to categorical batch by batch. The
    generator will loop over the data infinitely.
    Might be used to avoid memory issue when the range of y is large.
    '''
    X = np.array(X).astype('int32')
    y = np.array(y).astype('int32')
    n = len(X)
    batch_per_epoch = (n + batch_size - 1) / batch_size
    while True:
        range_idx = np.arange(n)
        for i in range(0, batch_per_epoch - 1):
            begin_idx = i * batch_size
            end_idx = begin_idx + batch_size
            batchX = X[range_idx[begin_idx:end_idx]]
            batchy = y[range_idx[begin_idx:end_idx]]
            batchy = to_categorical(batchy, max_index)  
            yield (batchX, batchy)
        begin_idx = (batch_per_epoch - 1) * batch_size
        end_idx = n
        batchX = X[range_idx[begin_idx:end_idx]]
        batchy = y[range_idx[begin_idx:end_idx]]
        batchy = to_categorical(batchy, max_index)  
        yield (batchX, batchy)

def get_categorical(y, max_index):
    return to_categorical(y, max_index)

def preprocess_text_ngram(strtext, order, vocabulary, is_training=False):
    '''
    Preprocess string text into N-gram dataset.
    
    strtext: input text as a string.
    vocabulary: VocabularyDict for storing index mapping and one-hot encoding words, 
          must be provided and properly configured before calling this
    is_training: whether the set should be used as training set. If true, vocabulary
          will be updated.
    '''
    if is_training:
        vocabulary.train_vocabulary(strtext)
        idxtext = vocabulary.encode_text(strtext, True)
    else:
        idxtext = vocabulary.encode_text(strtext, False)
    X = list()
    y = list()
    for idxlist in idxtext:
        for i in range(0, len(idxlist) - order + 1):
            X.append(idxlist[i:i + order - 1])
            y.append(idxlist[i + order - 1])
    X = np.array(X).astype('int32')
    y = np.array(y).astype('int32')
    return (X, y)

# TODO: doc lm_data_gen
