from tensorflow.python.keras import Sequential, regularizers, Input, Model
from tensorflow.python.keras.layers import Flatten, Dense, Conv2D, MaxPooling2D, Activation, BatchNormalization, Dropout
from varname.helpers import Wrapper

import datasets
from filterlayer.fileter_layer import FilterLayer
from train import train_model

weight_decay = 1e-4
no_of_filters = 48
blueprint = [2, 2]
dense_layer_factor = 1.5
dropout_factor = 0.2


def create_ext_cnn(input_shape, no_of_classes):
    global no_of_filters
    global dropout_factor
    global dense_layer_factor

    no_of_filters = 48
    dense_layer_factor = 1.5
    dropout_factor = 0.2

    input_layer = Input(shape=input_shape)

    output_tensor = input_layer

    output_tensor = FilterLayer(input_shape=input_shape, input_dim=input_shape[0])(output_tensor)
    output_tensor = BatchNormalization()(output_tensor)
    output_tensor = MaxPooling2D(pool_size=(2, 2))(output_tensor)

    for i in range(len(blueprint)):
        no_of_conv_layers = blueprint[i]
        for j in range(no_of_conv_layers):
            output_tensor = Conv2D(no_of_filters, (3, 3), padding='same',
                                   kernel_regularizer=regularizers.l2(weight_decay))(output_tensor)
            output_tensor = Activation('relu')(output_tensor)
            output_tensor = BatchNormalization()(output_tensor)
        output_tensor = MaxPooling2D(pool_size=(2, 2))(output_tensor)
        output_tensor = Dropout(dropout_factor)(output_tensor)

        no_of_filters = no_of_filters * 2
        dropout_factor = 0.5 if dropout_factor >= 0.5 else dropout_factor + 0.1

    output_tensor = Flatten()(output_tensor)
    output_tensor = Dense(no_of_classes * dense_layer_factor, activation='relu')(output_tensor)
    output_tensor = Dropout(0.25)(output_tensor)
    output_layer = Dense(no_of_classes, activation='softmax')(output_tensor)

    model = Model(input_layer, output_layer)

    model.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )

    return model


def train_ext_cnn_cifar10(epochs):
    ds_train, ds_validation, ds_test, ds_class_no, ds_class_names = datasets.get_cifar_10()
    input_shape = ds_train.element_spec[0].shape[1:]
    ext_cnn_cifar10 = Wrapper(create_ext_cnn(input_shape, ds_class_no))
    train_model(ext_cnn_cifar10.value, ext_cnn_cifar10.name,
                ds_train, ds_validation, ds_test, ds_class_names, epochs)


def train_ext_cnn_cifar100(epochs):
    ds_train, ds_validation, ds_test, ds_class_no, ds_class_names = datasets.get_cifar_100()
    input_shape = ds_train.element_spec[0].shape[1:]
    ext_cnn_cifar100 = Wrapper(create_ext_cnn(input_shape, ds_class_no))
    train_model(ext_cnn_cifar100.value, ext_cnn_cifar100.name,
                ds_train, ds_validation, ds_test, ds_class_names, epochs)


def train_ext_cnn_imagenet64(epochs):
    ds_train, ds_validation, ds_test, ds_class_no, ds_class_names = datasets.get_imagenet_64()
    input_shape = ds_train.element_spec[0].shape[1:]
    ext_cnn_imagenet64 = Wrapper(create_ext_cnn(input_shape, ds_class_no))
    train_model(ext_cnn_imagenet64.value, ext_cnn_imagenet64.name,
                ds_train, ds_validation, ds_test, ds_class_names, epochs)

from tensorflow.python.keras import Sequential, regularizers, Input, Model
from tensorflow.python.keras.layers import Flatten, Dense, Conv2D, MaxPooling2D, Activation, BatchNormalization, Dropout
from varname.helpers import Wrapper

import datasets
from filterlayer.fileter_layer import FilterLayer
from train import train_model

weight_decay = 1e-4
no_of_filters = 48
blueprint = [2, 2]
dense_layer_factor = 1.5
dropout_factor = 0.2


def create_ext_cnn(input_shape, no_of_classes):
    global no_of_filters
    global dropout_factor
    global dense_layer_factor

    no_of_filters = 48
    dense_layer_factor = 1.5
    dropout_factor = 0.2

    input_layer = Input(shape=input_shape)

    output_tensor = input_layer

    output_tensor = FilterLayer(input_shape=input_shape, input_dim=input_shape[0])(output_tensor)
    output_tensor = BatchNormalization()(output_tensor)
    output_tensor = MaxPooling2D(pool_size=(2, 2))(output_tensor)

    for i in range(len(blueprint)):
        no_of_conv_layers = blueprint[i]
        for j in range(no_of_conv_layers):
            output_tensor = Conv2D(no_of_filters, (3, 3), padding='same',
                                   kernel_regularizer=regularizers.l2(weight_decay))(output_tensor)
            output_tensor = Activation('relu')(output_tensor)
            output_tensor = BatchNormalization()(output_tensor)
        output_tensor = MaxPooling2D(pool_size=(2, 2))(output_tensor)
        output_tensor = Dropout(dropout_factor)(output_tensor)

        no_of_filters = no_of_filters * 2
        dropout_factor = 0.5 if dropout_factor >= 0.5 else dropout_factor + 0.1

    output_tensor = Flatten()(output_tensor)
    output_tensor = Dense(no_of_classes * dense_layer_factor, activation='relu')(output_tensor)
    output_tensor = Dropout(0.25)(output_tensor)
    output_layer = Dense(no_of_classes, activation='softmax')(output_tensor)

    model = Model(input_layer, output_layer)

    model.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        metrics=['accuracy']
    )

    return model


def train_ext_cnn_cifar10(epochs):
    ds_train, ds_validation, ds_test, ds_class_no, ds_class_names = datasets.get_cifar_10()
    input_shape = ds_train.element_spec[0].shape[1:]
    ext_cnn_cifar10 = Wrapper(create_ext_cnn(input_shape, ds_class_no))
    train_model(ext_cnn_cifar10.value, ext_cnn_cifar10.name,
                ds_train, ds_validation, ds_test, ds_class_names, epochs)


def train_ext_cnn_cifar100(epochs):
    ds_train, ds_validation, ds_test, ds_class_no, ds_class_names = datasets.get_cifar_100()
    input_shape = ds_train.element_spec[0].shape[1:]
    ext_cnn_cifar100 = Wrapper(create_ext_cnn(input_shape, ds_class_no))
    train_model(ext_cnn_cifar100.value, ext_cnn_cifar100.name,
                ds_train, ds_validation, ds_test, ds_class_names, epochs)


def train_ext_cnn_imagenet64(epochs):
    ds_train, ds_validation, ds_test, ds_class_no, ds_class_names = datasets.get_imagenet_64()
    input_shape = ds_train.element_spec[0].shape[1:]
    ext_cnn_imagenet64 = Wrapper(create_ext_cnn(input_shape, ds_class_no))
    train_model(ext_cnn_imagenet64.value, ext_cnn_imagenet64.name,
                ds_train, ds_validation, ds_test, ds_class_names, epochs)
