# Copyright (C) 2021-2024, Mindee.

# This program is licensed under the Apache License 2.0.
# See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> for full license details.

# Greatly inspired by https://github.com/pytorch/vision/blob/master/torchvision/models/mobilenetv3.py

from copy import deepcopy
from typing import Any

import tensorflow as tf
from tensorflow.keras import layers
from tensorflow.keras.models import Sequential

from ....datasets import VOCABS
from ...utils import _build_model, conv_sequence, load_pretrained_params

__all__ = [
    "MobileNetV3",
    "mobilenet_v3_small",
    "mobilenet_v3_small_r",
    "mobilenet_v3_large",
    "mobilenet_v3_large_r",
    "mobilenet_v3_small_crop_orientation",
    "mobilenet_v3_small_page_orientation",
]


default_cfgs: dict[str, dict[str, Any]] = {
    "mobilenet_v3_large": {
        "mean": (0.694, 0.695, 0.693),
        "std": (0.299, 0.296, 0.301),
        "input_shape": (32, 32, 3),
        "classes": list(VOCABS["french"]),
        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/mobilenet_v3_large-d857506e.weights.h5&src=0",
    },
    "mobilenet_v3_large_r": {
        "mean": (0.694, 0.695, 0.693),
        "std": (0.299, 0.296, 0.301),
        "input_shape": (32, 32, 3),
        "classes": list(VOCABS["french"]),
        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/mobilenet_v3_large_r-eef2e3c6.weights.h5&src=0",
    },
    "mobilenet_v3_small": {
        "mean": (0.694, 0.695, 0.693),
        "std": (0.299, 0.296, 0.301),
        "input_shape": (32, 32, 3),
        "classes": list(VOCABS["french"]),
        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/mobilenet_v3_small-3fcebad7.weights.h5&src=0",
    },
    "mobilenet_v3_small_r": {
        "mean": (0.694, 0.695, 0.693),
        "std": (0.299, 0.296, 0.301),
        "input_shape": (32, 32, 3),
        "classes": list(VOCABS["french"]),
        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/mobilenet_v3_small_r-dd50218d.weights.h5&src=0",
    },
    "mobilenet_v3_small_crop_orientation": {
        "mean": (0.694, 0.695, 0.693),
        "std": (0.299, 0.296, 0.301),
        "input_shape": (128, 128, 3),
        "classes": [0, -90, 180, 90],
        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/mobilenet_v3_small_crop_orientation-ef019b6b.weights.h5&src=0",
    },
    "mobilenet_v3_small_page_orientation": {
        "mean": (0.694, 0.695, 0.693),
        "std": (0.299, 0.296, 0.301),
        "input_shape": (512, 512, 3),
        "classes": [0, -90, 180, 90],
        "url": "https://doctr-static.mindee.com/models?id=v0.9.0/mobilenet_v3_small_page_orientation-0071d55d.weights.h5&src=0",
    },
}


def hard_swish(x: tf.Tensor) -> tf.Tensor:
    return x * tf.nn.relu6(x + 3.0) / 6.0


def _make_divisible(v: float, divisor: int, min_value: int | None = None) -> int:
    if min_value is None:
        min_value = divisor
    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
    # Make sure that round down does not go down by more than 10%.
    if new_v < 0.9 * v:
        new_v += divisor
    return new_v


class SqueezeExcitation(Sequential):
    """Squeeze and Excitation."""

    def __init__(self, chan: int, squeeze_factor: int = 4) -> None:
        super().__init__([
            layers.GlobalAveragePooling2D(),
            layers.Dense(chan // squeeze_factor, activation="relu"),
            layers.Dense(chan, activation="hard_sigmoid"),
            layers.Reshape((1, 1, chan)),
        ])

    def call(self, inputs: tf.Tensor, **kwargs: Any) -> tf.Tensor:
        x = super().call(inputs, **kwargs)
        x = tf.math.multiply(inputs, x)
        return x


class InvertedResidualConfig:
    def __init__(
        self,
        input_channels: int,
        kernel: int,
        expanded_channels: int,
        out_channels: int,
        use_se: bool,
        activation: str,
        stride: int | tuple[int, int],
        width_mult: float = 1,
    ) -> None:
        self.input_channels = self.adjust_channels(input_channels, width_mult)
        self.kernel = kernel
        self.expanded_channels = self.adjust_channels(expanded_channels, width_mult)
        self.out_channels = self.adjust_channels(out_channels, width_mult)
        self.use_se = use_se
        self.use_hs = activation == "HS"
        self.stride = stride

    @staticmethod
    def adjust_channels(channels: int, width_mult: float):
        return _make_divisible(channels * width_mult, 8)


class InvertedResidual(layers.Layer):
    """InvertedResidual for mobilenet

    Args:
        conf: configuration object for inverted residual
    """

    def __init__(
        self,
        conf: InvertedResidualConfig,
        **kwargs: Any,
    ) -> None:
        _kwargs = {"input_shape": kwargs.pop("input_shape")} if isinstance(kwargs.get("input_shape"), tuple) else {}
        super().__init__(**kwargs)

        act_fn = hard_swish if conf.use_hs else tf.nn.relu

        _is_s1 = (isinstance(conf.stride, tuple) and conf.stride == (1, 1)) or conf.stride == 1
        self.use_res_connect = _is_s1 and conf.input_channels == conf.out_channels

        _layers = []
        # expand
        if conf.expanded_channels != conf.input_channels:
            _layers.extend(conv_sequence(conf.expanded_channels, act_fn, kernel_size=1, bn=True, **_kwargs))

        # depth-wise
        _layers.extend(
            conv_sequence(
                conf.expanded_channels,
                act_fn,
                kernel_size=conf.kernel,
                strides=conf.stride,
                bn=True,
                groups=conf.expanded_channels,
            )
        )

        if conf.use_se:
            _layers.append(SqueezeExcitation(conf.expanded_channels))

        # project
        _layers.extend(
            conv_sequence(
                conf.out_channels,
                None,
                kernel_size=1,
                bn=True,
            )
        )

        self.block = Sequential(_layers)

    def call(
        self,
        inputs: tf.Tensor,
        **kwargs: Any,
    ) -> tf.Tensor:
        out = self.block(inputs, **kwargs)
        if self.use_res_connect:
            out = tf.add(out, inputs)

        return out


class MobileNetV3(Sequential):
    """Implements MobileNetV3, inspired from both:
    <https://github.com/xiaochus/MobileNetV3/tree/master/model>`_.
    and <https://pytorch.org/vision/stable/_modules/torchvision/models/mobilenetv3.html>`_.
    """

    def __init__(
        self,
        layout: list[InvertedResidualConfig],
        include_top: bool = True,
        head_chans: int = 1024,
        num_classes: int = 1000,
        cfg: dict[str, Any] | None = None,
        input_shape: tuple[int, int, int] | None = None,
    ) -> None:
        _layers = [
            Sequential(
                conv_sequence(
                    layout[0].input_channels, hard_swish, True, kernel_size=3, strides=2, input_shape=input_shape
                ),
                name="stem",
            )
        ]

        for idx, conf in enumerate(layout):
            _layers.append(
                InvertedResidual(conf, name=f"inverted_{idx}"),
            )

        _layers.append(
            Sequential(conv_sequence(6 * layout[-1].out_channels, hard_swish, True, kernel_size=1), name="final_block")
        )

        if include_top:
            _layers.extend([
                layers.GlobalAveragePooling2D(),
                layers.Dense(head_chans, activation=hard_swish),
                layers.Dropout(0.2),
                layers.Dense(num_classes),
            ])

        super().__init__(_layers)
        self.cfg = cfg


def _mobilenet_v3(arch: str, pretrained: bool, rect_strides: bool = False, **kwargs: Any) -> MobileNetV3:
    kwargs["num_classes"] = kwargs.get("num_classes", len(default_cfgs[arch]["classes"]))
    kwargs["input_shape"] = kwargs.get("input_shape", default_cfgs[arch]["input_shape"])
    kwargs["classes"] = kwargs.get("classes", default_cfgs[arch]["classes"])

    _cfg = deepcopy(default_cfgs[arch])
    _cfg["num_classes"] = kwargs["num_classes"]
    _cfg["classes"] = kwargs["classes"]
    _cfg["input_shape"] = kwargs["input_shape"]
    kwargs.pop("classes")

    # cf. Table 1 & 2 of the paper
    if arch.startswith("mobilenet_v3_small"):
        inverted_residual_setting = [
            InvertedResidualConfig(16, 3, 16, 16, True, "RE", 2),  # C1
            InvertedResidualConfig(16, 3, 72, 24, False, "RE", (2, 1) if rect_strides else 2),  # C2
            InvertedResidualConfig(24, 3, 88, 24, False, "RE", 1),
            InvertedResidualConfig(24, 5, 96, 40, True, "HS", (2, 1) if rect_strides else 2),  # C3
            InvertedResidualConfig(40, 5, 240, 40, True, "HS", 1),
            InvertedResidualConfig(40, 5, 240, 40, True, "HS", 1),
            InvertedResidualConfig(40, 5, 120, 48, True, "HS", 1),
            InvertedResidualConfig(48, 5, 144, 48, True, "HS", 1),
            InvertedResidualConfig(48, 5, 288, 96, True, "HS", (2, 1) if rect_strides else 2),  # C4
            InvertedResidualConfig(96, 5, 576, 96, True, "HS", 1),
            InvertedResidualConfig(96, 5, 576, 96, True, "HS", 1),
        ]
        head_chans = 1024
    else:
        inverted_residual_setting = [
            InvertedResidualConfig(16, 3, 16, 16, False, "RE", 1),
            InvertedResidualConfig(16, 3, 64, 24, False, "RE", 2),  # C1
            InvertedResidualConfig(24, 3, 72, 24, False, "RE", 1),
            InvertedResidualConfig(24, 5, 72, 40, True, "RE", (2, 1) if rect_strides else 2),  # C2
            InvertedResidualConfig(40, 5, 120, 40, True, "RE", 1),
            InvertedResidualConfig(40, 5, 120, 40, True, "RE", 1),
            InvertedResidualConfig(40, 3, 240, 80, False, "HS", (2, 1) if rect_strides else 2),  # C3
            InvertedResidualConfig(80, 3, 200, 80, False, "HS", 1),
            InvertedResidualConfig(80, 3, 184, 80, False, "HS", 1),
            InvertedResidualConfig(80, 3, 184, 80, False, "HS", 1),
            InvertedResidualConfig(80, 3, 480, 112, True, "HS", 1),
            InvertedResidualConfig(112, 3, 672, 112, True, "HS", 1),
            InvertedResidualConfig(112, 5, 672, 160, True, "HS", (2, 1) if rect_strides else 2),  # C4
            InvertedResidualConfig(160, 5, 960, 160, True, "HS", 1),
            InvertedResidualConfig(160, 5, 960, 160, True, "HS", 1),
        ]
        head_chans = 1280

    kwargs["num_classes"] = _cfg["num_classes"]
    kwargs["input_shape"] = _cfg["input_shape"]

    # Build the model
    model = MobileNetV3(
        inverted_residual_setting,
        head_chans=head_chans,
        cfg=_cfg,
        **kwargs,
    )
    _build_model(model)

    # Load pretrained parameters
    if pretrained:
        # The number of classes is not the same as the number of classes in the pretrained model =>
        # skip the mismatching layers for fine tuning
        load_pretrained_params(
            model, default_cfgs[arch]["url"], skip_mismatch=kwargs["num_classes"] != len(default_cfgs[arch]["classes"])
        )

    return model


def mobilenet_v3_small(pretrained: bool = False, **kwargs: Any) -> MobileNetV3:
    """MobileNetV3-Small architecture as described in
    `"Searching for MobileNetV3",
    <https://arxiv.org/pdf/1905.02244.pdf>`_.

    >>> import tensorflow as tf
    >>> from doctr.models import mobilenet_v3_small
    >>> model = mobilenet_v3_small(pretrained=False)
    >>> input_tensor = tf.random.uniform(shape=[1, 512, 512, 3], maxval=1, dtype=tf.float32)
    >>> out = model(input_tensor)

    Args:
        pretrained: boolean, True if model is pretrained
        **kwargs: keyword arguments of the MobileNetV3 architecture

    Returns:
        a keras.Model
    """
    return _mobilenet_v3("mobilenet_v3_small", pretrained, False, **kwargs)


def mobilenet_v3_small_r(pretrained: bool = False, **kwargs: Any) -> MobileNetV3:
    """MobileNetV3-Small architecture as described in
    `"Searching for MobileNetV3",
    <https://arxiv.org/pdf/1905.02244.pdf>`_, with rectangular pooling.

    >>> import tensorflow as tf
    >>> from doctr.models import mobilenet_v3_small_r
    >>> model = mobilenet_v3_small_r(pretrained=False)
    >>> input_tensor = tf.random.uniform(shape=[1, 512, 512, 3], maxval=1, dtype=tf.float32)
    >>> out = model(input_tensor)

    Args:
        pretrained: boolean, True if model is pretrained
        **kwargs: keyword arguments of the MobileNetV3 architecture

    Returns:
        a keras.Model
    """
    return _mobilenet_v3("mobilenet_v3_small_r", pretrained, True, **kwargs)


def mobilenet_v3_large(pretrained: bool = False, **kwargs: Any) -> MobileNetV3:
    """MobileNetV3-Large architecture as described in
    `"Searching for MobileNetV3",
    <https://arxiv.org/pdf/1905.02244.pdf>`_.

    >>> import tensorflow as tf
    >>> from doctr.models import mobilenet_v3_large
    >>> model = mobilenet_v3_large(pretrained=False)
    >>> input_tensor = tf.random.uniform(shape=[1, 512, 512, 3], maxval=1, dtype=tf.float32)
    >>> out = model(input_tensor)

    Args:
        pretrained: boolean, True if model is pretrained
        **kwargs: keyword arguments of the MobileNetV3 architecture

    Returns:
        a keras.Model
    """
    return _mobilenet_v3("mobilenet_v3_large", pretrained, False, **kwargs)


def mobilenet_v3_large_r(pretrained: bool = False, **kwargs: Any) -> MobileNetV3:
    """MobileNetV3-Large architecture as described in
    `"Searching for MobileNetV3",
    <https://arxiv.org/pdf/1905.02244.pdf>`_.

    >>> import tensorflow as tf
    >>> from doctr.models import mobilenet_v3_large_r
    >>> model = mobilenet_v3_large_r(pretrained=False)
    >>> input_tensor = tf.random.uniform(shape=[1, 512, 512, 3], maxval=1, dtype=tf.float32)
    >>> out = model(input_tensor)

    Args:
        pretrained: boolean, True if model is pretrained
        **kwargs: keyword arguments of the MobileNetV3 architecture

    Returns:
        a keras.Model
    """
    return _mobilenet_v3("mobilenet_v3_large_r", pretrained, True, **kwargs)


def mobilenet_v3_small_crop_orientation(pretrained: bool = False, **kwargs: Any) -> MobileNetV3:
    """MobileNetV3-Small architecture as described in
    `"Searching for MobileNetV3",
    <https://arxiv.org/pdf/1905.02244.pdf>`_.

    >>> import tensorflow as tf
    >>> from doctr.models import mobilenet_v3_small_crop_orientation
    >>> model = mobilenet_v3_small_crop_orientation(pretrained=False)
    >>> input_tensor = tf.random.uniform(shape=[1, 512, 512, 3], maxval=1, dtype=tf.float32)
    >>> out = model(input_tensor)

    Args:
        pretrained: boolean, True if model is pretrained
        **kwargs: keyword arguments of the MobileNetV3 architecture

    Returns:
        a keras.Model
    """
    return _mobilenet_v3("mobilenet_v3_small_crop_orientation", pretrained, include_top=True, **kwargs)


def mobilenet_v3_small_page_orientation(pretrained: bool = False, **kwargs: Any) -> MobileNetV3:
    """MobileNetV3-Small architecture as described in
    `"Searching for MobileNetV3",
    <https://arxiv.org/pdf/1905.02244.pdf>`_.

    >>> import tensorflow as tf
    >>> from doctr.models import mobilenet_v3_small_page_orientation
    >>> model = mobilenet_v3_small_page_orientation(pretrained=False)
    >>> input_tensor = tf.random.uniform(shape=[1, 512, 512, 3], maxval=1, dtype=tf.float32)
    >>> out = model(input_tensor)

    Args:
        pretrained: boolean, True if model is pretrained
        **kwargs: keyword arguments of the MobileNetV3 architecture

    Returns:
        a keras.Model
    """
    return _mobilenet_v3("mobilenet_v3_small_page_orientation", pretrained, include_top=True, **kwargs)
