from keras import Model, Input
from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten
from keras.optimizers import Adam

from pypagai.models.base import KerasModel


class SimpleLSTM(KerasModel):
    """
    Use a simple lstm neural network
    """
    @staticmethod
    def default_config():
        config = KerasModel.default_config()
        config['hidden'] = 32

        return config

    def __init__(self, cfg):
        super().__init__(cfg)
        self._cfg_ = cfg

    def _create_network_(self):
        hidden = self._cfg_['hidden']
        story = Input((self._story_maxlen, ), name='story')
        question = Input((self._query_maxlen, ), name='question')

        conc = concatenate([story, question],)
        conc = Reshape((1, int(conc.shape[1])))(conc)
        conc = Permute((2, 1))(conc)

        response = LSTM(hidden, dropout=0.2, recurrent_dropout=0.2)(conc)
        response = Dense(self._vocab_size, activation='softmax')(response)

        self._model = Model(inputs=[story, question], outputs=response)
        self._model.compile(optimizer=Adam(lr=2e-4), loss='sparse_categorical_crossentropy', metrics=['accuracy'])


class EmbedLSTM(KerasModel):

    """
    Use a simple lstm neural network
    """
    @staticmethod
    def default_config():
        config = KerasModel.default_config()
        config['hidden'] = 32

        return config

    def __init__(self, cfg):
        super().__init__(cfg)
        self._cfg_ = cfg

    def _create_network_(self):
        hidden = self._cfg_['hidden']

        story = Input((self._story_maxlen, ), name='story')
        question = Input((self._query_maxlen, ), name='question')

        eb_story = Embedding(self._vocab_size, 64)(story)
        eb_story = Dropout(0.3)(eb_story)

        eb_question = Embedding(self._vocab_size, 64)(question)
        eb_question = Dropout(0.3)(eb_question)

        conc = concatenate([eb_story, eb_question], axis=1)

        response = LSTM(hidden, dropout=0.2, recurrent_dropout=0.2)(conc)
        response = Dense(self._vocab_size, activation='softmax')(response)

        self._model = Model(inputs=[story, question], outputs=response)
        self._model.compile(optimizer=Adam(lr=2e-4), loss='sparse_categorical_crossentropy', metrics=['accuracy'])


class ConvLSTM(KerasModel):

    """
    Use a simple lstm neural network
    """
    @staticmethod
    def default_config():
        config = KerasModel.default_config()
        config['hidden'] = 32

        return config

    def __init__(self, model_cfg):
        super().__init__(model_cfg)
        self._cfg = model_cfg

    def _create_network_(self):
        hidden = self._cfg['hidden']

        story = Input((self._story_maxlen, ), name='story')
        question = Input((self._query_maxlen, ), name='question')

        eb_story = Embedding(self._vocab_size, 64)(story)
        eb_story = Convolution1D(64, 3, padding='same')(eb_story)
        eb_story = Convolution1D(32, 3, padding='same')(eb_story)
        eb_story = Convolution1D(16, 3, padding='same')(eb_story)
        # eb_story = Flatten()(eb_story)

        eb_question = Embedding(self._vocab_size, 64)(question)
        eb_question = Convolution1D(64, 3, padding='same')(eb_question)
        eb_question = Convolution1D(32, 3, padding='same')(eb_question)
        eb_question = Convolution1D(16, 3, padding='same')(eb_question)
        # eb_question = Flatten()(eb_question)

        conc = concatenate([eb_story, eb_question], axis=1)

        response = LSTM(hidden, dropout=0.2, recurrent_dropout=0.2)(conc)
        response = Dense(self._vocab_size, activation='softmax')(response)

        self._model = Model(inputs=[story, question], outputs=response)
        self._model.compile(optimizer=Adam(lr=2e-4), loss='sparse_categorical_crossentropy', metrics=['accuracy'])

[((14, 17, 14, 44), 'pypagai.models.base.KerasModel.default_config', 'KerasModel.default_config', ({}, {}), '()', False, 'from pypagai.models.base import KerasModel\n'), ((25, 16, 25, 59), 'keras.Input', 'Input', (), '', False, 'from keras import Model, Input\n'), ((26, 19, 26, 65), 'keras.Input', 'Input', (), '', False, 'from keras import Model, Input\n'), ((28, 15, 28, 46), 'keras.layers.concatenate', 'concatenate', ({(28, 27, 28, 44): '[story, question]'}, {}), '([story, question])', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((35, 22, 35, 71), 'keras.Model', 'Model', (), '', False, 'from keras import Model, Input\n'), ((46, 17, 46, 44), 'pypagai.models.base.KerasModel.default_config', 'KerasModel.default_config', ({}, {}), '()', False, 'from pypagai.models.base import KerasModel\n'), ((58, 16, 58, 59), 'keras.Input', 'Input', (), '', False, 'from keras import Model, Input\n'), ((59, 19, 59, 65), 'keras.Input', 'Input', (), '', False, 'from keras import Model, Input\n'), ((67, 15, 67, 59), 'keras.layers.concatenate', 'concatenate', (), '', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((72, 22, 72, 71), 'keras.Model', 'Model', (), '', False, 'from keras import Model, Input\n'), ((83, 17, 83, 44), 'pypagai.models.base.KerasModel.default_config', 'KerasModel.default_config', ({}, {}), '()', False, 'from pypagai.models.base import KerasModel\n'), ((95, 16, 95, 59), 'keras.Input', 'Input', (), '', False, 'from keras import Model, Input\n'), ((96, 19, 96, 65), 'keras.Input', 'Input', (), '', False, 'from keras import Model, Input\n'), ((110, 15, 110, 59), 'keras.layers.concatenate', 'concatenate', (), '', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((115, 22, 115, 71), 'keras.Model', 'Model', (), '', False, 'from keras import Model, Input\n'), ((30, 15, 30, 30), 'keras.layers.Permute', 'Permute', ({(30, 23, 30, 29): '(2, 1)'}, {}), '((2, 1))', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((32, 19, 32, 67), 'keras.layers.LSTM', 'LSTM', (), '', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((33, 19, 33, 64), 'keras.layers.Dense', 'Dense', (), '', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((61, 19, 61, 50), 'keras.layers.Embedding', 'Embedding', ({(61, 29, 61, 45): 'self._vocab_size', (61, 47, 61, 49): '64'}, {}), '(self._vocab_size, 64)', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((62, 19, 62, 31), 'keras.layers.Dropout', 'Dropout', ({(62, 27, 62, 30): '0.3'}, {}), '(0.3)', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((64, 22, 64, 53), 'keras.layers.Embedding', 'Embedding', ({(64, 32, 64, 48): 'self._vocab_size', (64, 50, 64, 52): '64'}, {}), '(self._vocab_size, 64)', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((65, 22, 65, 34), 'keras.layers.Dropout', 'Dropout', ({(65, 30, 65, 33): '0.3'}, {}), '(0.3)', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((69, 19, 69, 67), 'keras.layers.LSTM', 'LSTM', (), '', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((70, 19, 70, 64), 'keras.layers.Dense', 'Dense', (), '', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((98, 19, 98, 50), 'keras.layers.Embedding', 'Embedding', ({(98, 29, 98, 45): 'self._vocab_size', (98, 47, 98, 49): '64'}, {}), '(self._vocab_size, 64)', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((99, 19, 99, 55), 'keras.layers.Convolution1D', 'Convolution1D', (), '', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((100, 19, 100, 55), 'keras.layers.Convolution1D', 'Convolution1D', (), '', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((101, 19, 101, 55), 'keras.layers.Convolution1D', 'Convolution1D', (), '', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((104, 22, 104, 53), 'keras.layers.Embedding', 'Embedding', ({(104, 32, 104, 48): 'self._vocab_size', (104, 50, 104, 52): '64'}, {}), '(self._vocab_size, 64)', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((105, 22, 105, 58), 'keras.layers.Convolution1D', 'Convolution1D', (), '', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((106, 22, 106, 58), 'keras.layers.Convolution1D', 'Convolution1D', (), '', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((107, 22, 107, 58), 'keras.layers.Convolution1D', 'Convolution1D', (), '', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((112, 19, 112, 67), 'keras.layers.LSTM', 'LSTM', (), '', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((113, 19, 113, 64), 'keras.layers.Dense', 'Dense', (), '', False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((36, 38, 36, 51), 'keras.optimizers.Adam', 'Adam', (), '', False, 'from keras.optimizers import Adam\n'), ((73, 38, 73, 51), 'keras.optimizers.Adam', 'Adam', (), '', False, 'from keras.optimizers import Adam\n'), ((116, 38, 116, 51), 'keras.optimizers.Adam', 'Adam', (), '', False, 'from keras.optimizers import Adam\n')]