from keras.models import Sequential
from keras.layers import Dense, Dropout, LSTM

from utils import N_UNITS, DROPOUT_RATE


class Model(Sequential):

    def __init__(self, x, y):
        super().__init__()
        self.add(LSTM(N_UNITS, input_shape=(x.shape[1], x.shape[2]), return_sequences=True))
        self.add(Dropout(DROPOUT_RATE))
        self.add(LSTM(N_UNITS))
        self.add(Dropout(DROPOUT_RATE))
        self.add(Dense(y.shape[1], activation='softmax'))

    def compile_model(self):
        self.compile(loss='categorical_crossentropy', optimizer='adam')

    def fit_model(self, x, y, batch_size, epochs):
        self.fit(x, y, batch_size, epochs)