# Copyright (c) Aman Urumbekov and other contributors.
import math
import warnings
from typing import Any, Tuple, TypeVar

import torch
from torch import nn


# --- Tuple Utilities ---
T = TypeVar("T")


def _to_ntuple(n: int) -> Callable[[Any], Tuple[T, ...]]:
    """
    Creates a function that converts an input into a tuple of a specific size.
    """

    def parse(x: Any) -> Tuple[T, ...]:
        """
        Parses an input into a tuple of length n.

        - If x is a list/tuple of length n, it's returned as is.
        - If x is a list/tuple of length 1, its element is repeated n times.
        - If x is a scalar, it's repeated n times.
        """
        if isinstance(x, (list, tuple)):
            if len(x) == n:
                return tuple(x)
            if len(x) == 1:
                return (x[0],) * n
            raise ValueError(f"Expected a sequence of length 1 or {n}, but got {len(x)}")
        # Handle scalars (non-list/tuple types)
        return (x,) * n

    return parse


to_1tuple = _to_ntuple(1)
to_2tuple = _to_ntuple(2)
to_3tuple = _to_ntuple(3)
to_4tuple = _to_ntuple(4)


# --- Weight Initialization ---


def _no_grad_trunc_normal_(tensor: torch.Tensor, mean: float, std: float, a: float, b: float) -> torch.Tensor:
    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
    def norm_cdf(x: float) -> float:
        # Computes standard normal cumulative distribution function
        return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0

    if (mean < a - 2 * std) or (mean > b + 2 * std):
        warnings.warn(
            "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
            "The distribution of values may be incorrect.",
            stacklevel=2,
        )

    with torch.no_grad():
        # Values are generated by using a truncated uniform distribution and
        # then using the inverse CDF for the normal distribution.
        low = norm_cdf((a - mean) / std)
        up = norm_cdf((b - mean) / std)

        # Uniformly fill tensor with values from [low, up], then translate to
        # [2*low-1, 2*up-1].
        tensor.uniform_(2 * low - 1, 2 * up - 1)

        # Use inverse cdf transform for normal distribution to get truncated
        # standard normal
        tensor.erfinv_()

        # Transform to proper mean, std
        tensor.mul_(std * math.sqrt(2.0))
        tensor.add_(mean)

        # Clamp to ensure it's in the proper range
        tensor.clamp_(min=a, max=b)
        return tensor


def trunc_normal_(
    tensor: torch.Tensor, mean: float = 0.0, std: float = 1.0, a: float = -2.0, b: float = 2.0
) -> torch.Tensor:
    r"""Fills the input Tensor with values drawn from a truncated
    normal distribution.

    The values are effectively drawn from the
    normal distribution :math:`\mathcal{N}(\text{mean}, \text{std}^2)`
    with values outside :math:`[a, b]` redrawn until they are within
    the bounds. The method used for generating the random values works
    best when :math:`a \leq \text{mean} \leq b`.

    Args:
        tensor: an n-dimensional `torch.Tensor`.
        mean: the mean of the normal distribution.
        std: the standard deviation of the normal distribution.
        a: the minimum cutoff value.
        b: the maximum cutoff value.

    Returns:
        The initialized tensor.
    """
    return _no_grad_trunc_normal_(tensor, mean, std, a, b)


@torch.no_grad()
def default_init_weights(module_list, scale=1, bias_fill=0, **kwargs):
    """Initialize network weights.

    Args:
        module_list (list[nn.Module] | nn.Module): Modules to be initialized.
        scale (float): Scale initialized weights, especially for residual
            blocks. Default: 1.
        bias_fill (float): The value to fill bias. Default: 0
        kwargs (dict): Other arguments for initialization function.
    """
    if not isinstance(module_list, list):
        module_list = [module_list]
    for module in module_list:
        for m in module.modules():
            if isinstance(m, nn.Conv2d):
                init.kaiming_normal_(m.weight, **kwargs)
                m.weight.data *= scale
                if m.bias is not None:
                    m.bias.data.fill_(bias_fill)
            elif isinstance(m, nn.Linear):
                init.kaiming_normal_(m.weight, **kwargs)
                m.weight.data *= scale
                if m.bias is not None:
                    m.bias.data.fill_(bias_fill)
            elif isinstance(m, _BatchNorm):
                init.constant_(m.weight, 1)
                if m.bias is not None:
                    m.bias.data.fill_(bias_fill)
