from keras.layers import Conv2D, Activation, Input, Concatenate, LeakyReLU, Lambda, AveragePooling2D, UpSampling2D, Convolution2D, BatchNormalization, Conv2DTranspose, Add
from keras.models import Model
from InstanceNorm import InstanceNormalization


def make_standard_UNET(channels,outs):

    def relu(x):
        return Activation('relu')(x)

    def concat(x):
        return Concatenate()(x)

    c0 = Convolution2D(filters=32, kernel_size=3, strides=1, padding='same', name='c0')
    c1 = Convolution2D(filters=64, kernel_size=4, strides=2, padding='same', name='c1')
    c2 = Convolution2D(filters=64, kernel_size=3, strides=1, padding='same', name='c2')
    c3 = Convolution2D(filters=128, kernel_size=4, strides=2, padding='same', name='c3')
    c4 = Convolution2D(filters=128, kernel_size=3, strides=1, padding='same', name='c4')
    c5 = Convolution2D(filters=256, kernel_size=4, strides=2, padding='same', name='c5')
    c6 = Convolution2D(filters=256, kernel_size=3, strides=1, padding='same', name='c6')
    c7 = Convolution2D(filters=512, kernel_size=4, strides=2, padding='same', name='c7')
    c8 = Convolution2D(filters=512, kernel_size=3, strides=1, padding='same', name='c8')

    bnc0 = BatchNormalization(axis=3, name='bnc0')
    bnc1 = BatchNormalization(axis=3, name='bnc1')
    bnc2 = BatchNormalization(axis=3, name='bnc2')
    bnc3 = BatchNormalization(axis=3, name='bnc3')
    bnc4 = BatchNormalization(axis=3, name='bnc4')
    bnc5 = BatchNormalization(axis=3, name='bnc5')
    bnc6 = BatchNormalization(axis=3, name='bnc6')
    bnc7 = BatchNormalization(axis=3, name='bnc7')
    bnc8 = BatchNormalization(axis=3, name='bnc8')

    dc8 = Conv2DTranspose(filters=512, kernel_size=4, strides=2, padding='same', name='dc8_')
    dc7 = Convolution2D(filters=256, kernel_size=3, strides=1, padding='same', name='dc7')
    dc6 = Conv2DTranspose(filters=256, kernel_size=4, strides=2, padding='same', name='dc6_')
    dc5 = Convolution2D(filters=128, kernel_size=3, strides=1, padding='same', name='dc5')
    dc4 = Conv2DTranspose(filters=128, kernel_size=4, strides=2, padding='same', name='dc4_')
    dc3 = Convolution2D(filters=64, kernel_size=3, strides=1, padding='same', name='dc3')
    dc2 = Conv2DTranspose(filters=64, kernel_size=4, strides=2, padding='same', name='dc2_')
    dc1 = Convolution2D(filters=32, kernel_size=3, strides=1, padding='same', name='dc1')
    dc0 = Convolution2D(filters=outs, kernel_size=3, strides=1, padding='same', name='dc0')

    bnd1 = BatchNormalization(axis=3, name='bnd1')
    bnd2 = BatchNormalization(axis=3, name='bnd2')
    bnd3 = BatchNormalization(axis=3, name='bnd3')
    bnd4 = BatchNormalization(axis=3, name='bnd4')
    bnd5 = BatchNormalization(axis=3, name='bnd5')
    bnd6 = BatchNormalization(axis=3, name='bnd6')
    bnd7 = BatchNormalization(axis=3, name='bnd7')
    bnd8 = BatchNormalization(axis=3, name='bnd8')

    x = Input(shape=(128, 128, channels))

    e0 = relu(bnc0(c0(x), training = False))
    e1 = relu(bnc1(c1(e0), training = False))
    e2 = relu(bnc2(c2(e1), training = False))
    e3 = relu(bnc3(c3(e2), training = False))
    e4 = relu(bnc4(c4(e3), training = False))
    e5 = relu(bnc5(c5(e4), training = False))
    e6 = relu(bnc6(c6(e5), training = False))
    e7 = relu(bnc7(c7(e6), training = False))
    e8 = relu(bnc8(c8(e7), training = False))

    d8 = relu(bnd8(dc8(concat([e7, e8])), training = False))
    d7 = relu(bnd7(dc7(d8), training = False))
    d6 = relu(bnd6(dc6(concat([e6, d7])), training = False))
    d5 = relu(bnd5(dc5(d6), training = False))
    d4 = relu(bnd4(dc4(concat([e4, d5])), training = False))
    d3 = relu(bnd3(dc3(d4), training = False))
    d2 = relu(bnd2(dc2(concat([e2, d3])), training = False))
    d1 = relu(bnd1(dc1(d2), training = False))
    d0 = dc0(concat([e0, d1]))

    model = Model(inputs=x,outputs=d0)

    return model


def make_diff_net():

    def conv(x, filters, name):
        return Conv2D(filters=filters, strides=(1, 1), kernel_size=(3, 3), padding='same', name=name)(x)

    def relu(x):
        return Activation('relu')(x)

    def lrelu(x):
        return LeakyReLU(alpha=0.1)(x)

    def r_block(x, filters, name=None):
        return relu(conv(relu(conv(x, filters, None if name is None else name + '_c1')), filters,
                         None if name is None else name + '_c2'))

    def cat(a, b):
        return Concatenate()([UpSampling2D((2, 2))(a), b])

    def dog(x):
        down = AveragePooling2D((2, 2))(x)
        up = UpSampling2D((2, 2))(down)
        diff = Lambda(lambda p: p[0] - p[1])([x, up])
        return down, diff

    ip = Input(shape=(512, 512, 3))

    c512 = r_block(ip, 16, 'c512')

    c256, l512 = dog(c512)
    c256 = r_block(c256, 32, 'c256')

    c128, l256 = dog(c256)
    c128 = r_block(c128, 64, 'c128')

    c64, l128 = dog(c128)
    c64 = r_block(c64, 128, 'c64')

    c32, l64 = dog(c64)
    c32 = r_block(c32, 256, 'c32')

    c16, l32 = dog(c32)
    c16 = r_block(c16, 512, 'c16')

    d32 = cat(c16, l32)
    d32 = r_block(d32, 256, 'd32')

    d64 = cat(d32, l64)
    d64 = r_block(d64, 128, 'd64')

    d128 = cat(d64, l128)
    d128 = r_block(d128, 64, 'd128')

    d256 = cat(d128, l256)
    d256 = r_block(d256, 32, 'd256')

    d512 = cat(d256, l512)
    d512 = r_block(d512, 16, 'd512')

    op = conv(d512, 1, 'op')

    return Model(inputs=ip, outputs=op)


def make_wnet256():

    def conv(x, filters):
        return Conv2D(filters=filters, strides=(1, 1), kernel_size=(3, 3), padding='same')(x)

    def relu(x):
        return Activation('relu')(x)

    def lrelu(x):
        return LeakyReLU(alpha=0.1)(x)

    def r_block(x, filters):
        return relu(conv(relu(conv(x, filters)), filters))

    def res_block(x, filters):
        return relu(Add()([x, conv(relu(conv(x, filters)), filters)]))

    def cat(a, b):
        return Concatenate()([UpSampling2D((2, 2))(a), b])

    def dog(x):
        down = AveragePooling2D((2, 2))(x)
        up = UpSampling2D((2, 2))(down)
        diff = Lambda(lambda p: p[0] - p[1])([x, up])
        return down, diff

    ip_sketch = Input(shape=(256, 256, 1))
    ip_color = Input(shape=(256, 256, 3))

    c256 = r_block(ip_sketch, 32)

    c128, l256 = dog(c256)
    c128 = r_block(c128, 64)

    c64, l128 = dog(c128)
    c64 = r_block(c64, 128)

    c32, l64 = dog(c64)
    c32 = r_block(Concatenate()([c32, AveragePooling2D((8, 8))(ip_color)]), 256)

    c32 = res_block(c32, 256)
    c32 = res_block(c32, 256)
    c32 = res_block(c32, 256)

    c32 = res_block(c32, 256)
    c32 = res_block(c32, 256)
    c32 = res_block(c32, 256)

    c32 = res_block(c32, 256)
    c32 = res_block(c32, 256)
    c32 = res_block(c32, 256)

    d64 = cat(c32, l64)
    d64 = r_block(d64, 128)

    d128 = cat(d64, l128)
    d128 = r_block(d128, 64)

    d256 = cat(d128, l256)
    d256 = r_block(d256, 32)

    op = conv(d256, 3)

    return Model(inputs=[ip_sketch, ip_color], outputs=op)


def make_unet512():

    def conv(x, filters, strides=(1, 1), kernel_size=(3, 3)):
        return Conv2D(filters=filters, strides=strides, kernel_size=kernel_size, padding='same')(x)

    def donv(x, filters, strides=(2, 2), kernel_size=(4, 4)):
        return Conv2DTranspose(filters=filters, strides=strides, kernel_size=kernel_size, padding='same')(x)

    def relu(x):
        return Activation('relu')(x)

    def sigmoid(x):
        return Activation('sigmoid')(x)

    def norm(x):
        return InstanceNormalization(axis=3)(x)

    def cat(a, b):
        return Concatenate()([a, b])

    def res(x, filters):
        c1 = relu(norm(conv(x, filters // 2)))
        c2 = norm(conv(c1, filters))
        ad = Add()([x, c2])
        return relu(ad)

    ip = Input(shape=(512, 512, 3))

    c512 = relu(norm(conv(ip, 16, strides=(1, 1), kernel_size=(3, 3))))
    c256 = relu(norm(conv(c512, 32, strides=(2, 2), kernel_size=(4, 4))))

    c128 = relu(norm(conv(c256, 64, strides=(2, 2), kernel_size=(4, 4))))
    c128 = res(c128, 64)

    c64 = relu(norm(conv(c128, 128, strides=(2, 2), kernel_size=(4, 4))))
    c64 = res(c64, 128)
    c64 = res(c64, 128)

    c32 = relu(norm(conv(c64, 256, strides=(2, 2), kernel_size=(4, 4))))
    c32 = res(c32, 256)
    c32 = res(c32, 256)
    c32 = res(c32, 256)
    c32 = res(c32, 256)
    c32 = res(c32, 256)
    c32 = res(c32, 256)
    c32 = res(c32, 256)
    c32 = res(c32, 256)

    c16 = relu(norm(conv(c32, 512, strides=(2, 2), kernel_size=(4, 4))))
    c16 = res(c16, 512)
    c16 = res(c16, 512)
    c16 = res(c16, 512)
    c16 = res(c16, 512)
    c16 = res(c16, 512)
    c16 = res(c16, 512)
    c16 = res(c16, 512)
    c16 = res(c16, 512)

    c8 = relu(norm(conv(c16, 1024, strides=(2, 2), kernel_size=(4, 4))))
    c8 = res(c8, 1024)
    c8 = res(c8, 1024)
    c8 = res(c8, 1024)
    c8 = res(c8, 1024)

    e16 = relu(norm(donv(c8, 512, strides=(2, 2), kernel_size=(4, 4))))
    e16 = cat(e16, c16)
    e16 = relu(norm(conv(e16, 512, strides=(1, 1), kernel_size=(3, 3))))

    e32 = relu(norm(donv(e16, 256, strides=(2, 2), kernel_size=(4, 4))))
    e32 = cat(e32, c32)
    e32 = relu(norm(conv(e32, 256, strides=(1, 1), kernel_size=(3, 3))))

    e64 = relu(norm(donv(e32, 128, strides=(2, 2), kernel_size=(4, 4))))
    e64 = cat(e64, c64)
    e64 = relu(norm(conv(e64, 128, strides=(1, 1), kernel_size=(3, 3))))

    e128 = relu(norm(donv(e64, 64, strides=(2, 2), kernel_size=(4, 4))))
    e128 = cat(e128, c128)
    e128 = relu(norm(conv(e128, 64, strides=(1, 1), kernel_size=(3, 3))))

    e256 = relu(norm(donv(e128, 32, strides=(2, 2), kernel_size=(4, 4))))
    e256 = cat(e256, c256)
    e256 = relu(norm(conv(e256, 32, strides=(1, 1), kernel_size=(3, 3))))

    e512 = relu(norm(donv(e256, 16, strides=(2, 2), kernel_size=(4, 4))))
    e512 = cat(e512, c512)
    e512 = relu(norm(conv(e512, 16, strides=(1, 1), kernel_size=(3, 3))))

    ot = sigmoid(conv(e512, 1))

    return Model(inputs=ip, outputs=ot)