# copyright (c) 2023 PaddlePaddle Authors. All Rights Reserve.
#
# 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.

# Code was based on https://github.com/micronDLA/MobileViTv3/blob/main/MobileViTv3-v1/cvnets/models/classification/mobilevit.py
# reference: https://arxiv.org/abs/2209.15159

import math
from functools import partial
from typing import Dict, Optional, Tuple, Union

import paddle
import paddle.nn as nn
import paddle.nn.functional as F

from ....utils.save_load import load_dygraph_pretrain, load_dygraph_pretrain_from_url

MODEL_URLS = {
    "MobileViTV3_XXS":
    "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileViTV3_XXS_pretrained.pdparams",
    "MobileViTV3_XS":
    "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileViTV3_XS_pretrained.pdparams",
    "MobileViTV3_S":
    "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileViTV3_S_pretrained.pdparams",
    "MobileViTV3_XXS_L2":
    "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileViTV3_XXS_L2_pretrained.pdparams",
    "MobileViTV3_XS_L2":
    "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileViTV3_XS_L2_pretrained.pdparams",
    "MobileViTV3_S_L2":
    "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileViTV3_S_L2_pretrained.pdparams",
    "MobileViTV3_x0_5":
    "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileViTV3_x0_5_pretrained.pdparams",
    "MobileViTV3_x0_75":
    "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileViTV3_x0_75_pretrained.pdparams",
    "MobileViTV3_x1_0":
    "https://paddle-imagenet-models-name.bj.bcebos.com/dygraph/MobileViTV3_x1_0_pretrained.pdparams",
}

layer_norm_2d = partial(nn.GroupNorm, num_groups=1)


def make_divisible(v, divisor=8, min_value=None):
    if min_value is None:
        min_value = divisor
    new_v = max(min_value, int(v + divisor / 2) // divisor * divisor)
    if new_v < 0.9 * v:
        new_v += divisor
    return new_v


class InvertedResidual(nn.Layer):
    """
    Inverted residual block (MobileNetv2): https://arxiv.org/abs/1801.04381
    """

    def __init__(self,
                 in_channels: int,
                 out_channels: int,
                 stride: int,
                 expand_ratio: Union[int, float],
                 dilation: int=1) -> None:
        assert stride in [1, 2]
        super(InvertedResidual, self).__init__()
        self.stride = stride

        hidden_dim = make_divisible(int(round(in_channels * expand_ratio)), 8)
        self.use_res_connect = self.stride == 1 and in_channels == out_channels

        block = nn.Sequential()
        if expand_ratio != 1:
            block.add_sublayer(
                name="exp_1x1",
                sublayer=nn.Sequential(
                    ('conv', nn.Conv2D(
                        in_channels, hidden_dim, 1, bias_attr=False)),
                    ('norm', nn.BatchNorm2D(hidden_dim)), ('act', nn.Silu())))

        block.add_sublayer(
            name="conv_3x3",
            sublayer=nn.Sequential(
                ('conv', nn.Conv2D(
                    hidden_dim,
                    hidden_dim,
                    3,
                    bias_attr=False,
                    stride=stride,
                    padding=dilation,
                    dilation=dilation,
                    groups=hidden_dim)), ('norm', nn.BatchNorm2D(hidden_dim)),
                ('act', nn.Silu())))

        block.add_sublayer(
            name="red_1x1",
            sublayer=nn.Sequential(
                ('conv', nn.Conv2D(
                    hidden_dim, out_channels, 1, bias_attr=False)),
                ('norm', nn.BatchNorm2D(out_channels))))

        self.block = block
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.exp = expand_ratio
        self.dilation = dilation

    def forward(self, x, *args, **kwargs):
        if self.use_res_connect:
            return x + self.block(x)
        else:
            return self.block(x)


class MultiHeadAttention(nn.Layer):
    def __init__(self,
                 dim,
                 num_heads=8,
                 qkv_bias=False,
                 qk_scale=None,
                 attn_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim**-0.5

        self.qkv_proj = nn.Linear(dim, dim * 3, bias_attr=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.out_proj = nn.Linear(dim, dim, bias_attr=qkv_bias)

    def forward(self, x):
        # B = paddle.shape(x)[0]
        N, C = x.shape[1:]
        qkv = self.qkv_proj(x).reshape((-1, N, 3, self.num_heads,
                                        C // self.num_heads)).transpose(
                                            (2, 0, 3, 1, 4))
        q, k, v = qkv[0], qkv[1], qkv[2]

        attn = (q.matmul(k.transpose((0, 1, 3, 2)))) * self.scale
        attn = nn.functional.softmax(attn, axis=-1)
        attn = self.attn_drop(attn)

        x = (attn.matmul(v)).transpose((0, 2, 1, 3)).reshape((-1, N, C))
        x = self.out_proj(x)
        return x


class TransformerEncoder(nn.Layer):
    """
        This class defines the Transformer encoder (pre-norm) as described in "Attention is all you need" paper
            https://arxiv.org/abs/1706.03762
    """

    def __init__(self,
                 embed_dim: int,
                 ffn_latent_dim: int,
                 num_heads: Optional[int]=8,
                 attn_dropout: Optional[float]=0.0,
                 dropout: Optional[float]=0.1,
                 ffn_dropout: Optional[float]=0.0,
                 transformer_norm_layer: nn.Layer=nn.LayerNorm):
        super(TransformerEncoder, self).__init__()

        self.pre_norm_mha = nn.Sequential(
            transformer_norm_layer(embed_dim),
            MultiHeadAttention(
                embed_dim, num_heads, attn_drop=attn_dropout, qkv_bias=True),
            nn.Dropout(p=dropout))

        self.pre_norm_ffn = nn.Sequential(
            transformer_norm_layer(embed_dim),
            nn.Linear(embed_dim, ffn_latent_dim),
            nn.Silu(),
            nn.Dropout(p=ffn_dropout),
            nn.Linear(ffn_latent_dim, embed_dim),
            nn.Dropout(p=dropout))
        self.embed_dim = embed_dim
        self.ffn_dim = ffn_latent_dim
        self.ffn_dropout = ffn_dropout

    def forward(self, x):
        # Multi-head attention
        x = x + self.pre_norm_mha(x)

        # Feed forward network
        x = x + self.pre_norm_ffn(x)
        return x


class MobileViTV3Block(nn.Layer):
    """
        MobileViTV3 block
    """

    def __init__(self,
                 in_channels: int,
                 transformer_dim: int,
                 ffn_dim: int,
                 n_transformer_blocks: Optional[int]=2,
                 head_dim: Optional[int]=32,
                 attn_dropout: Optional[float]=0.1,
                 dropout: Optional[int]=0.1,
                 ffn_dropout: Optional[int]=0.1,
                 patch_h: Optional[int]=8,
                 patch_w: Optional[int]=8,
                 transformer_norm_layer: nn.Layer=nn.LayerNorm,
                 conv_ksize: Optional[int]=3,
                 dilation: Optional[int]=1,
                 var_ffn: Optional[bool]=False,
                 no_fusion: Optional[bool]=False):

        # For MobileViTV3: Normal 3x3 convolution --> Depthwise 3x3 convolution
        padding = (conv_ksize - 1) // 2 * dilation
        conv_3x3_in = nn.Sequential(
            ('conv', nn.Conv2D(
                in_channels,
                in_channels,
                conv_ksize,
                bias_attr=False,
                padding=padding,
                dilation=dilation,
                groups=in_channels)), ('norm', nn.BatchNorm2D(in_channels)),
            ('act', nn.Silu()))
        conv_1x1_in = nn.Sequential(('conv', nn.Conv2D(
            in_channels, transformer_dim, 1, bias_attr=False)))

        conv_1x1_out = nn.Sequential(
            ('conv', nn.Conv2D(
                transformer_dim, in_channels, 1, bias_attr=False)),
            ('norm', nn.BatchNorm2D(in_channels)), ('act', nn.Silu()))
        conv_3x3_out = None

        # For MobileViTV3: input+global --> local+global
        if not no_fusion:
            #input_ch = tr_dim + in_ch
            conv_3x3_out = nn.Sequential(
                ('conv', nn.Conv2D(
                    transformer_dim + in_channels,
                    in_channels,
                    1,
                    bias_attr=False)), ('norm', nn.BatchNorm2D(in_channels)),
                ('act', nn.Silu()))

        super().__init__()
        self.local_rep = nn.Sequential()
        self.local_rep.add_sublayer(name="conv_3x3", sublayer=conv_3x3_in)
        self.local_rep.add_sublayer(name="conv_1x1", sublayer=conv_1x1_in)

        assert transformer_dim % head_dim == 0
        num_heads = transformer_dim // head_dim

        ffn_dims = [ffn_dim] * n_transformer_blocks

        global_rep = [
            TransformerEncoder(
                embed_dim=transformer_dim,
                ffn_latent_dim=ffn_dims[block_idx],
                num_heads=num_heads,
                attn_dropout=attn_dropout,
                dropout=dropout,
                ffn_dropout=ffn_dropout,
                transformer_norm_layer=transformer_norm_layer)
            for block_idx in range(n_transformer_blocks)
        ]
        global_rep.append(transformer_norm_layer(transformer_dim))
        self.global_rep = nn.Sequential(*global_rep)

        self.conv_proj = conv_1x1_out

        self.fusion = conv_3x3_out

        self.patch_h = patch_h
        self.patch_w = patch_w
        self.patch_area = self.patch_w * self.patch_h

        self.cnn_in_dim = in_channels
        self.cnn_out_dim = transformer_dim
        self.n_heads = num_heads
        self.ffn_dim = ffn_dim
        self.dropout = dropout
        self.attn_dropout = attn_dropout
        self.ffn_dropout = ffn_dropout
        self.dilation = dilation
        self.ffn_max_dim = ffn_dims[0]
        self.ffn_min_dim = ffn_dims[-1]
        self.var_ffn = var_ffn
        self.n_blocks = n_transformer_blocks
        self.conv_ksize = conv_ksize

    def unfolding(self, feature_map):
        patch_w, patch_h = self.patch_w, self.patch_h
        patch_area = int(patch_w * patch_h)
        batch_size, in_channels, orig_h, orig_w = feature_map.shape

        new_h = int(math.ceil(orig_h / self.patch_h) * self.patch_h)
        new_w = int(math.ceil(orig_w / self.patch_w) * self.patch_w)

        interpolate = False
        if new_w != orig_w or new_h != orig_h:
            # Note: Padding can be done, but then it needs to be handled in attention function.
            feature_map = F.interpolate(
                feature_map,
                size=(new_h, new_w),
                mode="bilinear",
                align_corners=False)
            interpolate = True

        # number of patches along width and height
        num_patch_w = new_w // patch_w  # n_w
        num_patch_h = new_h // patch_h  # n_h
        num_patches = num_patch_h * num_patch_w  # N

        # [B, C, H, W] --> [B * C * n_h, p_h, n_w, p_w]
        reshaped_fm = feature_map.reshape([
            batch_size * in_channels * num_patch_h, patch_h, num_patch_w,
            patch_w
        ])
        # [B * C * n_h, p_h, n_w, p_w] --> [B * C * n_h, n_w, p_h, p_w]
        transposed_fm = reshaped_fm.transpose([0, 2, 1, 3])
        # [B * C * n_h, n_w, p_h, p_w] --> [B, C, N, P] where P = p_h * p_w and N = n_h * n_w
        reshaped_fm = transposed_fm.reshape(
            [batch_size, in_channels, num_patches, patch_area])
        # [B, C, N, P] --> [B, P, N, C]
        transposed_fm = reshaped_fm.transpose([0, 3, 2, 1])
        # [B, P, N, C] --> [BP, N, C]
        patches = transposed_fm.reshape(
            [batch_size * patch_area, num_patches, in_channels])

        info_dict = {
            "orig_size": (orig_h, orig_w),
            "batch_size": batch_size,
            "interpolate": interpolate,
            "total_patches": num_patches,
            "num_patches_w": num_patch_w,
            "num_patches_h": num_patch_h
        }

        return patches, info_dict

    def folding(self, patches, info_dict):
        n_dim = patches.dim()
        assert n_dim == 3, "Tensor should be of shape BPxNxC. Got: {}".format(
            patches.shape)
        # [BP, N, C] --> [B, P, N, C]
        patches = patches.reshape([
            info_dict["batch_size"], self.patch_area,
            info_dict["total_patches"], patches.shape[2]
        ])

        batch_size, pixels, num_patches, channels = patches.shape
        num_patch_h = info_dict["num_patches_h"]
        num_patch_w = info_dict["num_patches_w"]

        # [B, P, N, C] --> [B, C, N, P]
        patches = patches.transpose([0, 3, 2, 1])

        # [B, C, N, P] --> [B*C*n_h, n_w, p_h, p_w]
        feature_map = patches.reshape([
            batch_size * channels * num_patch_h, num_patch_w, self.patch_h,
            self.patch_w
        ])
        # [B*C*n_h, n_w, p_h, p_w] --> [B*C*n_h, p_h, n_w, p_w]
        feature_map = feature_map.transpose([0, 2, 1, 3])
        # [B*C*n_h, p_h, n_w, p_w] --> [B, C, H, W]
        feature_map = feature_map.reshape([
            batch_size, channels, num_patch_h * self.patch_h,
            num_patch_w * self.patch_w
        ])
        if info_dict["interpolate"]:
            feature_map = F.interpolate(
                feature_map,
                size=info_dict["orig_size"],
                mode="bilinear",
                align_corners=False)
        return feature_map

    def forward(self, x):
        res = x

        # For MobileViTV3: Normal 3x3 convolution --> Depthwise 3x3 convolution
        fm_conv = self.local_rep(x)

        # convert feature map to patches
        patches, info_dict = self.unfolding(fm_conv)

        # learn global representations
        patches = self.global_rep(patches)

        # [B x Patch x Patches x C] --> [B x C x Patches x Patch]
        fm = self.folding(patches=patches, info_dict=info_dict)

        fm = self.conv_proj(fm)

        if self.fusion is not None:
            # For MobileViTV3: input+global --> local+global
            fm = self.fusion(paddle.concat((fm_conv, fm), axis=1))

        # For MobileViTV3: Skip connection
        fm = fm + res

        return fm


class LinearSelfAttention(nn.Layer):
    def __init__(self, embed_dim, attn_dropout=0.0, bias=True):
        super().__init__()
        self.embed_dim = embed_dim
        self.qkv_proj = nn.Conv2D(
            embed_dim, 1 + (2 * embed_dim), 1, bias_attr=bias)
        self.attn_dropout = nn.Dropout(p=attn_dropout)
        self.out_proj = nn.Conv2D(embed_dim, embed_dim, 1, bias_attr=bias)

    def forward(self, x):
        # [B, C, P, N] --> [B, h + 2d, P, N]
        qkv = self.qkv_proj(x)

        # Project x into query, key and value
        # Query --> [B, 1, P, N]
        # value, key --> [B, d, P, N]
        query, key, value = paddle.split(
            qkv, [1, self.embed_dim, self.embed_dim], axis=1)

        # apply softmax along N dimension
        context_scores = F.softmax(query, axis=-1)
        # Uncomment below line to visualize context scores
        # self.visualize_context_scores(context_scores=context_scores)
        context_scores = self.attn_dropout(context_scores)

        # Compute context vector
        # [B, d, P, N] x [B, 1, P, N] -> [B, d, P, N]
        context_vector = key * context_scores
        # [B, d, P, N] --> [B, d, P, 1]
        context_vector = paddle.sum(context_vector, axis=-1, keepdim=True)

        # combine context vector with values
        # [B, d, P, N] * [B, d, P, 1] --> [B, d, P, N]
        out = F.relu(value) * context_vector
        out = self.out_proj(out)
        return out


class LinearAttnFFN(nn.Layer):
    def __init__(self,
                 embed_dim: int,
                 ffn_latent_dim: int,
                 attn_dropout: Optional[float]=0.0,
                 dropout: Optional[float]=0.1,
                 ffn_dropout: Optional[float]=0.0,
                 norm_layer: Optional[str]=layer_norm_2d) -> None:
        super().__init__()
        attn_unit = LinearSelfAttention(
            embed_dim=embed_dim, attn_dropout=attn_dropout, bias=True)

        self.pre_norm_attn = nn.Sequential(
            norm_layer(num_channels=embed_dim),
            attn_unit,
            nn.Dropout(p=dropout))

        self.pre_norm_ffn = nn.Sequential(
            norm_layer(num_channels=embed_dim),
            nn.Conv2D(embed_dim, ffn_latent_dim, 1),
            nn.Silu(),
            nn.Dropout(p=ffn_dropout),
            nn.Conv2D(ffn_latent_dim, embed_dim, 1),
            nn.Dropout(p=dropout))

    def forward(self, x):
        # self-attention
        x = x + self.pre_norm_attn(x)
        # Feed forward network
        x = x + self.pre_norm_ffn(x)
        return x


class MobileViTV3BlockV2(nn.Layer):
    """
    This class defines the `MobileViTV3 block`
    """

    def __init__(self,
                 in_channels: int,
                 attn_unit_dim: int,
                 ffn_multiplier: float=2.0,
                 n_attn_blocks: Optional[int]=2,
                 attn_dropout: Optional[float]=0.0,
                 dropout: Optional[float]=0.0,
                 ffn_dropout: Optional[float]=0.0,
                 patch_h: Optional[int]=8,
                 patch_w: Optional[int]=8,
                 conv_ksize: Optional[int]=3,
                 dilation: Optional[int]=1,
                 attn_norm_layer: Optional[str]=layer_norm_2d):
        cnn_out_dim = attn_unit_dim

        padding = (conv_ksize - 1) // 2 * dilation
        conv_3x3_in = nn.Sequential(
            ('conv', nn.Conv2D(
                in_channels,
                in_channels,
                conv_ksize,
                bias_attr=False,
                padding=padding,
                dilation=dilation,
                groups=in_channels)), ('norm', nn.BatchNorm2D(in_channels)),
            ('act', nn.Silu()))
        conv_1x1_in = nn.Sequential(('conv', nn.Conv2D(
            in_channels, cnn_out_dim, 1, bias_attr=False)))

        super().__init__()
        self.local_rep = nn.Sequential(conv_3x3_in, conv_1x1_in)

        self.global_rep, attn_unit_dim = self._build_attn_layer(
            d_model=attn_unit_dim,
            ffn_mult=ffn_multiplier,
            n_layers=n_attn_blocks,
            attn_dropout=attn_dropout,
            dropout=dropout,
            ffn_dropout=ffn_dropout,
            attn_norm_layer=attn_norm_layer)

        # MobileViTV3: input changed from just global to local+global
        self.conv_proj = nn.Sequential(
            ('conv', nn.Conv2D(
                2 * cnn_out_dim, in_channels, 1, bias_attr=False)),
            ('norm', nn.BatchNorm2D(in_channels)))

        self.patch_h = patch_h
        self.patch_w = patch_w

    def _build_attn_layer(self,
                          d_model: int,
                          ffn_mult: float,
                          n_layers: int,
                          attn_dropout: float,
                          dropout: float,
                          ffn_dropout: float,
                          attn_norm_layer: nn.Layer):

        # ensure that dims are multiple of 16
        ffn_dims = [ffn_mult * d_model // 16 * 16] * n_layers

        global_rep = [
            LinearAttnFFN(
                embed_dim=d_model,
                ffn_latent_dim=ffn_dims[block_idx],
                attn_dropout=attn_dropout,
                dropout=dropout,
                ffn_dropout=ffn_dropout,
                norm_layer=attn_norm_layer) for block_idx in range(n_layers)
        ]
        global_rep.append(attn_norm_layer(num_channels=d_model))

        return nn.Sequential(*global_rep), d_model

    def unfolding(self, feature_map):
        batch_size, in_channels, img_h, img_w = feature_map.shape

        # [B, C, H, W] --> [B, C, P, N]
        patches = F.unfold(
            feature_map,
            kernel_sizes=[self.patch_h, self.patch_w],
            strides=[self.patch_h, self.patch_w])
        n_patches = img_h * img_w // (self.patch_h * self.patch_w)
        patches = patches.reshape(
            [batch_size, in_channels, self.patch_h * self.patch_w, n_patches])

        return patches, (img_h, img_w)

    def folding(self, patches, output_size: Tuple[int, int]):
        batch_size, in_dim, patch_size, n_patches = patches.shape

        # [B, C, P, N]
        patches = patches.reshape([batch_size, in_dim * patch_size, n_patches])

        feature_map = F.fold(
            patches,
            output_size,
            kernel_sizes=[self.patch_h, self.patch_w],
            strides=[self.patch_h, self.patch_w])

        return feature_map

    def forward(self, x):
        fm_conv = self.local_rep(x)

        # convert feature map to patches
        patches, output_size = self.unfolding(fm_conv)

        # learn global representations on all patches
        patches = self.global_rep(patches)

        # [B x Patch x Patches x C] --> [B x C x Patches x Patch]
        fm = self.folding(patches=patches, output_size=output_size)

        # MobileViTV3: local+global instead of only global
        fm = self.conv_proj(paddle.concat((fm, fm_conv), axis=1))

        # MobileViTV3: skip connection
        fm = fm + x

        return fm


class MobileViTV3(nn.Layer):
    """
        MobileViTV3:
    """

    def __init__(self,
                 mobilevit_config: Dict,
                 dropout=0.1,
                 class_num=1000,
                 classifier_dropout=0.1,
                 output_stride=None,
                 mobilevit_v2_based=False):
        super().__init__()
        self.round_nearest = 8
        self.dilation = 1
        self.dropout = dropout
        self.mobilevit_v2_based = mobilevit_v2_based

        dilate_l4 = dilate_l5 = False
        if output_stride == 8:
            dilate_l4 = True
            dilate_l5 = True
        elif output_stride == 16:
            dilate_l5 = True

        # store model configuration in a dictionary
        in_channels = mobilevit_config["layer0"]["img_channels"]
        out_channels = mobilevit_config["layer0"]["out_channels"]
        self.conv_1 = nn.Sequential(
            ('conv', nn.Conv2D(
                in_channels,
                out_channels,
                3,
                bias_attr=False,
                stride=2,
                padding=1)), ('norm', nn.BatchNorm2D(out_channels)),
            ('act', nn.Silu()))

        in_channels = out_channels
        self.layer_1, out_channels = self._make_layer(
            input_channel=in_channels, cfg=mobilevit_config["layer1"])

        in_channels = out_channels
        self.layer_2, out_channels = self._make_layer(
            input_channel=in_channels, cfg=mobilevit_config["layer2"])

        in_channels = out_channels
        self.layer_3, out_channels = self._make_layer(
            input_channel=in_channels, cfg=mobilevit_config["layer3"])

        in_channels = out_channels
        self.layer_4, out_channels = self._make_layer(
            input_channel=in_channels,
            cfg=mobilevit_config["layer4"],
            dilate=dilate_l4)

        in_channels = out_channels
        self.layer_5, out_channels = self._make_layer(
            input_channel=in_channels,
            cfg=mobilevit_config["layer5"],
            dilate=dilate_l5)

        if self.mobilevit_v2_based:
            self.conv_1x1_exp = nn.Identity()
        else:
            in_channels = out_channels
            out_channels = min(mobilevit_config["last_layer_exp_factor"] *
                               in_channels, 960)
            self.conv_1x1_exp = nn.Sequential(
                ('conv', nn.Conv2D(
                    in_channels, out_channels, 1, bias_attr=False)),
                ('norm', nn.BatchNorm2D(out_channels)), ('act', nn.Silu()))

        self.classifier = nn.Sequential()
        self.classifier.add_sublayer(
            name="global_pool",
            sublayer=nn.Sequential(nn.AdaptiveAvgPool2D(1), nn.Flatten()))
        if 0.0 < classifier_dropout < 1.0:
            self.classifier.add_sublayer(
                name="dropout", sublayer=nn.Dropout(p=classifier_dropout))
        self.classifier.add_sublayer(
            name="fc", sublayer=nn.Linear(out_channels, class_num))

        # weight initialization
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Conv2D):
            fan_in = m.weight.shape[1] * m.weight.shape[2] * m.weight.shape[3]
            fan_out = m.weight.shape[0] * m.weight.shape[2] * m.weight.shape[3]
            if self.mobilevit_v2_based:
                bound = 1.0 / fan_in**0.5
                nn.initializer.Uniform(-bound, bound)(m.weight)
                if m.bias is not None:
                    nn.initializer.Uniform(-bound, bound)(m.bias)
            else:
                nn.initializer.KaimingNormal(fan_in=fan_out)(m.weight)
                if m.bias is not None:
                    nn.initializer.Constant(0)(m.bias)
        elif isinstance(m, nn.BatchNorm2D):
            nn.initializer.Constant(1)(m.weight)
            nn.initializer.Constant(0)(m.bias)
        elif isinstance(m, nn.Linear):
            if self.mobilevit_v2_based:
                nn.initializer.XavierUniform()(m.weight)
            else:
                nn.initializer.TruncatedNormal(std=.02)(m.weight)
            if m.bias is not None:
                nn.initializer.Constant(0)(m.bias)

    def _make_layer(self, input_channel, cfg, dilate=False):
        block_type = cfg.get("block_type", "mobilevit")
        if block_type.lower() == "mobilevit":
            return self._make_mit_layer(
                input_channel=input_channel, cfg=cfg, dilate=dilate)
        else:
            return self._make_mobilenet_layer(
                input_channel=input_channel, cfg=cfg)

    def _make_mit_layer(self, input_channel, cfg, dilate=False):
        prev_dilation = self.dilation
        block = []
        stride = cfg.get("stride", 1)

        if stride == 2:
            if dilate:
                self.dilation *= 2
                stride = 1

            layer = InvertedResidual(
                in_channels=input_channel,
                out_channels=cfg.get("out_channels"),
                stride=stride,
                expand_ratio=cfg.get("mv_expand_ratio", 4),
                dilation=prev_dilation)

            block.append(layer)
            input_channel = cfg.get("out_channels")

        if self.mobilevit_v2_based:
            block.append(
                MobileViTV3BlockV2(
                    in_channels=input_channel,
                    attn_unit_dim=cfg["attn_unit_dim"],
                    ffn_multiplier=cfg.get("ffn_multiplier"),
                    n_attn_blocks=cfg.get("attn_blocks", 1),
                    ffn_dropout=0.,
                    attn_dropout=0.,
                    dilation=self.dilation,
                    patch_h=cfg.get("patch_h", 2),
                    patch_w=cfg.get("patch_w", 2)))
        else:
            head_dim = cfg.get("head_dim", 32)
            transformer_dim = cfg["transformer_channels"]
            ffn_dim = cfg.get("ffn_dim")
            if head_dim is None:
                num_heads = cfg.get("num_heads", 4)
                if num_heads is None:
                    num_heads = 4
                head_dim = transformer_dim // num_heads

            assert transformer_dim % head_dim == 0, (
                "Transformer input dimension should be divisible by head dimension. "
                "Got {} and {}.".format(transformer_dim, head_dim))

            block.append(
                MobileViTV3Block(
                    in_channels=input_channel,
                    transformer_dim=transformer_dim,
                    ffn_dim=ffn_dim,
                    n_transformer_blocks=cfg.get("transformer_blocks", 1),
                    patch_h=cfg.get("patch_h", 2),
                    patch_w=cfg.get("patch_w", 2),
                    dropout=self.dropout,
                    ffn_dropout=0.,
                    attn_dropout=0.,
                    head_dim=head_dim))

        return nn.Sequential(*block), input_channel

    def _make_mobilenet_layer(self, input_channel, cfg):
        output_channels = cfg.get("out_channels")
        num_blocks = cfg.get("num_blocks", 2)
        expand_ratio = cfg.get("expand_ratio", 4)
        block = []

        for i in range(num_blocks):
            stride = cfg.get("stride", 1) if i == 0 else 1

            layer = InvertedResidual(
                in_channels=input_channel,
                out_channels=output_channels,
                stride=stride,
                expand_ratio=expand_ratio)
            block.append(layer)
            input_channel = output_channels
        return nn.Sequential(*block), input_channel

    def extract_features(self, x):
        x = self.conv_1(x)
        x = self.layer_1(x)
        x = self.layer_2(x)
        x = self.layer_3(x)

        x = self.layer_4(x)
        x = self.layer_5(x)
        x = self.conv_1x1_exp(x)
        return x

    def forward(self, x):
        x = self.extract_features(x)
        x = self.classifier(x)
        return x


def _load_pretrained(pretrained, model, model_url, use_ssld=False):
    if pretrained is False:
        pass
    elif pretrained is True:
        load_dygraph_pretrain_from_url(model, model_url, use_ssld=use_ssld)
    elif isinstance(pretrained, str):
        load_dygraph_pretrain(model, pretrained)
    else:
        raise RuntimeError(
            "pretrained type is not available. Please use `string` or `boolean` type."
        )


def MobileViTV3_S(pretrained=False, use_ssld=False, **kwargs):
    mv2_exp_mult = 4
    mobilevit_config = {
        "layer0": {
            "img_channels": 3,
            "out_channels": 16,
        },
        "layer1": {
            "out_channels": 32,
            "expand_ratio": mv2_exp_mult,
            "num_blocks": 1,
            "stride": 1,
            "block_type": "mv2"
        },
        "layer2": {
            "out_channels": 64,
            "expand_ratio": mv2_exp_mult,
            "num_blocks": 3,
            "stride": 2,
            "block_type": "mv2"
        },
        "layer3": {  # 28x28
            "out_channels": 128,
            "transformer_channels": 144,
            "ffn_dim": 288,
            "transformer_blocks": 2,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "layer4": {  # 14x14
            "out_channels": 256,
            "transformer_channels": 192,
            "ffn_dim": 384,
            "transformer_blocks": 4,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "layer5": {  # 7x7
            "out_channels": 320,
            "transformer_channels": 240,
            "ffn_dim": 480,
            "transformer_blocks": 3,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "last_layer_exp_factor": 4
    }

    model = MobileViTV3(mobilevit_config, **kwargs)

    _load_pretrained(
        pretrained, model, MODEL_URLS["MobileViTV3_S"], use_ssld=use_ssld)
    return model


def MobileViTV3_XS(pretrained=False, use_ssld=False, **kwargs):
    mv2_exp_mult = 4
    mobilevit_config = {
        "layer0": {
            "img_channels": 3,
            "out_channels": 16,
        },
        "layer1": {
            "out_channels": 32,
            "expand_ratio": mv2_exp_mult,
            "num_blocks": 1,
            "stride": 1,
            "block_type": "mv2"
        },
        "layer2": {
            "out_channels": 48,
            "expand_ratio": mv2_exp_mult,
            "num_blocks": 3,
            "stride": 2,
            "block_type": "mv2"
        },
        "layer3": {  # 28x28
            "out_channels": 96,
            "transformer_channels": 96,
            "ffn_dim": 192,
            "transformer_blocks": 2,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "layer4": {  # 14x14
            "out_channels": 160,
            "transformer_channels": 120,
            "ffn_dim": 240,
            "transformer_blocks": 4,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "layer5": {  # 7x7
            "out_channels": 160,
            "transformer_channels": 144,
            "ffn_dim": 288,
            "transformer_blocks": 3,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "last_layer_exp_factor": 4
    }

    model = MobileViTV3(mobilevit_config, **kwargs)

    _load_pretrained(
        pretrained, model, MODEL_URLS["MobileViTV3_XS"], use_ssld=use_ssld)
    return model


def MobileViTV3_XXS(pretrained=False, use_ssld=False, **kwargs):
    mv2_exp_mult = 2
    mobilevit_config = {
        "layer0": {
            "img_channels": 3,
            "out_channels": 16,
        },
        "layer1": {
            "out_channels": 16,
            "expand_ratio": mv2_exp_mult,
            "num_blocks": 1,
            "stride": 1,
            "block_type": "mv2"
        },
        "layer2": {
            "out_channels": 24,
            "expand_ratio": mv2_exp_mult,
            "num_blocks": 3,
            "stride": 2,
            "block_type": "mv2"
        },
        "layer3": {  # 28x28
            "out_channels": 64,
            "transformer_channels": 64,
            "ffn_dim": 128,
            "transformer_blocks": 2,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "layer4": {  # 14x14
            "out_channels": 80,
            "transformer_channels": 80,
            "ffn_dim": 160,
            "transformer_blocks": 4,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "layer5": {  # 7x7
            "out_channels": 128,
            "transformer_channels": 96,
            "ffn_dim": 192,
            "transformer_blocks": 3,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "last_layer_exp_factor": 4
    }

    model = MobileViTV3(mobilevit_config, **kwargs)

    _load_pretrained(
        pretrained, model, MODEL_URLS["MobileViTV3_XXS"], use_ssld=use_ssld)
    return model


def MobileViTV3_S_L2(pretrained=False, use_ssld=False, **kwargs):
    mv2_exp_mult = 4
    mobilevit_config = {
        "layer0": {
            "img_channels": 3,
            "out_channels": 16,
        },
        "layer1": {
            "out_channels": 32,
            "expand_ratio": mv2_exp_mult,
            "num_blocks": 1,
            "stride": 1,
            "block_type": "mv2"
        },
        "layer2": {
            "out_channels": 64,
            "expand_ratio": mv2_exp_mult,
            "num_blocks": 3,
            "stride": 2,
            "block_type": "mv2"
        },
        "layer3": {  # 28x28
            "out_channels": 128,
            "transformer_channels": 144,
            "ffn_dim": 288,
            "transformer_blocks": 2,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "layer4": {  # 14x14
            "out_channels": 256,
            "transformer_channels": 192,
            "ffn_dim": 384,
            "transformer_blocks": 2,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "layer5": {  # 7x7
            "out_channels": 320,
            "transformer_channels": 240,
            "ffn_dim": 480,
            "transformer_blocks": 3,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "last_layer_exp_factor": 4
    }

    model = MobileViTV3(mobilevit_config, **kwargs)

    _load_pretrained(
        pretrained, model, MODEL_URLS["MobileViTV3_S_L2"], use_ssld=use_ssld)
    return model


def MobileViTV3_XS_L2(pretrained=False, use_ssld=False, **kwargs):
    mv2_exp_mult = 4
    mobilevit_config = {
        "layer0": {
            "img_channels": 3,
            "out_channels": 16,
        },
        "layer1": {
            "out_channels": 32,
            "expand_ratio": mv2_exp_mult,
            "num_blocks": 1,
            "stride": 1,
            "block_type": "mv2"
        },
        "layer2": {
            "out_channels": 48,
            "expand_ratio": mv2_exp_mult,
            "num_blocks": 3,
            "stride": 2,
            "block_type": "mv2"
        },
        "layer3": {  # 28x28
            "out_channels": 96,
            "transformer_channels": 96,
            "ffn_dim": 192,
            "transformer_blocks": 2,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "layer4": {  # 14x14
            "out_channels": 160,
            "transformer_channels": 120,
            "ffn_dim": 240,
            "transformer_blocks": 2,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "layer5": {  # 7x7
            "out_channels": 160,
            "transformer_channels": 144,
            "ffn_dim": 288,
            "transformer_blocks": 3,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "last_layer_exp_factor": 4
    }

    model = MobileViTV3(mobilevit_config, **kwargs)

    _load_pretrained(
        pretrained, model, MODEL_URLS["MobileViTV3_XS_L2"], use_ssld=use_ssld)
    return model


def MobileViTV3_XXS_L2(pretrained=False, use_ssld=False, **kwargs):
    mv2_exp_mult = 2
    mobilevit_config = {
        "layer0": {
            "img_channels": 3,
            "out_channels": 16,
        },
        "layer1": {
            "out_channels": 16,
            "expand_ratio": mv2_exp_mult,
            "num_blocks": 1,
            "stride": 1,
            "block_type": "mv2"
        },
        "layer2": {
            "out_channels": 24,
            "expand_ratio": mv2_exp_mult,
            "num_blocks": 3,
            "stride": 2,
            "block_type": "mv2"
        },
        "layer3": {  # 28x28
            "out_channels": 64,
            "transformer_channels": 64,
            "ffn_dim": 128,
            "transformer_blocks": 2,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "layer4": {  # 14x14
            "out_channels": 80,
            "transformer_channels": 80,
            "ffn_dim": 160,
            "transformer_blocks": 2,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "layer5": {  # 7x7
            "out_channels": 128,
            "transformer_channels": 96,
            "ffn_dim": 192,
            "transformer_blocks": 3,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": mv2_exp_mult,
            "head_dim": None,
            "num_heads": 4,
            "block_type": "mobilevit"
        },
        "last_layer_exp_factor": 4
    }

    model = MobileViTV3(mobilevit_config, **kwargs)

    _load_pretrained(
        pretrained, model, MODEL_URLS["MobileViTV3_XXS_L2"], use_ssld=use_ssld)
    return model


def MobileViTV3_x1_0(pretrained=False, use_ssld=False, **kwargs):
    mobilevit_config = {
        "layer0": {
            "img_channels": 3,
            "out_channels": 32,
        },
        "layer1": {
            "out_channels": 64,
            "expand_ratio": 2,
            "num_blocks": 1,
            "stride": 1,
            "block_type": "mv2",
        },
        "layer2": {
            "out_channels": 128,
            "expand_ratio": 2,
            "num_blocks": 2,
            "stride": 2,
            "block_type": "mv2",
        },
        "layer3": {  # 28x28
            "out_channels": 256,
            "attn_unit_dim": 128,
            "ffn_multiplier": 2,
            "attn_blocks": 2,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": 2,
            "block_type": "mobilevit",
        },
        "layer4": {  # 14x14
            "out_channels": 384,
            "attn_unit_dim": 192,
            "ffn_multiplier": 2,
            "attn_blocks": 4,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": 2,
            "block_type": "mobilevit",
        },
        "layer5": {  # 7x7
            "out_channels": 512,
            "attn_unit_dim": 256,
            "ffn_multiplier": 2,
            "attn_blocks": 3,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": 2,
            "block_type": "mobilevit",
        },
        "last_layer_exp_factor": 4,
    }

    model = MobileViTV3(mobilevit_config, mobilevit_v2_based=True, **kwargs)

    _load_pretrained(
        pretrained, model, MODEL_URLS["MobileViTV3_x1_0"], use_ssld=use_ssld)
    return model


def MobileViTV3_x0_75(pretrained=False, use_ssld=False, **kwargs):
    mobilevit_config = {
        "layer0": {
            "img_channels": 3,
            "out_channels": 24,
        },
        "layer1": {
            "out_channels": 48,
            "expand_ratio": 2,
            "num_blocks": 1,
            "stride": 1,
            "block_type": "mv2",
        },
        "layer2": {
            "out_channels": 96,
            "expand_ratio": 2,
            "num_blocks": 2,
            "stride": 2,
            "block_type": "mv2",
        },
        "layer3": {  # 28x28
            "out_channels": 192,
            "attn_unit_dim": 96,
            "ffn_multiplier": 2,
            "attn_blocks": 2,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": 2,
            "block_type": "mobilevit",
        },
        "layer4": {  # 14x14
            "out_channels": 288,
            "attn_unit_dim": 144,
            "ffn_multiplier": 2,
            "attn_blocks": 4,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": 2,
            "block_type": "mobilevit",
        },
        "layer5": {  # 7x7
            "out_channels": 384,
            "attn_unit_dim": 192,
            "ffn_multiplier": 2,
            "attn_blocks": 3,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": 2,
            "block_type": "mobilevit",
        },
        "last_layer_exp_factor": 4,
    }

    model = MobileViTV3(mobilevit_config, mobilevit_v2_based=True, **kwargs)

    _load_pretrained(
        pretrained, model, MODEL_URLS["MobileViTV3_x0_75"], use_ssld=use_ssld)
    return model


def MobileViTV3_x0_5(pretrained=False, use_ssld=False, **kwargs):
    mobilevit_config = {
        "layer0": {
            "img_channels": 3,
            "out_channels": 16,
        },
        "layer1": {
            "out_channels": 32,
            "expand_ratio": 2,
            "num_blocks": 1,
            "stride": 1,
            "block_type": "mv2",
        },
        "layer2": {
            "out_channels": 64,
            "expand_ratio": 2,
            "num_blocks": 2,
            "stride": 2,
            "block_type": "mv2",
        },
        "layer3": {  # 28x28
            "out_channels": 128,
            "attn_unit_dim": 64,
            "ffn_multiplier": 2,
            "attn_blocks": 2,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": 2,
            "block_type": "mobilevit",
        },
        "layer4": {  # 14x14
            "out_channels": 192,
            "attn_unit_dim": 96,
            "ffn_multiplier": 2,
            "attn_blocks": 4,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": 2,
            "block_type": "mobilevit",
        },
        "layer5": {  # 7x7
            "out_channels": 256,
            "attn_unit_dim": 128,
            "ffn_multiplier": 2,
            "attn_blocks": 3,
            "patch_h": 2,
            "patch_w": 2,
            "stride": 2,
            "mv_expand_ratio": 2,
            "block_type": "mobilevit",
        },
        "last_layer_exp_factor": 4,
    }

    model = MobileViTV3(mobilevit_config, mobilevit_v2_based=True, **kwargs)

    _load_pretrained(
        pretrained, model, MODEL_URLS["MobileViTV3_x0_5"], use_ssld=use_ssld)
    return model
