import sys

import tensorflow.keras as K
import ssl
ssl._create_default_https_context = ssl._create_unverified_context
import json




def main():
    savefolder=sys.argv[1]
    modelname = sys.argv[2]
    customjson=sys.argv[3]
    tf_finetune=sys.argv[4]
    freezemodel=sys.argv[5]
    dropoutrate=sys.argv[6]
    predictclass=sys.argv[7]
    weights=sys.argv[8]
    print(sys.argv)
    print(customjson)
    f = open(customjson, 'r')
    data = json.load(f)
    r = json.dumps(data)
    kerasmodel = K.models.model_from_json(r)
    input = kerasmodel.input
    pretrainmodel=""
    top_dropout_rate=float(dropoutrate)
    if tf_finetune=="false":
        if modelname == "ResNet50":
            if len(weights) == 1:
                pretrainmodel = K.applications.ResNet50(weights="imagenet")
        #        else:
        #            pretrainmodel = K.applications.ResNet50(weights=weights)
        #            pretrainmodel = ResNet50(weights=weights)
        if modelname == "VGG16":
            if len(weights) == 1:
                pretrainmodel = K.applications.VGG16(weights="imagenet")
        #        else:
        #            import tensorflow.keras.applications.VGG16
        #            pretrainmodel = VGG16(weights=weights)
        if modelname == "VGG19":
            if len(weights) == 1:
                pretrainmodel = K.applications.VGG19(weights="imagenet")
        #        else:
        #            import tensorflow.keras.applications.VGG19
        #            pretrainmodel = VGG19(weights=weights)
        if modelname == "Xception":
            if len(weights) == 1:
                pretrainmodel = K.applications.Xception(weights="imagenet")
        #        else:
        #            import tensorflow.keras.applications.Xception
        #            pretrainmodel = Xception(weights=weights)
        if modelname == "EfficientNetB0":
            if len(weights) == 1:
                pretrainmodel = K.applications.EfficientNetB0(weights="imagenet", drop_connect_rate=top_dropout_rate)

        if modelname == "EfficientNetB1":
            if len(weights) == 1:
                pretrainmodel = K.applications.EfficientNetB1(weights="imagenet", drop_connect_rate=top_dropout_rate)

        if modelname == "EfficientNetB2":
            if len(weights) == 1:
                pretrainmodel = K.applications.EfficientNetB2(weights="imagenet", drop_connect_rate=top_dropout_rate)

        if modelname == "EfficientNetB3":
            if len(weights) == 1:
                pretrainmodel = K.applications.EfficientNetB3(weights="imagenet", drop_connect_rate=top_dropout_rate)

        if modelname == "EfficientNetB4":
            if len(weights) == 1:
                pretrainmodel = K.applications.EfficientNetB4(weights="imagenet", drop_connect_rate=top_dropout_rate)

        if modelname == "EfficientNetB5":
            if len(weights) == 1:
                pretrainmodel = K.applications.EfficientNetB5(weights="imagenet", drop_connect_rate=top_dropout_rate)

        if modelname == "EfficientNetB6":
            if len(weights) == 1:
                pretrainmodel = K.applications.EfficientNetB6(weights="imagenet", drop_connect_rate=top_dropout_rate)

        if modelname == "EfficientNetB7":
            if len(weights) == 1:
                pretrainmodel = K.applications.EfficientNetB7(weights="imagenet", drop_connect_rate=top_dropout_rate)
    else:
        print("transfer_learning is true")
        if modelname == "ResNet50":
            if len(weights) == 1:
                try:
                    print('get pretrainmodel')
                    pretrainmodel = K.applications.ResNet50(include_top=False,input_tensor=input,weights="imagenet")
                except:
                    print('cannot get ResNet50 with no top')
        #        else:
        #            pretrainmodel = K.applications.ResNet50(weights=weights)
        #            pretrainmodel = ResNet50(weights=weights)
        if modelname == "VGG16":
            if len(weights) == 1:
                pretrainmodel = K.applications.VGG16(include_top=False,input_tensor=input,weights="imagenet")
        #        else:
        #            import tensorflow.keras.applications.VGG16
        #            pretrainmodel = VGG16(weights=weights)
        if modelname == "VGG19":
            if len(weights) == 1:
                pretrainmodel = K.applications.VGG19(include_top=False,weights="imagenet")
        #        else:
        #            import tensorflow.keras.applications.VGG19
        #            pretrainmodel = VGG19(weights=weights)
        if modelname == "Xception":
            if len(weights) == 1:
                pretrainmodel = K.applications.Xception(include_top=False,weights="imagenet")
        #        else:
        #            import tensorflow.keras.applications.Xception
        #            pretrainmodel = Xception(weights=weights)
        if modelname == "EfficientNetB0":
            if len(weights) == 1:
                pretrainmodel = K.applications.EfficientNetB0(include_top=False,weights="imagenet",input_tensor=input, drop_connect_rate=top_dropout_rate)

        if modelname == "EfficientNetB1":
            if len(weights) == 1:
                pretrainmodel = K.applications.EfficientNetB1(include_top=False,weights="imagenet", drop_connect_rate=top_dropout_rate)

        if modelname == "EfficientNetB2":
            if len(weights) == 1:
                pretrainmodel = K.applications.EfficientNetB2(include_top=False,weights="imagenet", drop_connect_rate=top_dropout_rate)

        if modelname == "EfficientNetB3":
            if len(weights) == 1:
                pretrainmodel = K.applications.EfficientNetB3(include_top=False,weights="imagenet", drop_connect_rate=top_dropout_rate)

        if modelname == "EfficientNetB4":
            if len(weights) == 1:
                pretrainmodel = K.applications.EfficientNetB4(include_top=False,weights="imagenet", drop_connect_rate=top_dropout_rate)

        if modelname == "EfficientNetB5":
            if len(weights) == 1:
                pretrainmodel = K.applications.EfficientNetB5(include_top=False,weights="imagenet", drop_connect_rate=top_dropout_rate)

        if modelname == "EfficientNetB6":
            if len(weights) == 1:
                pretrainmodel = K.applications.EfficientNetB6(include_top=False,weights="imagenet", drop_connect_rate=top_dropout_rate)

        if modelname == "EfficientNetB7":
            if len(weights) == 1:
                pretrainmodel = K.applications.EfficientNetB7(include_top=False,weights="imagenet", drop_connect_rate=top_dropout_rate)
        #    if modelname == "usermodel":
        #        pretrainmodel=kerasmodel
    print(type('pretrainmodel type:'), pretrainmodel)
    if type(pretrainmodel) != type(""):
        if freezemodel == "true":
            pretrainmodel.trainable = False
            for l in pretrainmodel.layers[-20:]:
                if not isinstance(l, K.layers.BatchNormalization):
                    l.trainable = True
        if tf_finetune == "true":
            pretrainlastlayer = pretrainmodel.layers.pop(-1)
            print(pretrainmodel.summary())
            newlayer = kerasmodel.layers.pop(-1)
            print(pretrainlastlayer.output.shape[1], newlayer.output.shape[1])
            x = K.layers.GlobalAveragePooling2D(name="avg_pool")(pretrainmodel.layers[-2].output)
            x = K.layers.BatchNormalization()(x)

            x = K.layers.Dropout(top_dropout_rate)(x)
            NUM_CLASSES = int(predictclass)
            outputs = K.layers.Dense(NUM_CLASSES, activation="softmax", name="pred")(x)
            kerasmodel = K.models.Model(input, outputs)

        else:
            pretrainlastlayer = pretrainmodel.layers.pop(-1)
            newlayer = kerasmodel.layers.pop(-1)
            print("custom_model.json output shape:",newlayer.output.shape[1])
            # if pretrainlastlayer.output.shape[1] == newlayer.output.shape[1]:
            #     kerasmodel = pretrainmodel
            # else:
            input = pretrainmodel.input
            out = newlayer(pretrainmodel.layers[-2].output)
            kerasmodel = K.models.Model(input, out)
            #            if len(weights)>0:
            #                kerasmodel.load_weights(weights)
        print(kerasmodel.summary())
    else:
        kerasmodel.load_weights(weights)
    kerasmodel.save(savefolder)
    print("save model to disk:",savefolder)


if __name__ == '__main__':
    main()
