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


def cnn_net(input_shape, output_shape=5710):
    _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(64, (3, 3), padding='same', use_bias=use_bias)(down0_pool)
    down1 = BatchNormalization()(down1)
    down1 = LeakyReLU(alpha=0.1)(down1)
    down1 = Conv2D(64, (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(128, (3, 3), padding='same', use_bias=use_bias)(down1_pool)
    down2 = BatchNormalization()(down2)
    down2 = LeakyReLU(alpha=0.1)(down2)
    down2 = Conv2D(128, (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(256, (3, 3))(down2_pool)
    bn = BatchNormalization()(conv)
    rl = LeakyReLU(alpha=0.1)(bn)
    conv = Conv2D(256, (3, 3))(rl)
    bn = BatchNormalization()(conv)
    rl = LeakyReLU(alpha=0.1)(bn)

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

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