# 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
#
#     https://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 ResNet."""

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


@keras_tuner_export("keras_tuner.applications.HyperResNet")
class HyperResNet(hypermodel.HyperModel):
    """A ResNet hypermodel.

    Models built by `HyperResNet` 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:
        include_top: Boolean, whether to include the fully-connected layer at
            the top of the network.
        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.
        **kwargs: Additional keyword arguments that apply to all hypermodels.
            See `keras_tuner.HyperModel`.
    """

    def __init__(
        self,
        include_top=True,
        input_shape=None,
        input_tensor=None,
        classes=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        if include_top and classes is None:
            raise ValueError(
                "You must specify `classes` when `include_top=True`"
            )

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

        self.include_top = include_top
        self.input_shape = input_shape
        self.input_tensor = input_tensor
        self.classes = classes

    def build(self, hp):
        version = hp.Choice("version", ["v1", "v2", "next"], default="v2")
        conv3_depth = hp.Choice("conv3_depth", [4, 8])
        conv4_depth = hp.Choice("conv4_depth", [6, 23, 36])

        # Version-conditional fixed parameters
        preact = version == "v2"
        use_bias = version != "next"

        # Model definition.
        bn_axis = (
            3 if keras.backend.image_data_format() == "channels_last" else 1
        )

        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

        # Initial conv2d block.
        x = layers.ZeroPadding2D(padding=((3, 3), (3, 3)), name="conv1_pad")(x)
        x = layers.Conv2D(
            64, 7, strides=2, use_bias=use_bias, name="conv1_conv"
        )(x)
        if not preact:
            x = layers.BatchNormalization(
                axis=bn_axis, epsilon=1.001e-5, name="conv1_bn"
            )(x)
            x = layers.Activation("relu", name="conv1_relu")(x)
        x = layers.ZeroPadding2D(padding=((1, 1), (1, 1)), name="pool1_pad")(x)
        x = layers.MaxPooling2D(3, strides=2, name="pool1_pool")(x)

        # Middle hypertunable stack.
        if version == "v1":
            x = stack1(x, 64, 3, stride1=1, name="conv2")
            x = stack1(x, 128, conv3_depth, name="conv3")
            x = stack1(x, 256, conv4_depth, name="conv4")
            x = stack1(x, 512, 3, name="conv5")
        elif version == "v2":
            x = stack2(x, 64, 3, name="conv2")
            x = stack2(x, 128, conv3_depth, name="conv3")
            x = stack2(x, 256, conv4_depth, name="conv4")
            x = stack2(x, 512, 3, stride1=1, name="conv5")
        elif version == "next":
            x = stack3(x, 64, 3, name="conv2")
            x = stack3(x, 256, conv3_depth, name="conv3")
            x = stack3(x, 512, conv4_depth, name="conv4")
            x = stack3(x, 1024, 3, stride1=1, name="conv5")

        # Top of the model.
        if preact:
            x = layers.BatchNormalization(
                axis=bn_axis, epsilon=1.001e-5, name="post_bn"
            )(x)
            x = layers.Activation("relu", name="post_relu")(x)

        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)

        if not self.include_top:
            return keras.Model(inputs, x, name="ResNet")
        x = layers.Dense(self.classes, activation="softmax", name="probs")(x)
        model = keras.Model(inputs, x, name="ResNet")
        optimizer_name = hp.Choice(
            "optimizer", ["adam", "rmsprop", "sgd"], default="adam"
        )
        optimizer = keras.optimizers.get(optimizer_name)
        optimizer.learning_rate = hp.Choice(
            "learning_rate", [0.1, 0.01, 0.001], default=0.01
        )
        model.compile(
            optimizer=optimizer,
            loss="categorical_crossentropy",
            metrics=["accuracy"],
        )
        return model


def block1(x, filters, kernel_size=3, stride=1, conv_shortcut=True, name=None):
    """A residual block.

    Args:
        x: input tensor.
        filters: integer, filters of the bottleneck layer.
        kernel_size: default 3, kernel size of the bottleneck layer.
        stride: default 1, stride of the first layer.
        conv_shortcut: default True, use convolution shortcut if True,
            otherwise identity shortcut.
        name: string, block label.

    Returns:
        Output tensor for the residual block.
    """
    bn_axis = 3 if keras.backend.image_data_format() == "channels_last" else 1

    if conv_shortcut is True:
        shortcut = layers.Conv2D(
            4 * filters, 1, strides=stride, name=f"{name}_0_conv"
        )(x)

        shortcut = layers.BatchNormalization(
            axis=bn_axis, epsilon=1.001e-5, name=f"{name}_0_bn"
        )(shortcut)

    else:
        shortcut = x

    x = layers.Conv2D(filters, 1, strides=stride, name=f"{name}_1_conv")(x)
    x = layers.BatchNormalization(
        axis=bn_axis, epsilon=1.001e-5, name=f"{name}_1_bn"
    )(x)

    x = layers.Activation("relu", name=f"{name}_1_relu")(x)

    x = layers.Conv2D(
        filters, kernel_size, padding="same", name=f"{name}_2_conv"
    )(x)

    x = layers.BatchNormalization(
        axis=bn_axis, epsilon=1.001e-5, name=f"{name}_2_bn"
    )(x)

    x = layers.Activation("relu", name=f"{name}_2_relu")(x)

    x = layers.Conv2D(4 * filters, 1, name=f"{name}_3_conv")(x)
    x = layers.BatchNormalization(
        axis=bn_axis, epsilon=1.001e-5, name=f"{name}_3_bn"
    )(x)

    x = layers.Add(name=f"{name}_add")([shortcut, x])
    x = layers.Activation("relu", name=f"{name}_out")(x)
    return x


def stack1(x, filters, blocks, stride1=2, name=None):
    """A set of stacked residual blocks.

    Args:
        x: input tensor.
        filters: integer, filters of the bottleneck layer in a block.
        blocks: integer, blocks in the stacked blocks.
        stride1: default 2, stride of the first layer in the first block.
        name: string, stack label.

    Returns:
        Output tensor for the stacked blocks.
    """
    x = block1(x, filters, stride=stride1, name=f"{name}_block1")
    for i in range(2, blocks + 1):
        x = block1(
            x, filters, conv_shortcut=False, name=f"{name}_block{str(i)}"
        )
    return x


def block2(x, filters, kernel_size=3, stride=1, conv_shortcut=False, name=None):
    """A residual block.

    Args:
        x: input tensor.
        filters: integer, filters of the bottleneck layer.
        kernel_size: default 3, kernel size of the bottleneck layer.
        stride: default 1, stride of the first layer.
        conv_shortcut: default False, use convolution shortcut if True,
            otherwise identity shortcut.
        name: string, block label.

    Returns:
        Output tensor for the residual block.
    """
    bn_axis = 3 if keras.backend.image_data_format() == "channels_last" else 1

    preact = layers.BatchNormalization(
        axis=bn_axis, epsilon=1.001e-5, name=f"{name}_preact_bn"
    )(x)

    preact = layers.Activation("relu", name=f"{name}_preact_relu")(preact)

    if conv_shortcut is True:
        shortcut = layers.Conv2D(
            4 * filters, 1, strides=stride, name=f"{name}_0_conv"
        )(preact)

    else:
        shortcut = (
            layers.MaxPooling2D(1, strides=stride)(x) if stride > 1 else x
        )

    x = layers.Conv2D(
        filters, 1, strides=1, use_bias=False, name=f"{name}_1_conv"
    )(preact)

    x = layers.BatchNormalization(
        axis=bn_axis, epsilon=1.001e-5, name=f"{name}_1_bn"
    )(x)

    x = layers.Activation("relu", name=f"{name}_1_relu")(x)

    x = layers.ZeroPadding2D(padding=((1, 1), (1, 1)), name=f"{name}_2_pad")(x)
    x = layers.Conv2D(
        filters,
        kernel_size,
        strides=stride,
        use_bias=False,
        name=f"{name}_2_conv",
    )(x)

    x = layers.BatchNormalization(
        axis=bn_axis, epsilon=1.001e-5, name=f"{name}_2_bn"
    )(x)

    x = layers.Activation("relu", name=f"{name}_2_relu")(x)

    x = layers.Conv2D(4 * filters, 1, name=f"{name}_3_conv")(x)
    x = layers.Add(name=f"{name}_out")([shortcut, x])
    return x


def stack2(x, filters, blocks, stride1=2, name=None):
    """A set of stacked residual blocks.

    Args:
        x: input tensor.
        filters: integer, filters of the bottleneck layer in a block.
        blocks: integer, blocks in the stacked blocks.
        stride1: default 2, stride of the first layer in the first block.
        name: string, stack label.

    Returns:
        Output tensor for the stacked blocks.
    """
    x = block2(x, filters, conv_shortcut=True, name=f"{name}_block1")
    for i in range(2, blocks):
        x = block2(x, filters, name=f"{name}_block{str(i)}")
    x = block2(x, filters, stride=stride1, name=f"{name}_block{str(blocks)}")
    return x


def block3(
    x,
    filters,
    kernel_size=3,
    stride=1,
    groups=32,
    conv_shortcut=True,
    name=None,
):
    """A residual block.

    Args:
        x: input tensor.
        filters: integer, filters of the bottleneck layer.
        kernel_size: default 3, kernel size of the bottleneck layer.
        stride: default 1, stride of the first layer.
        groups: default 32, group size for grouped convolution.
        conv_shortcut: default True, use convolution shortcut if True,
            otherwise identity shortcut.
        name: string, block label.

    Returns:
        Output tensor for the residual block.
    """
    bn_axis = 3 if keras.backend.image_data_format() == "channels_last" else 1

    if conv_shortcut is True:
        shortcut = layers.Conv2D(
            (64 // groups) * filters,
            1,
            strides=stride,
            use_bias=False,
            name=f"{name}_0_conv",
        )(x)

        shortcut = layers.BatchNormalization(
            axis=bn_axis, epsilon=1.001e-5, name=f"{name}_0_bn"
        )(shortcut)

    else:
        shortcut = x

    x = layers.Conv2D(filters, 1, use_bias=False, name=f"{name}_1_conv")(x)
    x = layers.BatchNormalization(
        axis=bn_axis, epsilon=1.001e-5, name=f"{name}_1_bn"
    )(x)

    x = layers.Activation("relu", name=f"{name}_1_relu")(x)

    c = filters // groups
    x = layers.ZeroPadding2D(padding=((1, 1), (1, 1)), name=f"{name}_2_pad")(x)
    x = layers.DepthwiseConv2D(
        kernel_size,
        strides=stride,
        depth_multiplier=c,
        use_bias=False,
        name=f"{name}_2_conv",
    )(x)

    if bn_axis == 3:
        x_shape = ops.shape(x)[1:-1]
        x = layers.Reshape(x_shape + (groups, c, c))(x)
        output_shape = x_shape + (groups, c)
        x = layers.Lambda(
            lambda x: sum(x[:, :, :, :, i] for i in range(c)),
            name=f"{name}_2_reduce",
            output_shape=output_shape,
        )(x)

        x = layers.Reshape(x_shape + (filters,))(x)
    else:
        x_shape = ops.shape(x)[2:]
        x = layers.Reshape((groups, c, c) + x_shape)(x)
        output_shape = (groups, c) + x_shape
        x = layers.Lambda(
            lambda x: sum(x[:, :, i, :, :] for i in range(c)),
            name=f"{name}_2_reduce",
            output_shape=output_shape,
        )(x)

        x = layers.Reshape((filters,) + x_shape)(x)

    x = layers.BatchNormalization(
        axis=bn_axis, epsilon=1.001e-5, name=f"{name}_2_bn"
    )(x)

    x = layers.Activation("relu", name=f"{name}_2_relu")(x)

    x = layers.Conv2D(
        (64 // groups) * filters, 1, use_bias=False, name=f"{name}_3_conv"
    )(x)

    x = layers.BatchNormalization(
        axis=bn_axis, epsilon=1.001e-5, name=f"{name}_3_bn"
    )(x)

    x = layers.Add(name=f"{name}_add")([shortcut, x])
    x = layers.Activation("relu", name=f"{name}_out")(x)
    return x


def stack3(x, filters, blocks, stride1=2, groups=32, name=None):
    """A set of stacked residual blocks.

    Args:
        x: input tensor.
        filters: integer, filters of the bottleneck layer in a block.
        blocks: integer, blocks in the stacked blocks.
        stride1: default 2, stride of the first layer in the first block.
        groups: default 32, group size for grouped convolution.
        name: string, stack label.

    Returns:
        Output tensor for the stacked blocks.
    """
    x = block3(x, filters, stride=stride1, groups=groups, name=f"{name}_block1")

    for i in range(2, blocks + 1):
        x = block3(
            x,
            filters,
            groups=groups,
            conv_shortcut=False,
            name=f"{name}_block{str(i)}",
        )

    return x

