#!/usr/bin/env python3

from tensorflow.keras import Input, Model
from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten

def edsr_residual(x, num_filters, scaling):
    res = Conv2D(num_filters, 3, padding='same', activation='relu')(x)
    res = Conv2D(num_filters, 3, padding='same')(res)

    if scaling:
        res = Lambda(lambda t: t * scaling)(res)

    res = Add()([res, x])
    return res

def upsampling_block(model, num_filters, kernal_size, strides):
    model = Conv2D(num_filters, kernal_size, strides, padding="same")(model)
    model = UpSampling2D(size=2)(model)
    model = PReLU(alpha_initializer='zeros', alpha_regularizer=None, alpha_constraint=None, shared_axes=[1,2])(model)
    return model

def gridless_upsampling(model, num_filters, scale):
    def upsample(x, factor):
        x = UpSampling2D(size=factor)(model)
        x = Conv2D(num_filters, 3, padding='same')(x)
        return x

    if scale==2:
        model = upsample(model, 2)
    elif scale==3:
        model = upsample(model, 3)
    elif scale==4:
        model = upsample(model, 2)
        model = upsample(model, 2)

    return model

def build_edsr(input_shape, scale, num_filters, residual_blocks, res_block_scaling=0.1):
    input_layer = Input(shape=input_shape)

    gen1 = Conv2D(num_filters, kernel_size=9, padding='same')(input_layer)
    res = edsr_residual(gen1, num_filters, res_block_scaling)

    for i in range(residual_blocks - 1):
        res = edsr_residual(res, num_filters, res_block_scaling)

    gen2 = Conv2D(num_filters, kernel_size=3, padding='same')(res)
    model = Add()([gen2, gen1])

    for index in range(2):
        model = upsampling_block(model, num_filters, 3, 1)

    # model = gridless_upsampling(model, num_filters, scale)

    output = Conv2D(3, 9, padding='same')(model)
    output = Activation('tanh')(output)

    model = Model(inputs=[input_layer], outputs=[output], name='edsr_generator')
    return model

def discriminator_block(model, num_filters, kernel_size, strides):
    model = Conv2D(num_filters, kernel_size, strides, padding='same')(model)
    model = BatchNormalization(momentum=0.5)(model)
    model = LeakyReLU(alpha=0.2)(model)
    return model

def build_discriminator(input_shape, num_filters=64):
    input_layer = Input(shape=input_shape)

    dis1 = Conv2D(num_filters, 3, padding='same')(input_layer)
    dis1 = LeakyReLU(alpha=0.2)(dis1)

    dis2 = discriminator_block(dis1, num_filters, 3, 2)

    dis3 = discriminator_block(dis2, num_filters * 2, 3, 1)
    dis4 = discriminator_block(dis3, num_filters * 2, 3, 2)

    dis5 = discriminator_block(dis4, num_filters * 4, 3, 1)
    dis6 = discriminator_block(dis5, num_filters * 4, 3, 2)

    dis7 = discriminator_block(dis6, num_filters * 8, 3, 1)
    dis8 = discriminator_block(dis7, num_filters * 8, 3, 2)

    dis9 = Flatten()(dis8)
    dis9 = Dense(1024)(dis9)
    dis9 = LeakyReLU(alpha=0.2)(dis9)

    output = Dense(units=1)(dis9)
    output = Activation('sigmoid')(output)

    model = Model(inputs=[input_layer], outputs=[output], name='discriminator')
    return model

tensorflow.keras.layers.Activation tensorflow.keras.Input tensorflow.keras.layers.LeakyReLU tensorflow.keras.layers.Lambda tensorflow.keras.layers.PReLU tensorflow.keras.layers.Dense tensorflow.keras.layers.Conv2D tensorflow.keras.layers.UpSampling2D tensorflow.keras.Model tensorflow.keras.layers.BatchNormalization tensorflow.keras.layers.Add tensorflow.keras.layers.Flatten
[(39, 'tensorflow.keras.Input', 'Input', ([], {'shape': 'input_shape'}), False, 'from tensorflow.keras import Input, Model\n'), (58, 'tensorflow.keras.Model', 'Model', ([], {'inputs': '[input_layer]', 'outputs': '[output]', 'name': '"""edsr_generator"""'}), False, 'from tensorflow.keras import Input, Model\n'), (68, 'tensorflow.keras.Input', 'Input', ([], {'shape': 'input_shape'}), False, 'from tensorflow.keras import Input, Model\n'), (91, 'tensorflow.keras.Model', 'Model', ([], {'inputs': '[input_layer]', 'outputs': '[output]', 'name': '"""discriminator"""'}), False, 'from tensorflow.keras import Input, Model\n'), (7, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['num_filters', '(3)'], {'padding': '"""same"""', 'activation': '"""relu"""'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (8, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['num_filters', '(3)'], {'padding': '"""same"""'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (13, 'tensorflow.keras.layers.Add', 'Add', ([], {}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (17, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['num_filters', 'kernal_size', 'strides'], {'padding': '"""same"""'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (18, 'tensorflow.keras.layers.UpSampling2D', 'UpSampling2D', ([], {'size': '(2)'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (19, 'tensorflow.keras.layers.PReLU', 'PReLU', ([], {'alpha_initializer': '"""zeros"""', 'alpha_regularizer': 'None', 'alpha_constraint': 'None', 'shared_axes': '[1, 2]'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (41, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['num_filters'], {'kernel_size': '(9)', 'padding': '"""same"""'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (47, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['num_filters'], {'kernel_size': '(3)', 'padding': '"""same"""'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (48, 'tensorflow.keras.layers.Add', 'Add', ([], {}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (55, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(3)', '(9)'], {'padding': '"""same"""'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (56, 'tensorflow.keras.layers.Activation', 'Activation', (['"""tanh"""'], {}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (62, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['num_filters', 'kernel_size', 'strides'], {'padding': '"""same"""'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (63, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {'momentum': '(0.5)'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (64, 'tensorflow.keras.layers.LeakyReLU', 'LeakyReLU', ([], {'alpha': '(0.2)'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (70, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['num_filters', '(3)'], {'padding': '"""same"""'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (71, 'tensorflow.keras.layers.LeakyReLU', 'LeakyReLU', ([], {'alpha': '(0.2)'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (84, 'tensorflow.keras.layers.Flatten', 'Flatten', ([], {}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (85, 'tensorflow.keras.layers.Dense', 'Dense', (['(1024)'], {}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (86, 'tensorflow.keras.layers.LeakyReLU', 'LeakyReLU', ([], {'alpha': '(0.2)'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (88, 'tensorflow.keras.layers.Dense', 'Dense', ([], {'units': '(1)'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (89, 'tensorflow.keras.layers.Activation', 'Activation', (['"""sigmoid"""'], {}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (11, 'tensorflow.keras.layers.Lambda', 'Lambda', (['(lambda t: t * scaling)'], {}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (24, 'tensorflow.keras.layers.UpSampling2D', 'UpSampling2D', ([], {'size': 'factor'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n'), (25, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['num_filters', '(3)'], {'padding': '"""same"""'}), False, 'from tensorflow.keras.layers import Conv2D, Lambda, BatchNormalization, UpSampling2D, Activation, LeakyReLU, PReLU, Add, Dense, Flatten\n')]