from keras import Input, Model
from keras.layers import Lambda, Dense, Conv2D, Reshape, GlobalAveragePooling2D, BatchNormalization, \
    LeakyReLU, MaxPooling2D, Dropout, Flatten


def cnn_net_tiny(input_shape, output_shape=6270):
    _input = Input(input_shape)

    use_bias = False
    down0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(_input)
    down0 = BatchNormalization()(down0)
    down0 = LeakyReLU(alpha=0.1)(down0)
    down0 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(down0)
    down0 = BatchNormalization()(down0)
    down0 = LeakyReLU(alpha=0.1)(down0)
    down0_pool = MaxPooling2D((2, 2), strides=(2, 2))(down0)

    down1 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(down0_pool)
    down1 = BatchNormalization()(down1)
    down1 = LeakyReLU(alpha=0.1)(down1)
    down1 = Conv2D(32, (3, 3), padding='same', use_bias=use_bias)(down1)
    down1 = BatchNormalization()(down1)
    down1 = LeakyReLU(alpha=0.1)(down1)
    down1_pool = MaxPooling2D((2, 2), strides=(2, 2))(down1)

    down2 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(down1_pool)
    down2 = BatchNormalization()(down2)
    down2 = LeakyReLU(alpha=0.1)(down2)
    down2 = Conv2D(64, (3, 3), padding='same', use_bias=use_bias)(down2)
    down2 = BatchNormalization()(down2)
    down2 = LeakyReLU(alpha=0.1)(down2)
    down2_pool = MaxPooling2D((2, 2), strides=(2, 2))(down2)

    conv = Conv2D(64, (3, 3))(down2_pool)
    bn = BatchNormalization()(conv)
    rl = LeakyReLU(alpha=0.1)(bn)
    conv = Conv2D(64, (3, 3))(rl)
    bn = BatchNormalization()(conv)
    rl = LeakyReLU(alpha=0.1)(bn)

    # conv = Conv2D(output_shape, (1, 1), activation='softmax')(rl)
    # pool = GlobalAveragePooling2D()(conv)
    # x = Reshape((output_shape,))(pool)

    rl = Flatten()(rl)
    dense = Dense(16, activation="relu")(rl)
    drop = Dropout(0.2)(dense)
    dense = Dense(output_shape, activation="softmax")(drop)
    drop = Dropout(0.2)(dense)
    x = Reshape((output_shape,))(drop)

    model = Model(_input, x)
    model.summary()
    return model