"""
MobileNet v3, adapted from
https://github.com/pytorch/vision/blob/main/torchvision/models/mobilenetv3.py

Paper "Searching for MobileNetV3", https://arxiv.org/abs/1905.02244

Changes from original:
* Using nn.BatchNorm2d with eps 1e-5 instead of 1e-3
"""

# Reference license: BSD 3-Clause

from collections import OrderedDict
from collections.abc import Callable
from typing import Any
from typing import Optional

import torch
from torch import nn
from torchvision.ops import Conv2dNormActivation
from torchvision.ops import SqueezeExcitation

from birder.model_registry import registry
from birder.net.base import DetectorBackbone
from birder.net.base import make_divisible


class InvertedResidualConfig:
    """
    Stores information listed at Tables 1 and 2 of the MobileNet v3 paper
    """

    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        expanded_channels: int,
        kernel_size: tuple[int, int],
        stride: tuple[int, int],
        padding: tuple[int, int],
        width_multi: float,
        use_se: bool,
        activation: Callable[..., nn.Module] = nn.ReLU,
    ):
        self.in_channels = make_divisible(in_channels * width_multi, 8)
        self.out_channels = make_divisible(out_channels * width_multi, 8)
        self.expanded_channels = make_divisible(expanded_channels * width_multi, 8)
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.use_se = use_se
        self.activation = activation


class InvertedResidual(nn.Module):
    def __init__(self, cnf: InvertedResidualConfig) -> None:
        super().__init__()
        if cnf.stride[0] == 1 and cnf.stride[1] == 1 and cnf.in_channels == cnf.out_channels:
            self.shortcut = True
        else:
            self.shortcut = False

        layers = []
        # Expand
        if cnf.expanded_channels != cnf.in_channels:
            layers.append(
                Conv2dNormActivation(
                    cnf.in_channels,
                    cnf.expanded_channels,
                    kernel_size=(1, 1),
                    stride=(1, 1),
                    padding=(0, 0),
                    bias=False,
                    activation_layer=cnf.activation,
                )
            )

        # Depthwise
        layers.append(
            Conv2dNormActivation(
                cnf.expanded_channels,
                cnf.expanded_channels,
                kernel_size=cnf.kernel_size,
                stride=cnf.stride,
                padding=cnf.padding,
                groups=cnf.expanded_channels,
                bias=False,
                activation_layer=cnf.activation,
            )
        )

        if cnf.use_se is True:
            squeeze_channels = make_divisible(cnf.expanded_channels // 4, 8)
            layers.append(SqueezeExcitation(cnf.expanded_channels, squeeze_channels, scale_activation=nn.Hardsigmoid))

        # Project
        layers.append(
            Conv2dNormActivation(
                cnf.expanded_channels,
                cnf.out_channels,
                kernel_size=(1, 1),
                stride=(1, 1),
                padding=(0, 0),
                bias=False,
                activation_layer=None,
            )
        )

        self.block = nn.Sequential(*layers)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        if self.shortcut is True:
            return x + self.block(x)

        return self.block(x)


# pylint: disable=invalid-name
class MobileNet_v3(DetectorBackbone):
    def __init__(
        self,
        input_channels: int,
        num_classes: int,
        *,
        config: Optional[dict[str, Any]] = None,
        size: Optional[tuple[int, int]] = None,
    ) -> None:
        super().__init__(input_channels, num_classes, config=config, size=size)
        assert self.config is not None, "must set config"

        alpha: float = self.config["alpha"]
        large: bool = self.config["large"]
        self.mlp_head = self.config.get("mlp_head", True)

        if large is True:
            last_channels = int(round(1280 * max(1.0, alpha)))
            net_settings = [
                InvertedResidualConfig(16, 16, 16, (3, 3), (1, 1), (1, 1), alpha, False, nn.ReLU),
                InvertedResidualConfig(16, 24, 64, (3, 3), (2, 2), (1, 1), alpha, False, nn.ReLU),
                InvertedResidualConfig(24, 24, 72, (3, 3), (1, 1), (1, 1), alpha, False, nn.ReLU),
                InvertedResidualConfig(24, 40, 72, (5, 5), (2, 2), (2, 2), alpha, True, nn.ReLU),
                InvertedResidualConfig(40, 40, 120, (5, 5), (1, 1), (2, 2), alpha, True, nn.ReLU),
                InvertedResidualConfig(40, 40, 120, (5, 5), (1, 1), (2, 2), alpha, True, nn.ReLU),
                InvertedResidualConfig(40, 80, 240, (3, 3), (2, 2), (1, 1), alpha, False, nn.Hardswish),
                InvertedResidualConfig(80, 80, 200, (3, 3), (1, 1), (1, 1), alpha, False, nn.Hardswish),
                InvertedResidualConfig(80, 80, 184, (3, 3), (1, 1), (1, 1), alpha, False, nn.Hardswish),
                InvertedResidualConfig(80, 80, 184, (3, 3), (1, 1), (1, 1), alpha, False, nn.Hardswish),
                InvertedResidualConfig(80, 112, 480, (3, 3), (1, 1), (1, 1), alpha, True, nn.Hardswish),
                InvertedResidualConfig(112, 112, 672, (3, 3), (1, 1), (1, 1), alpha, True, nn.Hardswish),
                InvertedResidualConfig(112, 160, 672, (5, 5), (2, 2), (2, 2), alpha, True, nn.Hardswish),
                InvertedResidualConfig(160, 160, 960, (5, 5), (1, 1), (2, 2), alpha, True, nn.Hardswish),
                InvertedResidualConfig(160, 160, 960, (5, 5), (1, 1), (2, 2), alpha, True, nn.Hardswish),
            ]
        else:
            last_channels = int(round(1024 * max(1.0, alpha)))
            net_settings = [
                InvertedResidualConfig(16, 16, 16, (3, 3), (2, 2), (1, 1), alpha, True, nn.ReLU),
                InvertedResidualConfig(16, 24, 72, (3, 3), (2, 2), (1, 1), alpha, False, nn.ReLU),
                InvertedResidualConfig(24, 24, 88, (3, 3), (1, 1), (1, 1), alpha, False, nn.ReLU),
                InvertedResidualConfig(24, 40, 96, (5, 5), (2, 2), (2, 2), alpha, True, nn.Hardswish),
                InvertedResidualConfig(40, 40, 240, (5, 5), (1, 1), (2, 2), alpha, True, nn.Hardswish),
                InvertedResidualConfig(40, 40, 240, (5, 5), (1, 1), (2, 2), alpha, True, nn.Hardswish),
                InvertedResidualConfig(40, 48, 120, (5, 5), (1, 1), (2, 2), alpha, True, nn.Hardswish),
                InvertedResidualConfig(48, 48, 144, (5, 5), (1, 1), (2, 2), alpha, True, nn.Hardswish),
                InvertedResidualConfig(48, 96, 288, (5, 5), (2, 2), (2, 2), alpha, True, nn.Hardswish),
                InvertedResidualConfig(96, 96, 576, (5, 5), (1, 1), (2, 2), alpha, True, nn.Hardswish),
                InvertedResidualConfig(96, 96, 576, (5, 5), (1, 1), (2, 2), alpha, True, nn.Hardswish),
            ]

        self.stem = Conv2dNormActivation(
            self.input_channels,
            net_settings[0].in_channels,
            kernel_size=(3, 3),
            stride=(2, 2),
            padding=(1, 1),
            bias=False,
            activation_layer=nn.Hardswish,
        )

        layers: list[nn.Module] = []
        i = 0
        stages: OrderedDict[str, nn.Module] = OrderedDict()
        return_channels: list[int] = []
        for idx, block_settings in enumerate(net_settings):
            if block_settings.stride[0] > 1 or block_settings.stride[1] > 1:
                stages[f"stage{i}"] = nn.Sequential(*layers)
                return_channels.append(net_settings[idx - 1].out_channels)
                layers = []
                i += 1

            layers.append(InvertedResidual(block_settings))

        stages[f"stage{i}"] = nn.Sequential(*layers)
        return_channels.append(net_settings[-1].out_channels)

        self.body = nn.Sequential(stages)
        self.features = nn.Sequential(
            Conv2dNormActivation(
                net_settings[-1].out_channels,
                net_settings[-1].out_channels * 6,
                kernel_size=(1, 1),
                stride=(1, 1),
                padding=(0, 0),
                bias=False,
                activation_layer=nn.Hardswish,
            ),
            nn.AdaptiveAvgPool2d(output_size=(1, 1)),
            nn.Flatten(1),
        )
        self.return_channels = return_channels[1:5]
        self.embedding_size = net_settings[-1].out_channels * 6
        self.last_channels = last_channels
        self.classifier = self.create_classifier()

        # Weight initialization
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out")
                if m.bias is not None:
                    nn.init.zeros_(m.bias)

            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)

            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.zeros_(m.bias)

    def detection_features(self, x: torch.Tensor) -> dict[str, torch.Tensor]:
        x = self.stem(x)

        out = {}
        for name, module in self.body.named_children():
            x = module(x)
            if name in self.return_stages:
                out[name] = x

        return out

    def freeze_stages(self, up_to_stage: int) -> None:
        for param in self.stem.parameters():
            param.requires_grad_(False)

        for idx, module in enumerate(self.body.children()):
            if idx >= up_to_stage:
                break

            for param in module.parameters():
                param.requires_grad_(False)

    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
        x = self.stem(x)
        return self.body(x)

    def embedding(self, x: torch.Tensor) -> torch.Tensor:
        x = self.forward_features(x)
        return self.features(x)

    def create_classifier(
        self, embed_dim: Optional[int] = None, head_bias: Optional[bool] = None, mlp_head: Optional[bool] = None
    ) -> nn.Module:
        if self.num_classes == 0:
            return nn.Identity()

        if embed_dim is None:
            embed_dim = self.embedding_size
        if head_bias is None:
            head_bias = self.head_bias
        if mlp_head is None:
            mlp_head = self.mlp_head

        if mlp_head is True:
            return nn.Sequential(
                nn.Linear(embed_dim, self.last_channels),
                nn.Hardswish(inplace=True),
                nn.Dropout(p=0.2, inplace=True),
                nn.Linear(self.last_channels, self.num_classes, bias=head_bias),
            )

        return nn.Linear(embed_dim, self.num_classes, bias=head_bias)


registry.register_model_config("mobilenet_v3_small_0_25", MobileNet_v3, config={"alpha": 0.25, "large": False})
registry.register_model_config("mobilenet_v3_small_0_5", MobileNet_v3, config={"alpha": 0.5, "large": False})
registry.register_model_config("mobilenet_v3_small_0_75", MobileNet_v3, config={"alpha": 0.75, "large": False})
registry.register_model_config("mobilenet_v3_small_1_0", MobileNet_v3, config={"alpha": 1.0, "large": False})
registry.register_model_config("mobilenet_v3_small_1_25", MobileNet_v3, config={"alpha": 1.25, "large": False})
registry.register_model_config("mobilenet_v3_small_1_5", MobileNet_v3, config={"alpha": 1.5, "large": False})
registry.register_model_config("mobilenet_v3_small_1_75", MobileNet_v3, config={"alpha": 1.75, "large": False})
registry.register_model_config("mobilenet_v3_small_2_0", MobileNet_v3, config={"alpha": 2.0, "large": False})

registry.register_model_config("mobilenet_v3_large_0_25", MobileNet_v3, config={"alpha": 0.25, "large": True})
registry.register_model_config("mobilenet_v3_large_0_5", MobileNet_v3, config={"alpha": 0.5, "large": True})
registry.register_model_config("mobilenet_v3_large_0_75", MobileNet_v3, config={"alpha": 0.75, "large": True})
registry.register_model_config("mobilenet_v3_large_1_0", MobileNet_v3, config={"alpha": 1.0, "large": True})
registry.register_model_config("mobilenet_v3_large_1_25", MobileNet_v3, config={"alpha": 1.25, "large": True})
registry.register_model_config("mobilenet_v3_large_1_5", MobileNet_v3, config={"alpha": 1.5, "large": True})
registry.register_model_config("mobilenet_v3_large_1_75", MobileNet_v3, config={"alpha": 1.75, "large": True})
registry.register_model_config("mobilenet_v3_large_2_0", MobileNet_v3, config={"alpha": 2.0, "large": True})


registry.register_weights(
    "mobilenet_v3_small_1_0_il-common",
    {
        "description": "MobileNet v3 small (1.0 multiplier) model trained on the il-common dataset",
        "resolution": (256, 256),
        "formats": {
            "pt": {
                "file_size": 7.4,
                "sha256": "ac53227f7513fd0c0b5204ee57403de2ab6c74c4e4d1061b9168596c6b5cea48",
            }
        },
        "net": {"network": "mobilenet_v3_small_1_0", "tag": "il-common"},
    },
)
registry.register_weights(
    "mobilenet_v3_large_0_75_il-common",
    {
        "description": "MobileNet v3 large (0.75 multiplier) model trained on the il-common dataset",
        "resolution": (256, 256),
        "formats": {
            "pt": {
                "file_size": 12.1,
                "sha256": "92412316f3dbcc41e4f3186acb50027e87ce0ea3ea1d6b5a726ea883fea20b8e",
            }
        },
        "net": {"network": "mobilenet_v3_large_0_75", "tag": "il-common"},
    },
)
registry.register_weights(
    "mobilenet_v3_large_1_0_il-common",
    {
        "description": "MobileNet v3 large (1.0 multiplier) model trained on the il-common dataset",
        "resolution": (256, 256),
        "formats": {
            "pt": {
                "file_size": 18.1,
                "sha256": "85bc65482716368c4fd69d4c699b61330fdf806d99e6c88ffbeda54886c2f3d4",
            }
        },
        "net": {"network": "mobilenet_v3_large_1_0", "tag": "il-common"},
    },
)
