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'])

[((367, 394), 'pypagai.models.base.KerasModel.default_config', 'KerasModel.default_config', ([], {}), '()\n', (392, 394), False, 'from pypagai.models.base import KerasModel\n'), ((620, 662), 'keras.Input', 'Input', (['(self._story_maxlen,)'], {'name': '"""story"""'}), "((self._story_maxlen,), name='story')\n", (625, 662), False, 'from keras import Model, Input\n'), ((683, 728), 'keras.Input', 'Input', (['(self._query_maxlen,)'], {'name': '"""question"""'}), "((self._query_maxlen,), name='question')\n", (688, 728), False, 'from keras import Model, Input\n'), ((746, 776), 'keras.layers.concatenate', 'concatenate', (['[story, question]'], {}), '([story, question])\n', (757, 776), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((1042, 1091), 'keras.Model', 'Model', ([], {'inputs': '[story, question]', 'outputs': 'response'}), '(inputs=[story, question], outputs=response)\n', (1047, 1091), False, 'from keras import Model, Input\n'), ((1353, 1380), 'pypagai.models.base.KerasModel.default_config', 'KerasModel.default_config', ([], {}), '()\n', (1378, 1380), False, 'from pypagai.models.base import KerasModel\n'), ((1607, 1649), 'keras.Input', 'Input', (['(self._story_maxlen,)'], {'name': '"""story"""'}), "((self._story_maxlen,), name='story')\n", (1612, 1649), False, 'from keras import Model, Input\n'), ((1670, 1715), 'keras.Input', 'Input', (['(self._query_maxlen,)'], {'name': '"""question"""'}), "((self._query_maxlen,), name='question')\n", (1675, 1715), False, 'from keras import Model, Input\n'), ((1947, 1991), 'keras.layers.concatenate', 'concatenate', (['[eb_story, eb_question]'], {'axis': '(1)'}), '([eb_story, eb_question], axis=1)\n', (1958, 1991), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((2165, 2214), 'keras.Model', 'Model', ([], {'inputs': '[story, question]', 'outputs': 'response'}), '(inputs=[story, question], outputs=response)\n', (2170, 2214), False, 'from keras import Model, Input\n'), ((2475, 2502), 'pypagai.models.base.KerasModel.default_config', 'KerasModel.default_config', ([], {}), '()\n', (2500, 2502), False, 'from pypagai.models.base import KerasModel\n'), ((2745, 2787), 'keras.Input', 'Input', (['(self._story_maxlen,)'], {'name': '"""story"""'}), "((self._story_maxlen,), name='story')\n", (2750, 2787), False, 'from keras import Model, Input\n'), ((2808, 2853), 'keras.Input', 'Input', (['(self._query_maxlen,)'], {'name': '"""question"""'}), "((self._query_maxlen,), name='question')\n", (2813, 2853), False, 'from keras import Model, Input\n'), ((3497, 3541), 'keras.layers.concatenate', 'concatenate', (['[eb_story, eb_question]'], {'axis': '(1)'}), '([eb_story, eb_question], axis=1)\n', (3508, 3541), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((3715, 3764), 'keras.Model', 'Model', ([], {'inputs': '[story, question]', 'outputs': 'response'}), '(inputs=[story, question], outputs=response)\n', (3720, 3764), False, 'from keras import Model, Input\n'), ((847, 862), 'keras.layers.Permute', 'Permute', (['(2, 1)'], {}), '((2, 1))\n', (854, 862), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((889, 937), 'keras.layers.LSTM', 'LSTM', (['hidden'], {'dropout': '(0.2)', 'recurrent_dropout': '(0.2)'}), '(hidden, dropout=0.2, recurrent_dropout=0.2)\n', (893, 937), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((963, 1008), 'keras.layers.Dense', 'Dense', (['self._vocab_size'], {'activation': '"""softmax"""'}), "(self._vocab_size, activation='softmax')\n", (968, 1008), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((1737, 1768), 'keras.layers.Embedding', 'Embedding', (['self._vocab_size', '(64)'], {}), '(self._vocab_size, 64)\n', (1746, 1768), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((1795, 1807), 'keras.layers.Dropout', 'Dropout', (['(0.3)'], {}), '(0.3)\n', (1802, 1807), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((1841, 1872), 'keras.layers.Embedding', 'Embedding', (['self._vocab_size', '(64)'], {}), '(self._vocab_size, 64)\n', (1850, 1872), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((1905, 1917), 'keras.layers.Dropout', 'Dropout', (['(0.3)'], {}), '(0.3)\n', (1912, 1917), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((2012, 2060), 'keras.layers.LSTM', 'LSTM', (['hidden'], {'dropout': '(0.2)', 'recurrent_dropout': '(0.2)'}), '(hidden, dropout=0.2, recurrent_dropout=0.2)\n', (2016, 2060), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((2086, 2131), 'keras.layers.Dense', 'Dense', (['self._vocab_size'], {'activation': '"""softmax"""'}), "(self._vocab_size, activation='softmax')\n", (2091, 2131), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((2875, 2906), 'keras.layers.Embedding', 'Embedding', (['self._vocab_size', '(64)'], {}), '(self._vocab_size, 64)\n', (2884, 2906), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((2933, 2969), 'keras.layers.Convolution1D', 'Convolution1D', (['(64)', '(3)'], {'padding': '"""same"""'}), "(64, 3, padding='same')\n", (2946, 2969), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((2999, 3035), 'keras.layers.Convolution1D', 'Convolution1D', (['(32)', '(3)'], {'padding': '"""same"""'}), "(32, 3, padding='same')\n", (3012, 3035), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((3065, 3101), 'keras.layers.Convolution1D', 'Convolution1D', (['(16)', '(3)'], {'padding': '"""same"""'}), "(16, 3, padding='same')\n", (3078, 3101), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((3176, 3207), 'keras.layers.Embedding', 'Embedding', (['self._vocab_size', '(64)'], {}), '(self._vocab_size, 64)\n', (3185, 3207), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((3240, 3276), 'keras.layers.Convolution1D', 'Convolution1D', (['(64)', '(3)'], {'padding': '"""same"""'}), "(64, 3, padding='same')\n", (3253, 3276), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((3312, 3348), 'keras.layers.Convolution1D', 'Convolution1D', (['(32)', '(3)'], {'padding': '"""same"""'}), "(32, 3, padding='same')\n", (3325, 3348), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((3384, 3420), 'keras.layers.Convolution1D', 'Convolution1D', (['(16)', '(3)'], {'padding': '"""same"""'}), "(16, 3, padding='same')\n", (3397, 3420), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((3562, 3610), 'keras.layers.LSTM', 'LSTM', (['hidden'], {'dropout': '(0.2)', 'recurrent_dropout': '(0.2)'}), '(hidden, dropout=0.2, recurrent_dropout=0.2)\n', (3566, 3610), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((3636, 3681), 'keras.layers.Dense', 'Dense', (['self._vocab_size'], {'activation': '"""softmax"""'}), "(self._vocab_size, activation='softmax')\n", (3641, 3681), False, 'from keras.layers import Dense, concatenate, LSTM, Reshape, Permute, Embedding, Dropout, Convolution1D, Flatten\n'), ((1130, 1145), 'keras.optimizers.Adam', 'Adam', ([], {'lr': '(0.0002)'}), '(lr=0.0002)\n', (1134, 1145), False, 'from keras.optimizers import Adam\n'), ((2253, 2268), 'keras.optimizers.Adam', 'Adam', ([], {'lr': '(0.0002)'}), '(lr=0.0002)\n', (2257, 2268), False, 'from keras.optimizers import Adam\n'), ((3803, 3818), 'keras.optimizers.Adam', 'Adam', ([], {'lr': '(0.0002)'}), '(lr=0.0002)\n', (3807, 3818), False, 'from keras.optimizers import Adam\n')]