from typing import Dict, List, Optional, Set, Tuple, Type, Union, cast

import torch
from torch import nn

from torch_blue import vi

from ..base import VIModule
from ..distributions import MeanFieldNormal
from .common_types import VIkwargs, _dist_any_t
from .init import fixed_

__all__ = [
    "convert_norms",
    "convert_to_vimodule",
    "ban_convert",
]

# Modules in the banlist will not be converted. The default set contains modules that
# do not have Parameters and therefore would not change when converted. Additionally,
# norms are assumed to be not converted by default.
_banlist = {
    nn.Identity,
    # activations
    nn.Threshold,
    nn.ReLU,
    nn.RReLU,
    nn.Hardtanh,
    nn.ReLU6,
    nn.Sigmoid,
    nn.Hardsigmoid,
    nn.Tanh,
    nn.SiLU,
    nn.Mish,
    nn.Hardswish,
    nn.ELU,
    nn.CELU,
    nn.SELU,
    nn.GLU,
    nn.GELU,
    nn.Hardshrink,
    nn.LeakyReLU,
    nn.LogSigmoid,
    nn.Softplus,
    nn.Softshrink,
    nn.Softsign,
    nn.Tanhshrink,
    nn.Softmin,
    nn.Softmax,
    nn.Softmax2d,
    nn.LogSoftmax,
    # distances
    nn.PairwiseDistance,
    nn.CosineSimilarity,
    # dropout
    nn.Dropout,
    nn.Dropout1d,
    nn.Dropout2d,
    nn.Dropout3d,
    nn.AlphaDropout,
    nn.FeatureAlphaDropout,
    # flatten,
    nn.Flatten,
    nn.Unflatten,
    # fold
    nn.Fold,
    nn.Unfold,
    # padding
    nn.CircularPad1d,
    nn.CircularPad2d,
    nn.CircularPad3d,
    nn.ConstantPad1d,
    nn.ConstantPad2d,
    nn.ConstantPad3d,
    nn.ReflectionPad1d,
    nn.ReflectionPad2d,
    nn.ReflectionPad3d,
    nn.ReplicationPad1d,
    nn.ReplicationPad2d,
    nn.ReplicationPad3d,
    nn.ZeroPad1d,
    nn.ZeroPad2d,
    nn.ZeroPad3d,
    # pixel shuffle
    nn.PixelShuffle,
    nn.PixelUnshuffle,
    # pooling
    nn.MaxPool1d,
    nn.MaxPool2d,
    nn.MaxPool3d,
    nn.MaxUnpool1d,
    nn.MaxUnpool2d,
    nn.MaxUnpool3d,
    nn.AvgPool1d,
    nn.AvgPool2d,
    nn.AvgPool3d,
    nn.FractionalMaxPool2d,
    nn.FractionalMaxPool3d,
    nn.LPPool1d,
    nn.LPPool2d,
    nn.LPPool3d,
    nn.AdaptiveMaxPool1d,
    nn.AdaptiveMaxPool2d,
    nn.AdaptiveMaxPool3d,
    nn.AdaptiveAvgPool1d,
    nn.AdaptiveAvgPool2d,
    nn.AdaptiveAvgPool3d,
    # upsampling
    nn.Upsample,
    nn.UpsamplingNearest2d,
    nn.UpsamplingBilinear2d,
}
_torch_norms = {
    # norms
    nn.BatchNorm1d,
    nn.LazyBatchNorm1d,
    nn.BatchNorm2d,
    nn.LazyBatchNorm2d,
    nn.BatchNorm3d,
    nn.LazyBatchNorm3d,
    nn.SyncBatchNorm,
    nn.InstanceNorm1d,
    nn.LazyInstanceNorm1d,
    nn.InstanceNorm2d,
    nn.LazyInstanceNorm2d,
    nn.InstanceNorm3d,
    nn.LazyInstanceNorm3d,
    nn.LocalResponseNorm,
    nn.CrossMapLRN2d,
    nn.LayerNorm,
    nn.GroupNorm,
    nn.RMSNorm,
}
_banlist |= _torch_norms
_replace_banlist: Set[nn.Module] = set()
_reuse_banlist: Set[nn.Module] = set()
_submodule_banlist: Set[nn.Module] = set()


def convert_norms(mode: bool = True) -> None:
    """
    Set whether to auto-convert PyTorch norm modules.

    Since norms in neural networks are mostly for stability, we assume by default, that
    they should remain non-Bayesian. With `mode=True` this method makes it so norms from
    PyTorch are converted. With `mode=False` it resets to the default behavior.

    Parameters
    ----------
    mode: bool, default=True
        If `True` set auto-conversion to convert PyTorch norms. If `False` reset to
        default behavior of not converting PyTorch norms.

    Returns
    -------
    None

    """
    global _banlist
    if mode:
        _banlist.difference_update(_torch_norms)
    else:
        _banlist.update(_torch_norms)


def ban_convert(
    class_names: Union[Type[nn.Module], List[Type[nn.Module]]],
    ban_mode: str = "ban",
    unban: bool = False,
) -> None:
    """
    Ban given class or list of classes from conversions.

    This method adds one or more class names to one of the conversion banlists. There
    are several different ban modes as defined by the `ban_mode` argument.

    The default banlist contains a variety to PyTorch modules that do not have
    parameters and therefore would not change during conversion. Additionally, by
    default PyTorch normalization layers are not converted, since their purpose is only
    stability not learning. If you wish to change this behavior use
    :func:`~convert_norms` to add or remove all PyTorch norms from the banlist.

    Parameters
    ----------
    class_names: Union[nn.Module, List[nn.Module]]
        A module class or a list of module classes to add to a banlist.
    ban_mode: str, default="ban"
        "ban": The module is not changed by auto-conversion, but submodules are still
        converted.
        "submodule": Both the module and all its submodules are not converted.
        "reuse": The module is converted, but if it occurs multiple times the
        auto-generated class will not be reused.
        "replace": The module will be converted and an auto-generated class is created
        even if `torch_blue` implements a module with the same name.
    unban: bool, default=False
        If `True` this will unban the provided classes instead.
    """
    if ban_mode == "ban":
        global _banlist
        banlist = _banlist
    elif ban_mode == "reuse":
        global _reuse_banlist
        banlist = _reuse_banlist
    elif ban_mode == "replace":
        global _replace_banlist
        banlist = _replace_banlist
    elif ban_mode == "submodule":
        global _submodule_banlist
        banlist = _submodule_banlist
    else:
        raise ValueError(f"Unknown ban mode: {ban_mode}")

    if isinstance(class_names, type):
        class_names = [class_names]

    if unban:
        banlist.difference_update(class_names)
    else:
        banlist.update(class_names)


def _convert_module(
    module: nn.Module,
    variational_distribution: _dist_any_t = MeanFieldNormal(),
    prior: _dist_any_t = MeanFieldNormal(),
    rescale_prior: bool = False,
    kaiming_initialization: bool = True,
    prior_initialization: bool = False,
    return_log_probs: bool = True,
    keep_weights: bool = False,
) -> None:
    if module.__class__ in _banlist or isinstance(module, VIModule):
        return

    vikwargs: VIkwargs = dict(
        variational_distribution=variational_distribution,
        prior=prior,
        rescale_prior=rescale_prior,
        kaiming_initialization=kaiming_initialization,
        prior_initialization=prior_initialization,
        return_log_probs=return_log_probs,
        device=None,
        dtype=None,
    )
    module_class = module.__class__
    class_name = module_class.__name__
    if (module_class not in _replace_banlist) and hasattr(vi, "VI" + class_name):
        new_class = getattr(vi, "VI" + class_name)
    elif (module_class not in _reuse_banlist) and "AVI" + class_name in globals():
        new_class = globals()["AVI" + class_name]
    else:
        new_class_name = "AVI" + class_name
        new_class = type(new_class_name, (VIModule, module.__class__), dict())
        setattr(new_class, "forward", module.__class__.forward)
        globals()[new_class_name] = new_class

    module.__class__ = new_class
    module = cast(VIModule, module)
    if len(module._parameters) == 0:
        VIModule.__post_init__(module)
        return

    parameters = module._parameters
    module._parameters = dict()

    variable_shapes: Dict[str, Optional[Tuple[int, ...]]] = dict()
    variable_types: Dict[str, Tuple[torch.device, torch.dtype]] = dict()
    for name, parameter in parameters.items():
        if parameter is None:
            variable_shapes[name] = None
        else:
            variable_shapes[name] = tuple(parameter.shape)
            variable_types[name] = (parameter.device, parameter.dtype)

    devices = {spec[0] for spec in variable_types.values()}
    dtypes = {spec[1] for spec in variable_types.values()}

    types_set = False
    if len(devices) == 1 and len(dtypes) == 1:
        vikwargs["device"] = devices.pop()
        vikwargs["dtype"] = dtypes.pop()
        types_set = True

    VIModule.__init__(module, variable_shapes, convert_overwrite=True, **vikwargs)
    module.__class__.__post_init__(module)

    if not types_set:
        for var, (device, dtype) in variable_types.items():
            var_dist = module.variational_distribution[var]
            for param in var_dist.distribution_parameters:
                param_name = module.variational_parameter_name(var, param)
                param = getattr(module, param_name)
                param.data = param.to(device=device, dtype=dtype)

    if keep_weights:
        for name, parameter in parameters.items():
            if parameter is None:
                continue
            var_dist = module.variational_distribution[name]
            primary_parameter = var_dist.primary_parameter
            param_name = module.variational_parameter_name(name, primary_parameter)
            fixed_(getattr(module, param_name), parameter)


def convert_to_vimodule(
    module: nn.Module,
    keep_weights: bool = False,
    variational_distribution: _dist_any_t = MeanFieldNormal(),
    prior: _dist_any_t = MeanFieldNormal(),
    rescale_prior: bool = False,
    kaiming_initialization: bool = True,
    prior_initialization: bool = False,
    return_log_probs: bool = True,
) -> None:
    """
    Convert a PyTorch module to a VIModule.

    This method automatically converts a PyTorch module to a VIModule. This also works
    for any model that is compatible with :meth:`torch.vmap` (documentation
    `here <https://docs.pytorch.org/docs/stable/generated/torch.vmap.html>`__). Usually,
    this will be the case if you do not use the `+=`, `-=`, `*=`, and `/=` operators
    (their long form, i.e. `a = a + b` instead of `a += b` does not cause issues).

    To configure the model the usual :class:`~torch_blue.vi.VIkwargs` can be used,
    except `device` and `dtype` which are copied independently for each weight matrix
    from the original model. This allows to convert a distributed model and maintain
    the distributed structure. This is compatible with `torch.ddp`, but DDP needs to be
    applied after the conversion.

    For Bayesian pretraining or custom initialization schemes `keep_weights` can be set
    to `True`, which will maintain the original weights as value for the primary
    parameter of the weight distribution. While this can vary the primary parameter is
    typically the distribution mean.

    The model is converted inplace so to continue using the original model a copy needs
    to be made before conversion.

    A good method to verify correct model conversion is to set variational distribution
    and prior to :class:`~torch_blue.vi.distributions.NonBayesian` and `keep_weights`
    to `True`. This should make the converted model produce the same results as the
    original. (The output will have an additional sampling dimension and multiple copies
    of the same result unless you pass `samples=1` to the forward call)

    Standard PyTorch layers will automatically be converted to the optimized
    implementation of this library. Therefore, you should try to use class names that
    already exist in PyTorch (like "Transformer"). If you cannot do this, you can use
    :meth:`~.ban_torch_convert` to add classes to a banlist (or later remove them from
    it) that will make them be converted normally.

    Advanced note: Auto-conversion will ignore many standard modules from PyTorch since
    they do not have weights (making conversion irrelevant) or should not be converted
    (this mostly applies to norm layers).

    While not recommended you can use :meth:`~.convert_norms` to enable or disable
    conversion of all PyTorch norms.

    Furthermore, you can add (or remove) any class to the conversion ban list with
    :meth:`~.ban_convert`. By default, it will add the class to the global banlist so
    all instances of this class are left non-Bayesian. There are several more specific
    behaviors that can be disabled for specific classes with this method using the
    keyword argument `ban_mode`. Note, that it only affects the specific class not any
    subclasses, so if you wish to stop conversion of only one specific layer you can
    implement it as a subclass with different name but all methods unchanged.

    Auto-conversion will try to reuse auto-converted classes. This means if you
    implemented a custom layer type and used it multiple times all instances will still
    be instances of the same (converted) class. To disable this behavior you can use the
    keyword argument `ban_mode=reuse`.

    If there is a `torch_blue` implementation of a PyTorch layer auto-convert will use
    that instead of an automatically created class. Note, that this will cause problems
    if you name custom classes the same as PyTorch classes. To disable this behavior you
    can ban the replacement of the custom class with `ban_mode=replace`.

    Finally, you can stop conversion of a module and all of its submodules with
    `ban_mode=submodule`.

    Parameters
    ----------
    module: nn.Module
        The Pytorch module to convert.
    keep_weights: bool
        If `True` keep the original weights as value for the primary parameter.
    VIkwargs
        Several standard keyword arguments. See :class:`~.VIkwargs` for details.
        `device` and `dtype` cannot be used since they are copied from the original
        weights.
    """
    # device and dtype are not present because they are copied on a per-parameter basis
    # from the original modules
    vikwargs = dict(
        variational_distribution=variational_distribution,
        prior=prior,
        rescale_prior=rescale_prior,
        kaiming_initialization=kaiming_initialization,
        prior_initialization=prior_initialization,
        return_log_probs=return_log_probs,
    )
    # for m in module.modules():
    #    _convert_module(m, **vikwargs, keep_weights=keep_weights)
    global _submodule_banlist
    if module.__class__ in _submodule_banlist:
        return

    for m in module.children():
        convert_to_vimodule(m, keep_weights=keep_weights, **vikwargs)

    _convert_module(module, keep_weights=keep_weights, **vikwargs)
