# Copyright 2019 The KerasTuner Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Hypertunable version of EfficientNet based on Keras.applications."""

from keras_tuner.src.api_export import keras_tuner_export
from keras_tuner.src.backend import keras
from keras_tuner.src.backend.keras import layers
from keras_tuner.src.engine import hypermodel

efficientnet = keras.applications.efficientnet
EFFICIENTNET_MODELS = {
    "B0": efficientnet.EfficientNetB0,
    "B1": efficientnet.EfficientNetB1,
    "B2": efficientnet.EfficientNetB2,
    "B3": efficientnet.EfficientNetB3,
    "B4": efficientnet.EfficientNetB4,
    "B5": efficientnet.EfficientNetB5,
    "B6": efficientnet.EfficientNetB6,
    "B7": efficientnet.EfficientNetB7,
}

EFFICIENTNET_IMG_SIZE = {
    "B0": 224,
    "B1": 240,
    "B2": 260,
    "B3": 300,
    "B4": 380,
    "B5": 456,
    "B6": 528,
    "B7": 600,
}


@keras_tuner_export("keras_tuner.applications.HyperEfficientNet")
class HyperEfficientNet(hypermodel.HyperModel):
    """An EfficientNet hypermodel.

    Models built by `HyperEfficientNet` take images with shape (height, width,
    channels) as input. The output are one-hot encoded with the length matching
    the number of classes specified by the `classes` argument.

    Args:
        input_shape: Optional shape tuple, e.g. `(256, 256, 3)`.  One of
            `input_shape` or `input_tensor` must be specified.
        input_tensor: Optional Keras tensor (i.e. output of `layers.Input()`)
            to use as image input for the model.  One of `input_shape` or
            `input_tensor` must be specified.
        classes: Optional number of classes to classify images into, only to be
            specified if `include_top` is True, and if no `weights` argument is
            specified.
        augmentation_model: Optional `Model` or `HyperModel` instance for image
            augmentation.
        **kwargs: Additional keyword arguments that apply to all hypermodels.
            See `keras_tuner.HyperModel`.
    """

    def __init__(
        self,
        input_shape=None,
        input_tensor=None,
        classes=None,
        augmentation_model=None,
        **kwargs,
    ):
        if not isinstance(
            augmentation_model, (hypermodel.HyperModel, keras.Model, type(None))
        ):
            raise ValueError(
                "Keyword augmentation_model should be "
                "a `HyperModel`, a Keras `Model` or "
                f"empty. Received {augmentation_model}."
            )

        if not classes:
            raise ValueError("You must specify `classes`.")

        if input_shape is None and input_tensor is None:
            raise ValueError(
                "You must specify either `input_shape` or `input_tensor`."
            )

        self.input_shape = input_shape
        self.input_tensor = input_tensor
        self.classes = classes
        self.augmentation_model = augmentation_model

        super().__init__(**kwargs)

    def build(self, hp):
        if self.input_tensor is not None:
            inputs = keras.utils.get_source_inputs(self.input_tensor)
            x = self.input_tensor
        else:
            inputs = layers.Input(shape=self.input_shape)
            x = inputs

        if self.augmentation_model:
            if isinstance(self.augmentation_model, hypermodel.HyperModel):
                augmentation_model = self.augmentation_model.build(hp)
            elif isinstance(self.augmentation_model, keras.models.Model):
                augmentation_model = self.augmentation_model

            x = augmentation_model(x)

        # Select one of pre-trained EfficientNet as feature extractor
        version = hp.Choice(
            "version", [f"B{i}" for i in range(8)], default="B0"
        )
        img_size = EFFICIENTNET_IMG_SIZE[version]

        x = layers.Resizing(img_size, img_size, interpolation="bilinear")(x)
        efficientnet_model = EFFICIENTNET_MODELS[version](
            include_top=False, input_tensor=x
        )

        # Rebuild top layers of the model.
        x = efficientnet_model.output

        pooling = hp.Choice("pooling", ["avg", "max"], default="avg")
        if pooling == "avg":
            x = layers.GlobalAveragePooling2D(name="avg_pool")(x)
        elif pooling == "max":
            x = layers.GlobalMaxPooling2D(name="max_pool")(x)

        top_dropout_rate = hp.Float(
            "top_dropout_rate",
            min_value=0.2,
            max_value=0.8,
            step=0.2,
            default=0.2,
        )
        x = layers.Dropout(top_dropout_rate, name="top_dropout")(x)

        x = layers.Dense(self.classes, activation="softmax", name="probs")(x)

        # compile
        model = keras.Model(inputs, x, name="EfficientNet")
        self._compile(model, hp)

        return model

    def _compile(self, model, hp):
        """Compile model using hyperparameters in hp.

        When subclassing the hypermodel, this may be overridden to change
        behavior of compiling.
        """
        learning_rate = hp.Choice(
            "learning_rate", [0.1, 0.01, 0.001], default=0.01
        )
        optimizer = keras.optimizers.SGD(
            momentum=0.1, learning_rate=learning_rate
        )

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

