# -*- coding: utf-8 -*-
# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
# CRIAQ and ANITI - https://www.deel.ai/
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
# Copyright IRT Antoine de Saint Exupéry et Université Paul Sabatier Toulouse III - All
# rights reserved. DEEL is a research program operated by IVADO, IRT Saint Exupéry,
# CRIAQ and ANITI - https://www.deel.ai/
# =====================================================================================
"""
This module contains equivalents for Model and Sequential. These classes add support
for condensation and vanilla exportation.
"""

import abc
import copy
import warnings
import math
from collections import OrderedDict
from typing import Any

import numpy as np
import torch
import torch.nn as nn
import torch.nn.utils.parametrize as parametrize
from torch.nn import Sequential as TorchSequential


def _is_supported_1lip_layer(layer):
    """Return True if the layer is 1-Lipschitz. Note that in some cases, the layer
    is 1-Lipschitz for specific set of parameters.
    """
    supported_1lip_layers = (
        torch.nn.Softmax,
        torch.nn.Flatten,
        torch.nn.Identity,
        torch.nn.ReLU,
        torch.nn.Sigmoid,
        torch.nn.Tanh,
        torch.nn.Unflatten,
    )
    if isinstance(layer, supported_1lip_layers):
        return True
    elif isinstance(layer, torch.nn.MaxPool2d):
        return layer.kernel_size <= layer.stride
    return False


def _get_vanilla_module(module: nn.Module) -> nn.Module:
    """Return the vanilla torch module corresponding to a lipschitz module.

    Args:
        module (nn.Module): A lipschitz module

    Returns:
        nn.Module: The corresponding vanilla torch module
    """
    if isinstance(module, torch.nn.Conv2d):
        return torch.nn.Conv2d(
            in_channels=module.in_channels,
            out_channels=module.out_channels,
            kernel_size=module.kernel_size,
            stride=module.stride,
            padding=module.padding,
            dilation=module.dilation,
            groups=module.groups,
            bias=module.bias is not None,
            padding_mode=module.padding_mode,
        )
    elif isinstance(module, torch.nn.Linear):
        return torch.nn.Linear(
            in_features=module.in_features,
            out_features=module.out_features,
            bias=module.bias is not None,
        )
    elif isinstance(module, torch.nn.ConvTranspose2d):
        return torch.nn.ConvTranspose2d(
            in_channels=module.in_channels,
            out_channels=module.out_channels,
            kernel_size=module.kernel_size,
            stride=module.stride,
            padding=module.padding,
            output_padding=module.output_padding,
            dilation=module.dilation,
            groups=module.groups,
            bias=module.bias is not None,
            padding_mode=module.padding_mode,
        )
    return module


def vanilla_model(model: nn.Module, in_place=True, new_module=None) -> nn.Module:
    """Convert lipschitz modules into their non-lipschitz counterpart (for
    instance, SpectralConv2d layers become Conv2d layers).

    Warning: This function modifies the model in-place by default.
        Set `in_place=False` for a copy of the model.

    Args:
        model (nn.Module): Lipschitz neural network
        in_place (bool): If True, modify the model in place. If False, return a copy
        new_module (nn.Module): If provided, the function will fill this module
            instead of creating a new one. Only useful in recursion, should be None
            when calling the function.
    """
    if in_place is False and new_module is None:
        model.eval()
        new_module = copy.deepcopy(model)
        new_module.eval()
    if in_place is True:
        new_module = model

    for n, module in model.named_children():
        if isinstance(module, LipschitzModule):
            # torchlip modules
            setattr(new_module, n, module.vanilla_export())
        elif parametrize.is_parametrized(module) and isinstance(
            module, (nn.Conv2d, nn.Linear, nn.ConvTranspose2d)
        ):
            # compatibility with orthogonium modules
            nmod = _get_vanilla_module(module)
            nmod.weight.data = module.weight.data.clone()
            nmod.bias.data = (
                module.bias.data.clone() if module.bias is not None else None
            )
            setattr(new_module, n, nmod)
        elif len(list(module.children())) > 0:
            # compound module, go inside it
            _ = vanilla_model(
                module, in_place=in_place, new_module=new_module.__getattr__(n)
            )
    return new_module


class _LipschitzCoefMultiplication(nn.Module):
    """Parametrization module for lipschitz global coefficient multiplication."""

    def __init__(self, coef: float):
        super().__init__()
        self._coef = coef

    def forward(self, weight: torch.Tensor) -> torch.Tensor:
        return self._coef * weight


class LipschitzModule(abc.ABC):
    """
    This class allow to set lipschitz factor of a layer. Lipschitz layer must inherit
    this class to allow user to set the lipschitz factor.

    Warning:
         This class only regroup useful functions when developing new Lipschitz layers.
         But it does not ensure any property about the layer. This means that
         inheriting from this class won't ensure anything about the lipschitz constant.
    """

    # The target coefficient:
    _coefficient_lip: float

    def __init__(self, coefficient_lip: float = 1.0):
        self._coefficient_lip = coefficient_lip

    def apply_lipschitz_factor(self):
        """Multiply the layer weights by a lipschitz factor."""
        if self._coefficient_lip == 1.0:
            return
        parametrize.register_parametrization(
            self, "weight", _LipschitzCoefMultiplication(self._coefficient_lip)
        )

    @abc.abstractmethod
    def vanilla_export(self):
        """
        Convert this layer to a corresponding vanilla torch layer (when possible).

        Returns:
             A vanilla torch version of this layer.
        """
        pass


class Sequential(TorchSequential, LipschitzModule):
    def __init__(
        self,
        *args: Any,
        k_coef_lip: float = 1.0,
    ):
        """
        Equivalent of torch.Sequential but allow to set k-lip factor globally. Also
        support condensation and vanilla exportation.
        For now constant repartition is implemented (each layer
        get n_sqrt(k_lip_factor), where n is the number of layers)
        But in the future other repartition function may be implemented.

        Args:
            layers: list of layers to add to the model.
            name: name of the model, can be None
            k_coef_lip: the Lipschitz coefficient to ensure globally on the model.
        """
        TorchSequential.__init__(self, *args)
        LipschitzModule.__init__(self, k_coef_lip)

        # Force the Lipschitz coefficient:
        n_layers = np.sum(
            [isinstance(layer, LipschitzModule) for layer in self.children()]
        )
        for module in self.children():
            if isinstance(module, LipschitzModule):
                module._coefficient_lip = math.pow(k_coef_lip, 1 / n_layers)
            elif _is_supported_1lip_layer(module) is not True:
                warnings.warn(
                    "Sequential model contains a layer which is not a Lipschitz layer: {}".format(  # noqa: E501
                        module
                    )
                )

    def vanilla_export(self):
        """
        Exports this model to a vanilla torch Sequential.

        This method only works for flat sequential. Lipschitz modules are converted
        using their own `vanilla_export` method while non-Lipschitz modules are simply
        copied using `copy.deepcopy`.

        Returns:
            A Vanilla torch.nn.Sequential model.
        """
        layers = []
        for name, layer in self.named_children():
            if isinstance(layer, LipschitzModule):
                layers.append((name, layer.vanilla_export()))
            else:
                layers.append((name, copy.deepcopy(layer)))
        return TorchSequential(OrderedDict(layers))
