# -*- 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 List, Optional, 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.
    """

    def _is_exportable(module):
        return hasattr(module, "vanilla_export") and callable(
            getattr(module, "vanilla_export")
        )

    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 _is_exportable(module):
            # 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()
            if nmod.bias is not None:
                nmod.bias.data = module.bias.data.clone()
            else:
                assert module.bias is 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 SharedLipFactory:
    """
    Factory to share scaling factors between multiple layers.
    Register layers of type ScaledLipschitzModule
    that provide a get_scaling_factor method.
    Provide a method to get the product of all scaling factors.
    This assume that the network is sequential
    """

    def __init__(self):
        self.modules: List["ScaledLipschitzModule"] = []

    def register(self, module: "ScaledLipschitzModule"):
        self.modules.append(module)

    def get_current_product_value(self, training):
        """Retrieve the current product of the scaling factors."""
        if not self.modules:
            return torch.ones(())
        scalings = [m.get_scaling_factor(training=training) for m in self.modules]
        return torch.prod(torch.stack(scalings))


class ScaledLipschitzModule(abc.ABC):
    """
    This class allow to set learnable/fixed lipschitz parameter of a layer.
    args:
        scaling: whether the layer has a scaling factor or not.
        If not the scaling factor will return 1.0
        factory: Optional factory to share scaling factors between multiple layers.
    """

    def __init__(
        self, scaling: bool = False, factory: Optional[SharedLipFactory] = None
    ):
        assert (scaling) or (
            factory is None
        ), "Factory has to be none when scaling is False. "
        # Factory of factors
        self.factory = factory
        self.scaling = scaling
        if self.factory is not None:
            self.factory.register(self)

    """Retrieve the scaling_factor of the layer."""

    def get_scaling_factor(self, training: bool = False) -> torch.Tensor:
        if not self.scaling:
            return torch.ones((1,))
        return self.get_scaling(training=training)

    @abc.abstractmethod
    def get_scaling(self, training: bool = False) -> torch.Tensor:
        """Retrieve the scaling factor of the layer."""
        pass

    @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 _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(ScaledLipschitzModule):
    """
    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.
    args:
        coefficient_lip: multiplicative coefficient of the layer
        to form a K-Lipschitz layer.
        factory: from ScaledLipschitzModule,
    """

    # The target coefficient:
    _coefficient_lip: float

    def __init__(
        self,
        coefficient_lip: float = 1.0,
        factory: Optional[SharedLipFactory] = None,
    ):
        scaling = coefficient_lip != 1.0
        super().__init__(scaling=scaling, factory=factory)
        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)
        )

    def get_scaling(self, training=False, update=True):
        return self._coefficient_lip


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))
