import tensorflow as tf
import keras.backend as K
from keras.layers import Input, Dense, Embedding, Lambda, Concatenate, multiply, Conv1D, Reshape, BatchNormalization, Dropout, LeakyReLU, RepeatVector, multiply, Activation, MaxPooling1D, Add
from keras.models import Model, load_model
from keras.callbacks import Callback
from keras.callbacks import ModelCheckpoint


class WarmUpLearningRateScheduler(Callback):
    def __init__(self, warmup_batches, init_lr, verbose=0):
        super(WarmUpLearningRateScheduler, self).__init__()
        self.warmup_batches = warmup_batches
        self.init_lr = init_lr
        self.verbose = verbose
        self.batch_count = 0
        self.learning_rates = []

    def on_batch_end(self, batch, logs=None):
        self.batch_count = self.batch_count + 1
        lr = K.get_value(self.model.optimizer.lr)
        self.learning_rates.append(lr)

    def on_batch_begin(self, batch, logs=None):
        if self.batch_count <= self.warmup_batches:
            lr = self.batch_count*self.init_lr/self.warmup_batches
            K.set_value(self.model.optimizer.lr, lr)
            if self.verbose > 0:
                print('\nBatch %05d: WarmUpLearningRateScheduler setting learning '
                      'rate to %s.' % (self.batch_count + 1, lr))

def get_sum_model(nr, ncd, nd, neb, ne, nx, nfc = 128, nn=3):
    # nr: embedding size; ncd: compressed \ell; nd: maximum sketch dimensions;
    # neb: max column resolution (128); ne: individual cell embedding length (64);
    # nfc: FC layers size; nn: FC layers

    ndr = 0.3
    i = Input(shape=(ncd*2 + 4, nd + 1,))
    xv = Lambda(lambda x: K.repeat_elements(K.expand_dims(x[:, :, 0]), nr, axis=2))(i) # increase one dimension, size 128
    x = Lambda(lambda x: x[:, :, 1:])(i)
    print(x.shape, ncd, nd)

    x = Reshape(((ncd*2 + 4) * nd, 1))(x)
    x = Embedding(neb + 1 + 1 + 1 + 1, ne, mask_zero=False, name='phic')(x)
    x = Lambda(lambda x: x, output_shape=lambda s: s)(x)
    x = Reshape(((ncd*2 + 4), nd * ne))(x)
    for n in range(nn):
        x = Dense(nfc, use_bias=True, name='phir' + str(n))(x)
        x = LeakyReLU()(x)
        x = Dropout(ndr)(x)
    x = Dense(nr, name='phirf')(x)
    #x = Activation('elu')(x)
    x = LeakyReLU()(x)
    x = Dropout(ndr)(x)

    st = 0
    ed = ncd

    qr = Lambda(lambda x: x[:, -4:-2, :])(x)
    qr = Reshape((1, 2 * nr))(qr)
    vq = Concatenate(axis=-1)([qr, qr])
    for a in range(nx - 2):
        vq = Concatenate(axis=-1)([vq, qr])

    xk = Lambda(lambda x: x[:, st:ed, :])(x)
    xf = Lambda(lambda x: x[:, st:ed, :])(xv)
    xk = multiply([xk, xf])
    e = Lambda(lambda x: K.sum(x[:, :], axis=1), output_shape=(lambda shape: (shape[0], shape[2])))(xk)
    vd = Concatenate(axis=-1)([e, e])
    for a in range(nx-2):
        vd = Concatenate(axis=-1)([vd, e])

    st = ncd
    ed = 2*ncd
    xk = Lambda(lambda x: x[:, st:ed, :])(x)
    xf = Lambda(lambda x: x[:, st:ed, :])(xv)
    xk = multiply([xk, xf])

    qr = Lambda(lambda x: x[:, -2:, :])(x)
    qr = Reshape((1, 2 * nr))(qr)
    for a in range(nx):
        vq = Concatenate(axis=-1)([vq, qr])

    e = Lambda(lambda x: K.sum(x[:, :], axis=1), output_shape=(lambda shape: (shape[0], shape[2])))(xk)
    for a in range(nx):
        vd = Concatenate(axis=-1)([vd, e])


    for n in range(nn):
        vd = Dense(nr*nx, use_bias=True, name='phis' + str(n))(vd)
        vd = LeakyReLU()(vd)
        vd = Dropout(ndr)(vd)
    vd = Dense(nr*nx, name='phisf')(vd)
    vd = LeakyReLU()(vd)
    vd = Dropout(ndr)(vd)

    for n in range(nn):
        vq = Dense(2*nr*nx, use_bias=True, name='phip' + str(n))(vq)
        vq = LeakyReLU()(vq)
        vq = Dropout(ndr)(vq)
    vq = Dense(2*nr*nx, name='phipf')(vq)
    vq = LeakyReLU()(vq)
    vq = Dropout(ndr)(vq)

    for n in range(nn):
        vd = Dense(nfc, use_bias=True, name='dense1' + str(n))(vd)
        vd = LeakyReLU()(vd)
        vd = Dropout(ndr)(vd)
    vd = Dense(nfc, use_bias=True, name='dense1f')(vd)
    vd = LeakyReLU()(vd)
    vd = Dropout(ndr)(vd)

    for n in range(nn):
        vq = Dense(nfc, use_bias=True, name='dense2' + str(n))(vq)
        vq = LeakyReLU()(vq)
        vq = Dropout(ndr)(vq)
    vq = Dense(nfc, use_bias=True, name='dense2f')(vq)
    vq = LeakyReLU()(vq)
    vq = Dropout(ndr)(vq)
    vq = Lambda(lambda x: x[:, 0, :], output_shape=(lambda shape: (shape[0], shape[2])))(vq)

    cr = multiply([vd, vq])
    cr = Dense(nfc, use_bias=True, name='dense31')(cr)
    cr = LeakyReLU()(cr)
    cr = Dropout(ndr)(cr)
    cr = Dense(nfc, use_bias=True, name='dense32')(cr)
    cr = LeakyReLU()(cr)
    cr = Dropout(ndr)(cr)

    cr = Dense(1, name='dense33')(cr)

    counter = Model(i, cr)
    return counter

def get_query_model(nr, nx, nfc = 128, nn=3):
    ndr = 0.3
    sketch = Input(shape=(1, nr*nx))
    predicate_emb = Input(shape=(1, 2 * nfc*nx))

    vd = sketch
    vq = predicate_emb

    for n in range(nn):
        vd = Dense(nfc, use_bias=True, name='dense1' + str(n))(vd)
        vd = LeakyReLU()(vd)
        vd = Dropout(ndr)(vd)
    vd = Dense(nfc, use_bias=True, name='dense1f')(vd)
    vd = LeakyReLU()(vd)
    vd = Dropout(ndr)(vd)

    for n in range(nn):
        vq = Dense(nfc, use_bias=True, name='dense2' + str(n))(vq)
        vq = LeakyReLU()(vq)
        vq = Dropout(ndr)(vq)
    vq = Dense(nfc, use_bias=True, name='dense2f')(vq)
    vq = LeakyReLU()(vq)
    vq = Dropout(ndr)(vq)
    vq = Lambda(lambda x: x[:, 0, :], output_shape=(lambda shape: (shape[0], shape[2])))(vq)

    cr = multiply([vd, vq])

    cr = Dense(nfc, use_bias=True, name='dense31')(cr)
    cr = LeakyReLU()(cr)
    cr = Dropout(ndr)(cr)
    cr = Dense(nfc, use_bias=True, name='dense32')(cr)
    cr = LeakyReLU()(cr)
    cr = Dropout(ndr)(cr)

    cr = Dense(1, name='dense33')(cr)

    counter = Model([sketch, predicate_emb], cr)
    return counter

def get_data_emb_model(nr, nx, nfc = 128, nn=3):
    ndr = 0.3
    sketch = Input(shape=(1, 2*nr*nx))
    predicate_emb = Input(shape=(1, 4 * nfc * nx))

    vd = sketch
    for n in range(nn):
        vd = Dense(nr*nx, use_bias=True, name='phis' + str(n))(vd)
        vd = LeakyReLU()(vd)
        vd = Dropout(ndr)(vd)
    vd = Dense(nr*nx, name='phisf')(vd)
    vd = LeakyReLU()(vd)
    vd = Dropout(ndr)(vd)

    vq = predicate_emb
    for n in range(nn):
        vq = Dense(2*nr*nx, use_bias=True, name='phip' + str(n))(vq)
        vq = LeakyReLU()(vq)
        vq = Dropout(ndr)(vq)
    vq = Dense(2*nr*nx, name='phipf')(vq)
    vq = LeakyReLU()(vq)
    vq = Dropout(ndr)(vq)

    emb = Model([sketch, predicate_emb], [vd, vq])
    return emb


def get_merge_emb_model(nr, nfc = 128, nn=3):
    # nr: embedding size;
    # nfc: FC layers size; nn: FC layers

    ndr = 0.3
    i = Input(shape=(1, 2*nr))
    x = Lambda(lambda x: x)(i)

    for n in range(nn):
        x = Dense(nfc, use_bias=True, name='Phi' + str(n))(x)
        x = LeakyReLU()(x)
        x = Dropout(ndr)(x)
    x = Dense(nfc, name='Phif')(x)
    x = LeakyReLU()(x)
    x = Dropout(ndr)(x)

    emb = Model(i, x)

    return emb