# Path: seya/layers/recurrent.py
# class StatefulGRU(GRU):
#     def __init__(self, batch_size, output_dim=128,
#                  init='glorot_uniform', inner_init='orthogonal',
#                  activation='sigmoid', inner_activation='hard_sigmoid',
#                  weights=None, return_sequences=False,
#                  input_dim=None, input_length=None, **kwargs):
# 
#         self.batch_size = batch_size
#         self.input_dim = input_dim
#         self.input_length = input_length
#         if self.input_dim:
#             kwargs['input_shape'] = (self.input_length, self.input_dim)
# 
#         super(StatefulGRU, self).__init__(
#             output_dim, init=init, inner_init=inner_init,
#             activation=activation, inner_activation=inner_activation,
#             weights=weights,
#             return_sequences=return_sequences,
#             input_dim=input_dim, input_length=input_length, **kwargs)
# 
#     def build(self):
#         super(StatefulGRU, self).build()
#         self.h = K.zeros((self.batch_size, self.output_dim))  # Here is the state
# 
#     def get_output(self, train=False):
#         X = self.get_input(train)
#         padded_mask = self.get_padded_shuffled_mask(train, X, pad=1)
#         X = X.dimshuffle((1, 0, 2))
# 
#         x_z = T.dot(X, self.W_z) + self.b_z
#         x_r = T.dot(X, self.W_r) + self.b_r
#         x_h = T.dot(X, self.W_h) + self.b_h
#         outputs, updates = theano.scan(
#             self._step,
#             sequences=[x_z, x_r, x_h, padded_mask],
#             outputs_info=self.h[:X.shape[1]],
#             non_sequences=[self.U_z, self.U_r, self.U_h])
# 
#         self.updates = ((self.h, outputs[-1]), )  # initial state of next batch
#                                                   # is the last state of this
#                                                   # batch
#         if self.return_sequences:
#             return outputs.dimshuffle((1, 0, 2))
#         return outputs[-1]
# 
#     def init_updates(self):
#         self.get_output(train=True)
# 
#     def get_config(self):
#         return {"name": self.__class__.__name__,
#                 "input_dim": self.input_dim,
#                 "output_dim": self.output_dim,
#                 "init": self.init.__name__,
#                 "inner_init": self.inner_init.__name__,
#                 "activation": self.activation.__name__,
#                 "inner_activation": self.inner_activation.__name__,
#                 "return_sequences": self.return_sequences}
# 
# Path: seya/callbacks.py
# class ResetRNNState(Callback):
#     """
#     This is supposed to be used with stateful RNNs like
#     seya.layers.recurrent.StatefulGRU
# 
#     h: the rnn state
#     func: a function that returns true when the state should be reset to zero
#     """
#     def __init__(self, h, func):
#         self.h = h
#         self.func = func
# 
#     def on_batch_end(self, batch, logs={}):
#         if self.func(batch, logs):
#             self.h.set_value(self.h.get_value()*0)


from keras.models import Sequential
from keras.layers.core import Dense, Activation, Dropout
from keras.layers.recurrent import LSTM
from keras.datasets.data_utils import get_file
from seya.layers.recurrent import StatefulGRU as SGRU
from seya.callbacks import ResetRNNState
import numpy as np
import random, sys
maxlen = 20
step = 3
sentences = []
next_chars = []
for i in range(0, len(text) - maxlen, step):
    sentences.append(text[i : i + maxlen])
    next_chars.append(text[i + maxlen])
print('nb sequences:', len(sentences))

print('Vectorization...')
X = np.zeros((len(sentences), maxlen, len(chars)), dtype=np.bool)
y = np.zeros((len(sentences), len(chars)), dtype=np.bool)
for i, sentence in enumerate(sentences):
    for t, char in enumerate(sentence):
        X[i, t, char_indices[char]] = 1
    y[i, char_indices[next_chars[i]]] = 1


# build the model: 2 stacked Stateful GRUs
print('Build model...')
model = Sequential()
model.add(SGRU(input_shape=(None, len(chars)), batch_size=batch_size,
               output_dim=512, return_sequences=True))
model.add(Dropout(0.2))
model.add(SGRU(batch_size=batch_size, output_dim=512, return_sequences=False))
model.add(Dropout(0.2))
model.add(Dense(len(chars)))
model.add(Activation('softmax'))

# define callbacks to reset state after every 4 batches

# Path: seya/layers/recurrent.py
# class StatefulGRU(GRU):
#     def __init__(self, batch_size, output_dim=128,
#                  init='glorot_uniform', inner_init='orthogonal',
#                  activation='sigmoid', inner_activation='hard_sigmoid',
#                  weights=None, return_sequences=False,
#                  input_dim=None, input_length=None, **kwargs):
# 
#         self.batch_size = batch_size
#         self.input_dim = input_dim
#         self.input_length = input_length
#         if self.input_dim:
#             kwargs['input_shape'] = (self.input_length, self.input_dim)
# 
#         super(StatefulGRU, self).__init__(
#             output_dim, init=init, inner_init=inner_init,
#             activation=activation, inner_activation=inner_activation,
#             weights=weights,
#             return_sequences=return_sequences,
#             input_dim=input_dim, input_length=input_length, **kwargs)
# 
#     def build(self):
#         super(StatefulGRU, self).build()
#         self.h = K.zeros((self.batch_size, self.output_dim))  # Here is the state
# 
#     def get_output(self, train=False):
#         X = self.get_input(train)
#         padded_mask = self.get_padded_shuffled_mask(train, X, pad=1)
#         X = X.dimshuffle((1, 0, 2))
# 
#         x_z = T.dot(X, self.W_z) + self.b_z
#         x_r = T.dot(X, self.W_r) + self.b_r
#         x_h = T.dot(X, self.W_h) + self.b_h
#         outputs, updates = theano.scan(
#             self._step,
#             sequences=[x_z, x_r, x_h, padded_mask],
#             outputs_info=self.h[:X.shape[1]],
#             non_sequences=[self.U_z, self.U_r, self.U_h])
# 
#         self.updates = ((self.h, outputs[-1]), )  # initial state of next batch
#                                                   # is the last state of this
#                                                   # batch
#         if self.return_sequences:
#             return outputs.dimshuffle((1, 0, 2))
#         return outputs[-1]
# 
#     def init_updates(self):
#         self.get_output(train=True)
# 
#     def get_config(self):
#         return {"name": self.__class__.__name__,
#                 "input_dim": self.input_dim,
#                 "output_dim": self.output_dim,
#                 "init": self.init.__name__,
#                 "inner_init": self.inner_init.__name__,
#                 "activation": self.activation.__name__,
#                 "inner_activation": self.inner_activation.__name__,
#                 "return_sequences": self.return_sequences}
# 
# Path: seya/callbacks.py
# class ResetRNNState(Callback):
#     """
#     This is supposed to be used with stateful RNNs like
#     seya.layers.recurrent.StatefulGRU
# 
#     h: the rnn state
#     func: a function that returns true when the state should be reset to zero
#     """
#     def __init__(self, h, func):
#         self.h = h
#         self.func = func
# 
#     def on_batch_end(self, batch, logs={}):
#         if self.func(batch, logs):
#             self.h.set_value(self.h.get_value()*0)


# Path: examples/stateful_text_generation.py
from keras.models import Sequential
from keras.layers.core import Dense, Activation, Dropout
from keras.layers.recurrent import LSTM
from keras.datasets.data_utils import get_file
from seya.layers.recurrent import StatefulGRU as SGRU
from seya.callbacks import ResetRNNState
import numpy as np
import random, sys

maxlen = 20
step = 3
sentences = []
next_chars = []
for i in range(0, len(text) - maxlen, step):
    sentences.append(text[i : i + maxlen])
    next_chars.append(text[i + maxlen])
print('nb sequences:', len(sentences))

print('Vectorization...')
X = np.zeros((len(sentences), maxlen, len(chars)), dtype=np.bool)
y = np.zeros((len(sentences), len(chars)), dtype=np.bool)
for i, sentence in enumerate(sentences):
    for t, char in enumerate(sentence):
        X[i, t, char_indices[char]] = 1
    y[i, char_indices[next_chars[i]]] = 1


# build the model: 2 stacked Stateful GRUs
print('Build model...')
model = Sequential()
model.add(SGRU(input_shape=(None, len(chars)), batch_size=batch_size,
               output_dim=512, return_sequences=True))
model.add(Dropout(0.2))
model.add(SGRU(batch_size=batch_size, output_dim=512, return_sequences=False))
model.add(Dropout(0.2))
model.add(Dense(len(chars)))
model.add(Activation('softmax'))

# define callbacks to reset state after every 4 batches

h1_reset = ResetRNNState(model.layers[0].h, lambda batches, logs: batches%4 == 0)