from keras.models import Sequential, Model
from keras.layers.embeddings import Embedding
from keras.layers import Input, Activation, Dense, Permute, Dropout, add, dot, concatenate
from keras.layers import LSTM, TimeDistributed
from keras.layers import Conv1D, MaxPool1D
from keras.layers.wrappers import Bidirectional

class ConfigurableNetwork:
    defaultfile = 'default.cfg'
    def __init__(self, modelname):
        self._modelname = modelname

        # Set up the environment. load config file if it is there, if not, then create it form default
        # Load the weights

        # existential question: should this contain all of the handling for fitting and callbacks?
        # yeah I think that makes sense. One stop shop for configs, weights, and logging.


    def setup(self):
        pass

    @property
    def modelname(self): return self._modelname



class QueryableNet:
    def __init__(self, vocab_size=22, story_maxlen=68, query_maxlen=4):
        self.model = None
        self.vocab_size = vocab_size
        self.query_maxlen = query_maxlen
        self.story_maxlen = story_maxlen


    def query(self, storyvec, queryvec):
        storyvec = storyvec.reshape((-1, self.story_maxlen))
        queryvec = queryvec.reshape((-1, self.query_maxlen))
        ans = self.model.predict([storyvec, queryvec])
        return ans

class DeepMemNet(QueryableNet):
    """
    DeepMemNet for the Facebook bAbI context task.

    Model notes:
    Single context task:
      Regular LSTM (accuracy/val_acc):
        Run 1
            47/50% @ 7 epochs
            72/70% @ 49 epochs
            86/80% @ 61 epochs
            95/90% @ 88 epochs
        Run 2



      Bidirectional LSTM:
        Run 1
            50%/50% @ 6
            81%/80% @ 34
            87%/86% @ 48 (peak valacc)
            90%/86% @ 60 epochs - minor overfitting

        I think parameters were not configured right, these might be Single LSTM! need to rerun everything >_<
        Run 2
            73/71% @ 44
            83/80% @ 53
            94/90% @ 75
        Run 3
            71/70% @ 41
            83/80% @ 50
            94/90% @ 76

      Bidirectional + extra forward LSTM:
        55%/??% @ 80 (stalls)

      TDDense + Bidirectional:
        Run 1:
            76/73% @ 37
            80/80% @ 41
            91/90% @ 54 - new record!!

    Double Context task:
      Regular:
        50%/??% @ 35 epochs
        67%/??% @ 80 epochs
        70%/??% @ 100
        80%/??% @ 192
        84.7%/??% @ 260

      Bidirectional:
        50%/??% @ 26 epochs
        70%/??% @ 48 epochs - improvement!
        80%/??% @ 68 epochs - super improvement!
        90%/??% @ 110 epochs - smokin'!
        95%/??% @ 148 epochs - starting to level off
        97%/??% @ 200 epochs - i think it's starting to overfit
    """
    # todo: add performance logging
    def __init__(self, vocab_size=22, story_maxlen=68, query_maxlen=4, n_lstm=32, bidirect=True, tdd=True,
                 batch_size=None, matchconv=False, permute=False):
        """
        DeepMemNet

        Param note - changing parameters will require new model file (duh) - this isn't automatic yet
        :param vocab_size:
        :param story_maxlen:
        :param query_maxlen:
        :param n_lstm:
        :param bidirect:
        """

        # todo: config file for model hyperparams with logging link

        # self.vocab_size = vocab_size
        # self.story_maxlen = story_maxlen
        # self.query_maxlen = query_maxlen
        super().__init__(vocab_size=vocab_size, story_maxlen=story_maxlen, query_maxlen=query_maxlen)
        # placeholders
        input_sequence = Input((story_maxlen,), name='InputSeq')
        question = Input((query_maxlen,), name='Question')

        # Encoders - initial encoders are pretty much just projecting the input into a useful space
        # not much need to optimize here really
        input_encoder_m = Sequential(name='InputEncoderM')