from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.callbacks import EarlyStopping
from tensorflow.keras.callbacks import LearningRateScheduler
from tensorflow.keras.layers import Activation
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import GlobalAveragePooling2D
from tensorflow.keras.layers import Input
from tensorflow.keras.layers.experimental.preprocessing import Rescaling
from tensorflow.keras.layers.experimental.preprocessing import Resizing
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam

from tf_utils.callbacks import schedule_fn2
from tf_utils.dogsCatsDataAdvanced import DOGSCATS


IMAGENET_SIZE = 224
IMAGENET_DEPTH = 3
IMAGENET_SHAPE = (IMAGENET_SIZE, IMAGENET_SIZE, IMAGENET_DEPTH)


def build_model(img_shape, num_classes) -> Model:
    base_model = MobileNetV2(
        include_top=False,
        weights="imagenet",
        input_shape=IMAGENET_SHAPE
    )

    num_layers = len(base_model.layers)
    print(f"Number of layers in the base model: {num_layers}")
    fine_tune_at = num_layers - 10
    for layer in base_model.layers[:fine_tune_at]:
        layer.trainable = False

    input_img = Input(shape=img_shape)
    x = Rescaling(scale=2.0, offset=-1.0)(input_img)
    x = Resizing(height=IMAGENET_SIZE, width=IMAGENET_SIZE)(x)
    x = base_model(x)
    x = GlobalAveragePooling2D()(x)
    x = Dense(units=num_classes)(x)
    y_pred = Activation("softmax")(x)

    model = Model(
        inputs=[input_img],
        outputs=[y_pred]
    )

    model.summary()

    return model


if __name__ == "__main__":
    """
    Best model from chapter 9_2: 0.9034 accuracy
    Best model from chapter 9_7: 0.9614 accuracy
    """
    data = DOGSCATS()

    train_dataset = data.get_train_set()
    val_dataset = data.get_val_set()
    test_dataset = data.get_test_set()

    img_shape = data.img_shape
    num_classes = data.num_classes

    # Global params
    epochs = 100

    model = build_model(
        img_shape,
        num_classes
    )

    opt = Adam(learning_rate=5e-4)

    model.compile(
        loss="categorical_crossentropy",
        optimizer=opt,
        metrics=["accuracy"]
    )

    lrs_callback = LearningRateScheduler(
        schedule=schedule_fn2,
        verbose=1
    )

    es_callback = EarlyStopping(
        monitor="val_loss",
        patience=30,
        verbose=1,
        restore_best_weights=True
    )

    model.fit(
        train_dataset,
        verbose=1,
        epochs=epochs,
        callbacks=[lrs_callback, es_callback],
        validation_data=val_dataset,
    )

    scores = model.evaluate(
        val_dataset,
        verbose=0
    )
    print(f"Scores: {scores}")

tensorflow.keras.layers.GlobalAveragePooling2D tensorflow.keras.layers.Activation tensorflow.keras.models.Model tensorflow.keras.layers.Dense tensorflow.keras.callbacks.LearningRateScheduler tensorflow.keras.optimizers.Adam tensorflow.keras.layers.experimental.preprocessing.Rescaling tensorflow.keras.layers.experimental.preprocessing.Resizing tensorflow.keras.callbacks.EarlyStopping tensorflow.keras.applications.MobileNetV2 tensorflow.keras.layers.Input
[(23, 'tensorflow.keras.applications.MobileNetV2', 'MobileNetV2', ([], {'include_top': '(False)', 'weights': '"""imagenet"""', 'input_shape': 'IMAGENET_SHAPE'}), False, 'from tensorflow.keras.applications import MobileNetV2\n'), (35, 'tensorflow.keras.layers.Input', 'Input', ([], {'shape': 'img_shape'}), False, 'from tensorflow.keras.layers import Input\n'), (43, 'tensorflow.keras.models.Model', 'Model', ([], {'inputs': '[input_img]', 'outputs': '[y_pred]'}), False, 'from tensorflow.keras.models import Model\n'), (58, 'tf_utils.dogsCatsDataAdvanced.DOGSCATS', 'DOGSCATS', ([], {}), False, 'from tf_utils.dogsCatsDataAdvanced import DOGSCATS\n'), (75, 'tensorflow.keras.optimizers.Adam', 'Adam', ([], {'learning_rate': '(0.0005)'}), False, 'from tensorflow.keras.optimizers import Adam\n'), (83, 'tensorflow.keras.callbacks.LearningRateScheduler', 'LearningRateScheduler', ([], {'schedule': 'schedule_fn2', 'verbose': '(1)'}), False, 'from tensorflow.keras.callbacks import LearningRateScheduler\n'), (88, 'tensorflow.keras.callbacks.EarlyStopping', 'EarlyStopping', ([], {'monitor': '"""val_loss"""', 'patience': '(30)', 'verbose': '(1)', 'restore_best_weights': '(True)'}), False, 'from tensorflow.keras.callbacks import EarlyStopping\n'), (36, 'tensorflow.keras.layers.experimental.preprocessing.Rescaling', 'Rescaling', ([], {'scale': '(2.0)', 'offset': '(-1.0)'}), False, 'from tensorflow.keras.layers.experimental.preprocessing import Rescaling\n'), (37, 'tensorflow.keras.layers.experimental.preprocessing.Resizing', 'Resizing', ([], {'height': 'IMAGENET_SIZE', 'width': 'IMAGENET_SIZE'}), False, 'from tensorflow.keras.layers.experimental.preprocessing import Resizing\n'), (39, 'tensorflow.keras.layers.GlobalAveragePooling2D', 'GlobalAveragePooling2D', ([], {}), False, 'from tensorflow.keras.layers import GlobalAveragePooling2D\n'), (40, 'tensorflow.keras.layers.Dense', 'Dense', ([], {'units': 'num_classes'}), False, 'from tensorflow.keras.layers import Dense\n'), (41, 'tensorflow.keras.layers.Activation', 'Activation', (['"""softmax"""'], {}), False, 'from tensorflow.keras.layers import Activation\n')]