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, 'pypagai.models.base.KerasModel.default_config', 'KerasModel.default_config', ([], {}), False, 'from pypagai.models.base import KerasModel\n'), (25, 'keras.Input', 'Input', (['(self._story_maxlen,)'], {'name': '"""story"""'}), False, 'from keras import Model, Input\n'), (26, 'keras.Input', 'Input', (['(self._query_maxlen,)'], {'name': '"""question"""'}), False, 'from keras import Model, Input\n'), (28, 'keras.layers.concatenate', 'concatenate', (['[story, question]'], {}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (35, 'keras.Model', 'Model', ([], {'inputs': '[story, question]', 'outputs': 'response'}), False, 'from keras import Model, Input\n'), (46, 'pypagai.models.base.KerasModel.default_config', 'KerasModel.default_config', ([], {}), False, 'from pypagai.models.base import KerasModel\n'), (58, 'keras.Input', 'Input', (['(self._story_maxlen,)'], {'name': '"""story"""'}), False, 'from keras import Model, Input\n'), (59, 'keras.Input', 'Input', (['(self._query_maxlen,)'], {'name': '"""question"""'}), False, 'from keras import Model, Input\n'), (67, 'keras.layers.concatenate', 'concatenate', (['[eb_story, eb_question]'], {'axis': '(1)'}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (72, 'keras.Model', 'Model', ([], {'inputs': '[story, question]', 'outputs': 'response'}), False, 'from keras import Model, Input\n'), (83, 'pypagai.models.base.KerasModel.default_config', 'KerasModel.default_config', ([], {}), False, 'from pypagai.models.base import KerasModel\n'), (95, 'keras.Input', 'Input', (['(self._story_maxlen,)'], {'name': '"""story"""'}), False, 'from keras import Model, Input\n'), (96, 'keras.Input', 'Input', (['(self._query_maxlen,)'], {'name': '"""question"""'}), False, 'from keras import Model, Input\n'), (110, 'keras.layers.concatenate', 'concatenate', (['[eb_story, eb_question]'], {'axis': '(1)'}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (115, 'keras.Model', 'Model', ([], {'inputs': '[story, question]', 'outputs': 'response'}), False, 'from keras import Model, Input\n'), (30, 'keras.layers.Permute', 'Permute', (['(2, 1)'], {}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (32, 'keras.layers.LSTM', 'LSTM', (['hidden'], {'dropout': '(0.2)', 'recurrent_dropout': '(0.2)'}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (33, 'keras.layers.Dense', 'Dense', (['self._vocab_size'], {'activation': '"""softmax"""'}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (61, 'keras.layers.Embedding', 'Embedding', (['self._vocab_size', '(64)'], {}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (62, 'keras.layers.Dropout', 'Dropout', (['(0.3)'], {}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (64, 'keras.layers.Embedding', 'Embedding', (['self._vocab_size', '(64)'], {}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (65, 'keras.layers.Dropout', 'Dropout', (['(0.3)'], {}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (69, 'keras.layers.LSTM', 'LSTM', (['hidden'], {'dropout': '(0.2)', 'recurrent_dropout': '(0.2)'}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (70, 'keras.layers.Dense', 'Dense', (['self._vocab_size'], {'activation': '"""softmax"""'}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (98, 'keras.layers.Embedding', 'Embedding', (['self._vocab_size', '(64)'], {}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (99, 'keras.layers.Convolution1D', 'Convolution1D', (['(64)', '(3)'], {'padding': '"""same"""'}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (100, 'keras.layers.Convolution1D', 'Convolution1D', (['(32)', '(3)'], {'padding': '"""same"""'}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (101, 'keras.layers.Convolution1D', 'Convolution1D', (['(16)', '(3)'], {'padding': '"""same"""'}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (104, 'keras.layers.Embedding', 'Embedding', (['self._vocab_size', '(64)'], {}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (105, 'keras.layers.Convolution1D', 'Convolution1D', (['(64)', '(3)'], {'padding': '"""same"""'}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (106, 'keras.layers.Convolution1D', 'Convolution1D', (['(32)', '(3)'], {'padding': '"""same"""'}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (107, 'keras.layers.Convolution1D', 'Convolution1D', (['(16)', '(3)'], {'padding': '"""same"""'}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (112, 'keras.layers.LSTM', 'LSTM', (['hidden'], {'dropout': '(0.2)', 'recurrent_dropout': '(0.2)'}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (113, 'keras.layers.Dense', 'Dense', (['self._vocab_size'], {'activation': '"""softmax"""'}), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), (36, 'keras.optimizers.Adam', 'Adam', ([], {'lr': '(0.0002)'}), False, 'from keras.optimizers import Adam\n'), (73, 'keras.optimizers.Adam', 'Adam', ([], {'lr': '(0.0002)'}), False, 'from keras.optimizers import Adam\n'), (116, 'keras.optimizers.Adam', 'Adam', ([], {'lr': '(0.0002)'}), False, 'from keras.optimizers import Adam\n')]