import tensorflow as tf
from tensorflow.keras.applications.resnet50 import ResNet50
from tensorflow.keras.applications.vgg16 import VGG16
from tensorflow.keras.applications.inception_v3 import InceptionV3
from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate
from tensorflow.keras.layers import BatchNormalization, Dropout
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import SGD
import numpy as np
import tensorflow.keras.backend as K
from tensorflow.keras.callbacks import Callback, ModelCheckpoint
import sys

class ArticleType:
    '''This model is based off of VGG16 with the addition of BatchNorm layers and then branching '''

    def __init__(self, label):
        self.master_classes=4
        self.sub_classes=21
        self.art_classes=45
        
        input_image = Input(shape=(224,224,3),name="InputImg")

        #--- block 1 ---
        x = Conv2D(64, (3, 3), activation='relu', padding='same', name='block1_conv1')(input_image)
        x = BatchNormalization()(x)
        x = Conv2D(64, (3, 3), activation='relu', padding='same', name='block1_conv2')(x)
        x = BatchNormalization()(x)
        x = MaxPooling2D((2, 2), strides=(2, 2), name='block1_pool')(x)

        #--- block 2 ---
        x = Conv2D(128, (3, 3), activation='relu', padding='same', name='block2_conv1')(x)
        x = BatchNormalization()(x)
        x = Conv2D(128, (3, 3), activation='relu', padding='same', name='block2_conv2')(x)
        x = BatchNormalization()(x)
        x = MaxPooling2D((2, 2), strides=(2, 2), name='block2_pool')(x)

        #--- block 3 ---
        x = Conv2D(256, (3, 3), activation='relu', padding='same', name='block3_conv1')(x)
        x = BatchNormalization()(x)
        x = Conv2D(256, (3, 3), activation='relu', padding='same', name='block3_conv2')(x)
        x = BatchNormalization()(x)
        x = MaxPooling2D((2, 2), strides=(2, 2), name='block3_pool')(x)

        #--- block 4 ---
        x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block4_conv1')(x)
        x = BatchNormalization()(x)
        x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block4_conv2')(x)
        x = BatchNormalization()(x)
        x = MaxPooling2D((2, 2), strides=(2, 2), name='block4_pool')(x)

        #--- block 5 subCategory ---
        x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block5_conv1_art')(x)
        x = BatchNormalization()(x)
        x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block5_conv2_art')(x)
        x = BatchNormalization()(x)
        x = Conv2D(512, (3, 3), activation='relu', padding='same', name='block5_conv3_art')(x)
        x = BatchNormalization()(x)

        #Experiment with number of neurons in final layers for the baseline models.

        #--- subCategory prediction---
        x = Flatten(name='flatten')(x)
        x = Dense(4096, activation='relu', name='fc_art')(x)
        x = BatchNormalization()(x)
        x = Dropout(0.5)(x)
        x = Dense(4096, activation='relu', name='fc2_art')(x)
        x = BatchNormalization()(x)
        x = Dropout(0.5)(x)
        pred = Dense(self.art_classes, activation='softmax', name='art_output')(x)

        model = Model(
            inputs=input_image,
            outputs=pred,
            name="Baseline_articleType_CNN")

        trainable_params= np.sum([K.count_params(w) for w in model.trainable_weights])
        #trainable_params = tf.keras.backend.count_params(model.trainable_weights)
        print("Trainable paramaters: "+str(trainable_params))
        #Keras will automaticall use categorical accuracy when accuracy is used.
        model.compile(optimizer=SGD(lr=0.001, momentum=0.9), loss='categorical_crossentropy', metrics=['categorical_accuracy'])

        checkpoint = ModelCheckpoint("../weights/"+label+"_best_weights.h5", monitor='val_loss', verbose=1,
            save_best_only=True, save_weights_only=True,mode='auto')
        self.cbks = [checkpoint]
        self.model = model
tensorflow.keras.callbacks.ModelCheckpoint tensorflow.keras.models.Model tensorflow.keras.backend.count_params tensorflow.keras.layers.Dense tensorflow.keras.layers.Conv2D tensorflow.keras.optimizers.SGD tensorflow.keras.layers.BatchNormalization tensorflow.keras.layers.Dropout tensorflow.keras.layers.MaxPooling2D tensorflow.keras.layers.Flatten tensorflow.keras.layers.Input
[(22, 'tensorflow.keras.layers.Input', 'Input', ([], {'shape': '(224, 224, 3)', 'name': '"""InputImg"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (72, 'tensorflow.keras.models.Model', 'Model', ([], {'inputs': 'input_image', 'outputs': 'pred', 'name': '"""Baseline_articleType_CNN"""'}), False, 'from tensorflow.keras.models import Model\n'), (83, 'tensorflow.keras.callbacks.ModelCheckpoint', 'ModelCheckpoint', (["('../weights/' + label + '_best_weights.h5')"], {'monitor': '"""val_loss"""', 'verbose': '(1)', 'save_best_only': '(True)', 'save_weights_only': '(True)', 'mode': '"""auto"""'}), False, 'from tensorflow.keras.callbacks import Callback, ModelCheckpoint\n'), (25, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(64)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block1_conv1"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (26, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import BatchNormalization, Dropout\n'), (27, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(64)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block1_conv2"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (28, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import BatchNormalization, Dropout\n'), (29, 'tensorflow.keras.layers.MaxPooling2D', 'MaxPooling2D', (['(2, 2)'], {'strides': '(2, 2)', 'name': '"""block1_pool"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (32, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(128)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block2_conv1"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (33, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import BatchNormalization, Dropout\n'), (34, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(128)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block2_conv2"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (35, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import BatchNormalization, Dropout\n'), (36, 'tensorflow.keras.layers.MaxPooling2D', 'MaxPooling2D', (['(2, 2)'], {'strides': '(2, 2)', 'name': '"""block2_pool"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (39, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(256)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block3_conv1"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (40, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import BatchNormalization, Dropout\n'), (41, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(256)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block3_conv2"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (42, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import BatchNormalization, Dropout\n'), (43, 'tensorflow.keras.layers.MaxPooling2D', 'MaxPooling2D', (['(2, 2)'], {'strides': '(2, 2)', 'name': '"""block3_pool"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (46, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(512)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block4_conv1"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (47, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import BatchNormalization, Dropout\n'), (48, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(512)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block4_conv2"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (49, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import BatchNormalization, Dropout\n'), (50, 'tensorflow.keras.layers.MaxPooling2D', 'MaxPooling2D', (['(2, 2)'], {'strides': '(2, 2)', 'name': '"""block4_pool"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (53, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(512)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block5_conv1_art"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (54, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import BatchNormalization, Dropout\n'), (55, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(512)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block5_conv2_art"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (56, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import BatchNormalization, Dropout\n'), (57, 'tensorflow.keras.layers.Conv2D', 'Conv2D', (['(512)', '(3, 3)'], {'activation': '"""relu"""', 'padding': '"""same"""', 'name': '"""block5_conv3_art"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (58, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import BatchNormalization, Dropout\n'), (63, 'tensorflow.keras.layers.Flatten', 'Flatten', ([], {'name': '"""flatten"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (64, 'tensorflow.keras.layers.Dense', 'Dense', (['(4096)'], {'activation': '"""relu"""', 'name': '"""fc_art"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (65, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import BatchNormalization, Dropout\n'), (66, 'tensorflow.keras.layers.Dropout', 'Dropout', (['(0.5)'], {}), False, 'from tensorflow.keras.layers import BatchNormalization, Dropout\n'), (67, 'tensorflow.keras.layers.Dense', 'Dense', (['(4096)'], {'activation': '"""relu"""', 'name': '"""fc2_art"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (68, 'tensorflow.keras.layers.BatchNormalization', 'BatchNormalization', ([], {}), False, 'from tensorflow.keras.layers import BatchNormalization, Dropout\n'), (69, 'tensorflow.keras.layers.Dropout', 'Dropout', (['(0.5)'], {}), False, 'from tensorflow.keras.layers import BatchNormalization, Dropout\n'), (70, 'tensorflow.keras.layers.Dense', 'Dense', (['self.art_classes'], {'activation': '"""softmax"""', 'name': '"""art_output"""'}), False, 'from tensorflow.keras.layers import Dense, GlobalAveragePooling2D, MaxPooling2D, Flatten, Input, Conv2D, concatenate\n'), (77, 'tensorflow.keras.backend.count_params', 'K.count_params', (['w'], {}), True, 'import tensorflow.keras.backend as K\n'), (81, 'tensorflow.keras.optimizers.SGD', 'SGD', ([], {'lr': '(0.001)', 'momentum': '(0.9)'}), False, 'from tensorflow.keras.optimizers import SGD\n')]