"""Normalization layers."""

import functools
import numbers
import typing

import torch
from torch import Tensor, nn
from torch.nn.modules.batchnorm import BatchNorm2d, SyncBatchNorm
from torchvision.ops.misc import FrozenBatchNorm2d as _FrozenBatchNorm2d

from unipercept.types import Device, DType, Size
from unipercept.utils.inspect import locate_object

# -------------------- #
# Register norm layers #
# -------------------- #

type NormFactory = typing.Callable[[int], nn.Module]
type NormSpec = str | NormFactory | type[nn.Module] | nn.Module | None


def get_default_bias(  # noqa: PLR0911
    value: bool | None, norm: nn.Module | None, *, default: bool = True
) -> bool:
    """Check whether a norm module makes the bias a transforming module obsolete.

    A bias parameter is only needed of the norm module has been explicitly set to
    not include a bias term.
    """
    # If the value is a bool, return that value
    if isinstance(value, bool):
        return value
    if norm is None:
        return default

    assert isinstance(norm, nn.Module), f"{norm=} ({type(norm)}) is not a module"

    if hasattr(norm, "default_bias"):
        return typing.cast(bool, norm.default_bias)
    if hasattr(norm, "affine"):
        return not norm.affine
    if hasattr(norm, "bias"):
        return not norm.bias
    if hasattr(norm, "elementwise_affine"):
        return not norm.elementwise_affine

    return default


def get_norm(spec: NormSpec, num_channels: int, **kwargs) -> nn.Module:
    """Resolve a norm module from a string, a factory function or a module instance.

    Parameters
    ----------
    norm
        A string, a factory function or an instance of a norm module.
    num_channels
        Number of channels to normalize.
    **kwargs
        Additional keyword arguments to pass to the norm factory.

    Notes
    -----
    When a module instance is provided,
    the ``num_channels`` argument is ignored and the object is returned as-is.
    """
    if spec is None:
        spec = nn.Identity()
    elif isinstance(spec, str):
        spec = locate_object(spec)

    # If already a module instance, return that instance directly
    if isinstance(spec, nn.Module):
        return spec
    if callable(spec):
        return spec(num_channels, **kwargs)
    msg = f"Norm module {spec!r} not found."
    raise ValueError(msg)


# --------------------- #
# LayerNorm and RMSNorm #
# --------------------- #

try:
    from apex.normalization.fused_layer_norm import (  # noqa: I001 # type: ignore[import]
        fused_layer_norm_affine,
        fused_layer_norm,
        fused_rms_norm_affine,
        fused_rms_norm,
    )

    has_apex = True
except ImportError:
    has_apex = False

LAYERNORM_DEFAULT_EPS = 1e-5
LAYERNORM_DEFAULT_MEMORY_EFFICIENT = False


def layer_norm(
    input: Tensor,
    shape: Size,
    weight: Tensor | None = None,
    bias: Tensor | None = None,
    eps: float = LAYERNORM_DEFAULT_EPS,
    memory_efficient: bool = LAYERNORM_DEFAULT_MEMORY_EFFICIENT,
) -> Tensor:
    """Layer Normalization."""
    if (
        torch.jit.is_tracing()
        or torch.jit.is_scripting()
        or torch.compiler.is_compiling()
        or not input.is_cuda
        or not has_apex
    ):
        if torch.is_autocast_enabled():
            dtype = torch.get_autocast_dtype(input.device.type)
            input = input.to(dtype)
            if weight is not None:
                weight = weight.to(dtype)
            if bias is not None:
                bias = bias.to(dtype)
        with torch.autocast(input.device.type, enabled=False):
            return nn.functional.layer_norm(input, shape, weight, bias, eps)
    if weight is not None:
        assert bias is not None
        return fused_layer_norm_affine(
            input,
            weight,
            bias,
            shape,
            eps,
            memory_efficient,
        )
    assert bias is None
    return fused_layer_norm(input, shape, eps, memory_efficient)


class LayerNorm(nn.Module):
    """Layer Normalization."""

    shape: typing.Final[tuple[int, ...]]
    eps: typing.Final[float]
    elementwise_affine: typing.Final[bool]
    memory_efficient: typing.Final[bool]

    def __init__(
        self,
        shape: typing.Iterable[int] | numbers.Integral | int,
        *,
        eps: float = LAYERNORM_DEFAULT_EPS,
        elementwise_affine: bool = True,
        device: Device = None,
        dtype: DType | None = None,
    ) -> None:
        factory_kwargs = {"device": device, "dtype": dtype}
        super().__init__()
        if isinstance(shape, numbers.Integral | int):
            self.shape = typing.cast(tuple[int], (shape,))
        else:
            self.shape = tuple(shape)
        self.eps = eps
        self.elementwise_affine = elementwise_affine
        self.memory_efficient = False

        if self.elementwise_affine:
            self.weight = nn.Parameter(torch.empty(self.shape, **factory_kwargs))
            self.bias = nn.Parameter(torch.empty(self.shape, **factory_kwargs))
        else:
            self.register_parameter("weight", None)
            self.register_parameter("bias", None)

        self.reset_parameters()

    def reset_parameters(self) -> None:
        if self.elementwise_affine:
            nn.init.ones_(self.weight)
            nn.init.zeros_(self.bias)

    @typing.override
    def forward(self, input: Tensor) -> Tensor:
        """Forward pass."""
        return layer_norm(
            input, self.shape, self.weight, self.bias, self.eps, self.memory_efficient
        )

    @typing.override
    def extra_repr(self) -> str:
        return "{shape}, eps={eps}, elementwise_affine={elementwise_affine}".format(
            **self.__dict__
        )


RMSNORM_DEFAULT_EPS = LAYERNORM_DEFAULT_EPS
RMSNORM_DEFAULT_MEMORY_EFFICIENT = LAYERNORM_DEFAULT_MEMORY_EFFICIENT


def rms_norm(
    input: Tensor,
    shape: Size,
    weight: Tensor | None = None,
    eps: float = RMSNORM_DEFAULT_EPS,
    memory_efficient: bool = RMSNORM_DEFAULT_MEMORY_EFFICIENT,
) -> Tensor:
    """Root Mean Square (RMS) Normalization."""
    if (
        torch.jit.is_tracing()
        or torch.jit.is_scripting()
        or torch.compiler.is_compiling()
        or not input.is_cuda
        or not has_apex
    ):
        if torch.is_autocast_enabled():
            dtype = torch.get_autocast_dtype(input.device.type)
            input = input.to(dtype)
            if weight is not None:
                weight = weight.to(dtype)
        with torch.autocast(input.device.type, enabled=False):
            return nn.functional.rms_norm(input, shape, weight, eps)
    if weight is not None:
        return fused_rms_norm_affine(
            input,
            weight,
            shape,
            eps,
            memory_efficient,
        )
    return fused_rms_norm(input, shape, eps, memory_efficient)


class RMSNorm(nn.Module):
    """Root Mean Square (RMS) Normalization."""

    shape: typing.Final[tuple[int, ...]]
    eps: typing.Final[float]
    elementwise_affine: typing.Final[bool]
    memory_efficient: typing.Final[bool]

    def __init__(
        self,
        shape: typing.Iterable[int] | numbers.Integral | int,
        *,
        eps: float = RMSNORM_DEFAULT_EPS,
        elementwise_affine: bool = True,
        device: Device = None,
        dtype: DType | None = None,
    ) -> None:
        factory_kwargs = {"device": device, "dtype": dtype}
        super().__init__()
        if isinstance(shape, numbers.Integral | int):
            self.shape = typing.cast(tuple[int], (shape,))
        else:
            self.shape = tuple(shape)
        self.eps = eps
        self.elementwise_affine = elementwise_affine
        self.memory_efficient = False
        if self.elementwise_affine:
            self.weight = nn.Parameter(torch.empty(self.shape, **factory_kwargs))
        else:
            self.register_parameter("weight", None)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        """Reset the weights."""
        if self.elementwise_affine:
            nn.init.ones_(self.weight)

    @typing.override
    def forward(self, input: Tensor) -> Tensor:
        """Forward pass."""
        return rms_norm(input, self.shape, self.weight, self.eps, self.memory_efficient)

    @typing.override
    def extra_repr(self) -> str:
        """Extra information about the module."""
        return "{shape}, eps={eps}, elementwise_affine={elementwise_affine}".format(
            **self.__dict__
        )


# ---------- #
# Group Norm #
# ---------- #


class GroupNorm(nn.Module):
    r"""Group Normalization."""

    num_groups: typing.Final[int]
    num_channels: typing.Final[int]
    eps: typing.Final[float]
    affine: typing.Final[bool]

    def __init__(
        self,
        num_groups: int,
        num_channels: int,
        *,
        eps: float = 1e-5,
        affine: bool = True,
        device: Device = None,
        dtype: DType | None = None,
    ) -> None:
        factory_kwargs = {"device": device, "dtype": dtype}
        super().__init__()
        if num_channels % num_groups != 0:
            msg = f"{num_channels=} must be divisible by {num_groups=}"
            raise ValueError(msg)

        self.num_groups = num_groups
        self.num_channels = num_channels
        self.eps = eps
        self.affine = affine
        if self.affine:
            self.weight = nn.Parameter(torch.empty(num_channels, **factory_kwargs))
            self.bias = nn.Parameter(torch.empty(num_channels, **factory_kwargs))
        else:
            self.register_parameter("weight", None)
            self.register_parameter("bias", None)

        self.reset_parameters()

    def reset_parameters(self) -> None:
        if self.affine:
            nn.init.ones_(self.weight)
            nn.init.zeros_(self.bias)

    @typing.override
    def forward(self, input: Tensor) -> Tensor:
        """Forward pass."""
        return nn.functional.group_norm(
            input, self.num_groups, self.weight, self.bias, self.eps
        )

    @typing.override
    def extra_repr(self) -> str:
        """Extra information about the module."""
        return "{num_groups}, {num_channels}, eps={eps}, affine={affine}".format(
            **self.__dict__
        )


def GroupNorm32(num_channels: int, **kwargs) -> nn.GroupNorm:
    """
    GroupNorm with the number of groups equal to 32, like in Detectron2.

    Notes
    -----
    The amount of channels' value (32) originates from the optimal value found in
    the GroupNorm paper.
    """
    return nn.GroupNorm(32, num_channels=num_channels, **kwargs)


def GroupNormCG(num_channels: int, **kwargs: typing.Any) -> nn.GroupNorm:
    """GroupNorm with the number of groups equal to the number of channels."""
    return nn.GroupNorm(num_channels, num_channels=num_channels, **kwargs)


def layer_norm_2d(input: Tensor, weight: Tensor, bias: Tensor, eps: float) -> Tensor:
    """Layer normalization for 2D (CHW) inputs."""
    return nn.functional.group_norm(input, 1, weight, bias, eps)


def LayerNorm2d(num_channels: int, **kwargs) -> nn.GroupNorm:
    """LayerNorm with the number of groups equal to the number of channels."""
    return nn.GroupNorm(num_groups=1, num_channels=num_channels, **kwargs)


def GroupNormFactory(
    *, num_groups: int, **kwargs
) -> typing.Callable[[int], nn.GroupNorm]:
    """GroupNorm with the number of groups equal to the number of channels."""
    return functools.partial(nn.GroupNorm, num_groups, **kwargs)


# -------------------- #
# Global Response Norm #
# -------------------- #


def global_response_norm(
    x: Tensor,
    spatial_dim: tuple[int, int],
    channel_dim: int,
    bias: Tensor,
    weight: Tensor,
    eps: float,
) -> Tensor:
    """See :class:`GlobalResponseNorm` for details."""
    x_i = x.type_as(weight)
    x_g = x_i.norm(p=2, dim=spatial_dim, keepdim=True)
    x_n = x_g / (x_g.mean(dim=channel_dim, keepdim=True) + eps)
    y = x_i + torch.addcmul(bias, weight, x_i * x_n)

    return y.type_as(x)


class GlobalResponseNorm(nn.Module):
    """Global Response Normalization based on ConvNeXt-V2.

    Paper: https://arxiv.org/abs/2301.00808
    """

    def __init__(self, channels, eps=1e-6, channels_last=True) -> None:
        super().__init__()
        self.eps = eps
        if channels_last:
            self.spatial_dim = (1, 2)
            self.channel_dim = -1
            self.wb_shape = (1, 1, 1, -1)
        else:
            self.spatial_dim = (2, 3)
            self.channel_dim = 1
            self.wb_shape = (1, -1, 1, 1)

        self.weight = nn.Parameter(torch.zeros(channels))
        self.bias = nn.Parameter(torch.zeros(channels))

    @typing.override
    def forward(self, x):
        return global_response_norm(
            x,
            self.spatial_dim,
            self.channel_dim,
            self.bias.view(self.wb_shape),
            self.weight.view(self.wb_shape),
            self.eps,
        )


class GlobalResponseNorm2d(GlobalResponseNorm):
    """Gobal response norm for CHW tensors."""

    def __init__(self, *args, **kwargs) -> None:
        super().__init__(*args, channels_last=False, **kwargs)


# ---------------- #
# Frozen BatchNorm #
# ---------------- #


class FrozenBatchNorm2d(_FrozenBatchNorm2d):
    """BatchNorm2d with frozen parameters."""

    @classmethod
    def convert_from[M: nn.Module](
        cls, module: M, recursive: bool = True
    ) -> M | typing.Self:
        """Convert BatchNorm2d or SyncBatchNorm modules to FrozenBatchNorm2d.

        See Also
        --------
        - Based on `SyncBatchNorm <https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/batchnorm.py>`_.
        """
        res = module
        if isinstance(module, BatchNorm2d | SyncBatchNorm):
            res = cls(module.num_features)
            if module.affine:
                res.weight.data = module.weight.data.clone().detach()
                res.bias.data = module.bias.data.clone().detach()
            assert module.running_mean is not None
            assert module.running_var is not None
            res.running_mean.data = module.running_mean.data
            res.running_var.data = module.running_var.data
            res.eps = module.eps
            assert module.num_batches_tracked is not None
            res.num_batches_tracked = module.num_batches_tracked
        elif recursive:
            for name, child in module.named_children():
                new_child = cls.convert_from(child, recursive=True)
                if new_child is not child:
                    res.add_module(name, new_child)
        else:
            msg = f"Cannot convert {module.__class__.__name__} to {cls.__name__}."
            raise ValueError(msg)

        del module
        return res
