"""
B-cos DenseNet models

Modified from https://github.com/pytorch/vision/blob/0504df5ddf9431909130e7788faf05446bb8/torchvision/models/densenet.py

CIFAR10 modifications from
https://github.com/gpleiss/efficient_densenet_pytorch/blob/master/models/densenet.py
"""
import math
import re
from collections import OrderedDict
from typing import Any, Callable, List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as cp
from torch import Tensor
from torch.hub import load_state_dict_from_url

from bcos.common import BcosUtilMixin
from bcos.modules import BcosConv2d, LogitLayer, norms

__all__ = [
    "BcosDenseNet",
    "densenet121",
    "densenet169",
    "densenet201",
    "densenet161",
    # c10
    "cifar10_densenet250",
    "cifar10_densenet100",
    "cifar10_densenet190",
]


DEFAULT_NORM_LAYER = norms.NoBias(norms.DetachablePositionNorm2d)
DEFAULT_CONV_LAYER = BcosConv2d


class _DenseLayer(nn.Module):
    def __init__(
        self,
        num_input_features: int,
        growth_rate: int,
        bn_size: int,
        drop_rate: float,
        memory_efficient: bool = False,
        norm_layer: Callable[..., nn.Module] = DEFAULT_NORM_LAYER,
        conv_layer: Callable[..., nn.Module] = DEFAULT_CONV_LAYER,
    ) -> None:
        super(_DenseLayer, self).__init__()
        # Diff to torchvision: Removed ReLU and replaced BatchNorm with norm_layer
        self.norm1 = norm_layer(num_input_features)
        # Diff to torchvision: Replace Conv2d with ProjectionConv2d
        self.conv1 = conv_layer(
            num_input_features,
            bn_size * growth_rate,
            kernel_size=1,
            stride=1,
            padding=0,
        )
        self.norm2 = norm_layer(bn_size * growth_rate)
        self.conv2 = conv_layer(
            bn_size * growth_rate,
            growth_rate,
            kernel_size=3,
            stride=1,
            padding=1,
        )
        # Diff End
        # Diff End
        self.drop_rate = float(drop_rate)
        self.memory_efficient = memory_efficient

    def bn_function(self, inputs: List[Tensor]) -> Tensor:
        concated_features = torch.cat(inputs, 1)
        # Diff to torchvision: Deleted relu
        bottleneck_output = self.conv1(self.norm1(concated_features))  # noqa: T484
        # Diff End
        return bottleneck_output

    # todo: rewrite when torchscript supports any
    def any_requires_grad(self, input: List[Tensor]) -> bool:
        for tensor in input:
            if tensor.requires_grad:
                return True
        return False

    @torch.jit.unused  # noqa: T484
    def call_checkpoint_bottleneck(self, input: List[Tensor]) -> Tensor:
        def closure(*inputs):
            return self.bn_function(inputs)

        return cp.checkpoint(closure, *input)

    @torch.jit._overload_method  # noqa: F811
    def forward(self, input: List[Tensor]) -> Tensor:
        pass

    @torch.jit._overload_method  # noqa: F811
    def forward(self, input: Tensor) -> Tensor:  # noqa: F811
        pass

    # torchscript does not yet support *args, so we overload method
    # allowing it to take either a List[Tensor] or single Tensor
    def forward(self, input: Tensor) -> Tensor:  # noqa: F811
        if isinstance(input, Tensor):
            prev_features = [input]
        else:
            prev_features = input

        if self.memory_efficient and self.any_requires_grad(prev_features):
            if torch.jit.is_scripting():
                raise Exception("Memory Efficient not supported in JIT")

            bottleneck_output = self.call_checkpoint_bottleneck(prev_features)
        else:
            bottleneck_output = self.bn_function(prev_features)
        # Diff to torchvision: Deleted relu
        new_features = self.conv2(self.norm2(bottleneck_output))
        # Diff End

        if self.drop_rate > 0:
            new_features = F.dropout(
                new_features, p=self.drop_rate, training=self.training
            )
        return new_features


class _DenseBlock(nn.ModuleDict):
    _version = 2

    def __init__(
        self,
        num_layers: int,
        num_input_features: int,
        bn_size: int,
        growth_rate: int,
        drop_rate: float,
        memory_efficient: bool = False,
        norm_layer: Callable[..., nn.Module] = DEFAULT_NORM_LAYER,
        conv_layer: Callable[..., nn.Module] = DEFAULT_CONV_LAYER,
    ) -> None:
        super(_DenseBlock, self).__init__()
        for i in range(num_layers):
            layer = _DenseLayer(
                num_input_features + i * growth_rate,
                growth_rate=growth_rate,
                bn_size=bn_size,
                drop_rate=drop_rate,
                memory_efficient=memory_efficient,
                norm_layer=norm_layer,
                conv_layer=conv_layer,
            )
            self.add_module("denselayer%d" % (i + 1), layer)

    def forward(self, init_features: Tensor) -> Tensor:
        features = [init_features]
        for name, layer in self.items():
            new_features = layer(features)
            features.append(new_features)
        return torch.cat(features, 1)


class _Transition(nn.Sequential):
    def __init__(
        self,
        num_input_features: int,
        num_output_features: int,
        norm_layer: Callable[..., nn.Module] = DEFAULT_NORM_LAYER,
        conv_layer: Callable[..., nn.Module] = DEFAULT_CONV_LAYER,
    ) -> None:
        super(_Transition, self).__init__()
        # Diff to torchvision: Deleted relu and changed conv to ProjectionConv2d
        self.add_module("norm", norm_layer(num_input_features))
        self.conv = conv_layer(
            num_input_features,
            num_output_features,
            kernel_size=1,
            stride=1,
            padding=0,
        )
        # Diff End

        self.add_module("pool", nn.AvgPool2d(kernel_size=2, stride=2))


class BcosDenseNet(BcosUtilMixin, nn.Module):
    r"""Densenet-BC model class, based on
    `"Densely Connected Convolutional Networks" <https://arxiv.org/pdf/1608.06993.pdf>`_.

    Args:
        growth_rate (int) - how many filters to add each layer (`k` in paper)
        block_config (list of 4 ints) - how many layers in each pooling block
        num_init_features (int) - the number of filters to learn in the first convolution layer
        bn_size (int) - multiplicative factor for number of bottle neck layers
          (i.e. bn_size * k features in the bottleneck layer)
        drop_rate (float) - dropout rate after each dense layer
        num_classes (int) - number of classification classes
        memory_efficient (bool) - If True, uses checkpointing. Much more memory efficient,
          but slower. Default: *False*. See `"paper" <https://arxiv.org/pdf/1707.06990.pdf>`_.
    """

    def __init__(
        self,
        growth_rate: int = 32,
        block_config: Tuple[int, int, int, int] = (6, 12, 24, 16),
        num_init_features: int = 64,
        bn_size: int = 4,
        drop_rate: float = 0,
        num_classes: int = 1000,
        in_chans: int = 6,
        memory_efficient: bool = False,
        norm_layer: Callable[..., nn.Module] = DEFAULT_NORM_LAYER,
        conv_layer: Callable[..., nn.Module] = DEFAULT_CONV_LAYER,
        small_inputs: bool = False,  # True for 32x32 images from gpleiss' impl
        logit_bias: Optional[float] = None,
        logit_temperature: Optional[float] = None,
    ) -> None:
        super(BcosDenseNet, self).__init__()

        # First convolution
        # Diff to torchvision: Deleted ReLU, changed Conv2d for ProjectionConv2d and
        # MaxPool for AvgPool
        if small_inputs:
            self.features = nn.Sequential(
                OrderedDict(
                    [
                        (
                            "conv0",
                            conv_layer(
                                in_chans,
                                num_init_features,
                                kernel_size=3,
                                stride=1,
                                padding=1,
                            ),
                        )
                    ]
                )
            )
        else:
            self.features = nn.Sequential(
                OrderedDict(
                    [
                        (
                            "conv0",
                            conv_layer(
                                in_chans,
                                num_init_features,
                                kernel_size=7,
                                stride=2,
                                padding=3,
                            ),
                        ),
                        ("norm0", norm_layer(num_init_features)),
                        ("pool0", nn.AvgPool2d(kernel_size=3, stride=2, padding=1)),
                    ]
                )
            )
        # Diff End
        # Each denseblock
        num_features = num_init_features
        for i, num_layers in enumerate(block_config):
            block = _DenseBlock(
                num_layers=num_layers,
                num_input_features=num_features,
                bn_size=bn_size,
                growth_rate=growth_rate,
                drop_rate=drop_rate,
                memory_efficient=memory_efficient,
                norm_layer=norm_layer,
                conv_layer=conv_layer,
            )
            self.features.add_module("denseblock%d" % (i + 1), block)
            num_features = num_features + num_layers * growth_rate
            if i != len(block_config) - 1:
                trans = _Transition(
                    num_input_features=num_features,
                    num_output_features=num_features // 2,
                    norm_layer=norm_layer,
                    conv_layer=conv_layer,
                )
                self.features.add_module("transition%d" % (i + 1), trans)
                num_features = num_features // 2

        # Final norm layer
        self.features.add_module("norm5", norm_layer(num_features))

        # Diff to torchvision: changed Linear layer to BcosConv (conv classifier)
        self.classifier = conv_layer(num_features, num_classes, kernel_size=1)
        self.num_classes = num_classes
        self.logit_layer = LogitLayer(
            logit_temperature=logit_temperature,
            logit_bias=logit_bias or -math.log(num_classes - 1),
        )
        # Diff End

        # Official init from torch repo.
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight)
            elif isinstance(m, nn.BatchNorm2d):
                if m.weight is not None:
                    nn.init.constant_(m.weight, 1)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)

    def forward(self, x: Tensor) -> Tensor:
        # Diff to torchvision: Deleted relu
        features = self.features(x)
        # out = F.relu(features, inplace=True)
        out = self.classifier(features)
        out = F.adaptive_avg_pool2d(out, (1, 1))
        out = torch.flatten(out, 1)
        out = self.logit_layer(out)
        # Diff End
        return out

    def get_classifier(self) -> nn.Module:
        """Returns the classifier part of the model. Note this comes before global pooling."""
        return self.classifier

    def get_feature_extractor(self) -> nn.Module:
        """Returns the feature extractor part of the model. Without global pooling."""
        return self.features


def _load_state_dict(model: nn.Module, model_url: str, progress: bool) -> None:
    # '.'s are no longer allowed in module names, but previous _DenseLayer
    # has keys 'norm.1', 'relu.1', 'conv.1', 'norm.2', 'relu.2', 'conv.2'.
    # They are also in the checkpoints in model_urls. This pattern is used
    # to find such keys.
    pattern = re.compile(
        r"^(.*denselayer\d+\.(?:norm|relu|conv))\.((?:[12])\.(?:weight|bias|running_mean|running_var))$"
    )

    state_dict = load_state_dict_from_url(
        model_url, map_location="cpu", progress=progress, check_hash=True
    )
    for key in list(state_dict.keys()):
        res = pattern.match(key)
        if res:
            new_key = res.group(1) + res.group(2)
            state_dict[new_key] = state_dict[key]
            del state_dict[key]
    model.load_state_dict(state_dict)


def _densenet(
    arch: str,
    growth_rate: int,
    block_config: Union[Tuple[int, int, int], Tuple[int, int, int, int]],
    num_init_features: int,
    pretrained: bool,
    progress: bool,
    **kwargs: Any,
) -> BcosDenseNet:
    model = BcosDenseNet(growth_rate, block_config, num_init_features, **kwargs)
    if pretrained:
        raise ValueError(
            "If you want to load pretrained weights, then please use the entrypoints in "
            "bcos.pretrained or bcos.model.pretrained instead."
        )
    return model


def densenet121(
    pretrained: bool = False,
    progress: bool = True,
    num_init_features=64,
    growth_rate=32,
    **kwargs: Any,
) -> BcosDenseNet:
    return _densenet(
        "densenet121",
        growth_rate,
        (6, 12, 24, 16),
        num_init_features,
        pretrained,
        progress,
        **kwargs,
    )


def densenet161(
    pretrained: bool = False, progress: bool = True, **kwargs: Any
) -> BcosDenseNet:
    return _densenet(
        "densenet161", 48, (6, 12, 36, 24), 96, pretrained, progress, **kwargs
    )


def densenet169(
    pretrained: bool = False, progress: bool = True, **kwargs: Any
) -> BcosDenseNet:
    return _densenet(
        "densenet169", 32, (6, 12, 32, 32), 64, pretrained, progress, **kwargs
    )


def densenet201(
    pretrained: bool = False, progress: bool = True, **kwargs: Any
) -> BcosDenseNet:
    return _densenet(
        "densenet201", 32, (6, 12, 48, 32), 64, pretrained, progress, **kwargs
    )


# cifar helpers
def _get_cifar_block_config(depth: int) -> Tuple[int, int, int]:
    if (depth - 4) % 3:
        raise Exception("Invalid depth!")
    return ((depth - 4) // 6,) * 3


def _update_default_cifar(kwargs: dict) -> None:
    kwargs.setdefault("num_classes", 10)
    kwargs.setdefault("small_inputs", True)


def cifar10_densenet52(
    pretrained: bool = False, progress: bool = True, **kwargs: Any
) -> BcosDenseNet:
    L = 52
    k = 12
    block_config = _get_cifar_block_config(L)
    _update_default_cifar(kwargs)
    return _densenet(
        f"cifar10_densenet{L}", k, block_config, k * 2, pretrained, progress, **kwargs
    )


# cifar densenet configs from Sec. 3#Implementation detail from the paper
def cifar10_densenet100(
    pretrained: bool = False, progress: bool = True, **kwargs: Any
) -> BcosDenseNet:
    L = 100
    k = 12
    block_config = _get_cifar_block_config(L)
    _update_default_cifar(kwargs)
    return _densenet(
        f"cifar10_densenet{L}", k, block_config, k * 2, pretrained, progress, **kwargs
    )


def cifar10_densenet250(
    pretrained: bool = False, progress: bool = True, **kwargs: Any
) -> BcosDenseNet:
    L = 250
    k = 24
    block_config = _get_cifar_block_config(L)
    _update_default_cifar(kwargs)
    return _densenet(
        f"cifar10_densenet{L}", k, block_config, k * 2, pretrained, progress, **kwargs
    )


def cifar10_densenet190(
    pretrained: bool = False, progress: bool = True, **kwargs: Any
) -> BcosDenseNet:
    L = 190
    k = 40
    block_config = _get_cifar_block_config(L)
    _update_default_cifar(kwargs)
    return _densenet(
        f"cifar10_densenet{L}", k, block_config, k * 2, pretrained, progress, **kwargs
    )
