# Copyright (c) Microsoft, Inc. 2020
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
# Author: penhe@microsoft.com
# Date: 01/25/2021
#

import torch
import torch.nn.functional as F

__all__ = ["PerturbationLayer", "AdversarialLearner", "hook_sift_layer"]


class PerturbationLayer(torch.nn.Module):
    def __init__(self, hidden_size, learning_rate=1e-4, init_perturbation=1e-2):
        super().__init__()
        self.learning_rate = learning_rate
        self.init_perturbation = init_perturbation
        self.delta = None
        self.LayerNorm = torch.nn.LayerNorm(hidden_size, 1e-7, elementwise_affine=False)
        self.adversarial_mode = False

    def adversarial_(self, adversarial=True):
        self.adversarial_mode = adversarial
        if not adversarial:
            self.delta = None

    def forward(self, input):
        if not self.adversarial_mode:
            self.input = self.LayerNorm(input)
            return self.input
        else:
            if self.delta is None:
                self.update_delta(requires_grad=True)
            return self.perturbated_input

    def update_delta(self, requires_grad=False):
        if not self.adversarial_mode:
            return True
        if self.delta is None:
            delta = torch.clamp(
                self.input.new(self.input.size()).normal_(0, self.init_perturbation).float(),
                -2 * self.init_perturbation,
                2 * self.init_perturbation,
            )
        else:
            grad = self.delta.grad
            self.delta.grad = None
            delta = self.delta
            norm = grad.norm()
            if torch.isnan(norm) or torch.isinf(norm):
                return False
            eps = self.learning_rate
            with torch.no_grad():
                delta = delta + eps * grad / (1e-6 + grad.abs().max(-1, keepdim=True)[0])
        self.delta = delta.float().detach().requires_grad_(requires_grad)
        self.perturbated_input = (self.input.to(delta).detach() + self.delta).to(self.input)
        return True


def hook_sift_layer(
    model,
    hidden_size,
    learning_rate=1e-4,
    init_perturbation=1e-2,
    target_module="embeddings.LayerNorm",
):
    """
    Hook the sift perturbation layer to and existing model. With this method, you can apply adversarial training
    without changing the existing model implementation.

    Params:
    `model`: The model instance to apply adversarial training
    `hidden_size`: The dimmension size of the perturbated embedding
    `learning_rate`: The learning rate to update the perturbation
    `init_perturbation`: The initial range of perturbation
    `target_module`: The module to apply perturbation. It can be the name of the sub-module of the model or the sub-module instance.
    The perturbation layer will be inserted before the sub-module.

    Outputs:
    The perturbation layers.

    """

    if isinstance(target_module, str):
        _modules = [k for n, k in model.named_modules() if target_module in n]
    else:
        assert isinstance(target_module, torch.nn.Module), f"{type(target_module)} is not an instance of torch.nn.Module"
        _modules = [target_module]
    adv_modules = []
    for m in _modules:
        adv = PerturbationLayer(hidden_size, learning_rate, init_perturbation)

        def adv_hook(module, inputs):
            return adv(inputs[0])

        for h in list(m._forward_pre_hooks.keys()):
            if m._forward_pre_hooks[h].__name__ == "adv_hook":
                del m._forward_pre_hooks[h]
        m.register_forward_pre_hook(adv_hook)
        adv_modules.append(adv)
    return adv_modules


class AdversarialLearner:
    """Adversarial Learner
    This class is the helper class for adversarial training.

    Params:
    `model`: The model instance to apply adversarial training
    `perturbation_modules`: The sub modules in the model that will generate perturbations. If it's `None`,
    the constructor will detect sub-modules of type `PerturbationLayer` in the model.

    Example usage:
    ```python
    # Create DeBERTa model
    adv_modules = hook_sift_layer(model, hidden_size=768)
    adv = AdversarialLearner(model, adv_modules)
    def logits_fn(model, *wargs, **kwargs):
    logits,_ = model(*wargs, **kwargs)
    return logits
    logits,loss = model(**data)

    loss = loss + adv.loss(logits, logits_fn, **data)
    # Other steps is the same as general training.

    ```

    """

    def __init__(self, model, adv_modules=None):
        if adv_modules is None:
            self.adv_modules = [m for m in model.modules() if isinstance(m, PerturbationLayer)]
        else:
            self.adv_modules = adv_modules
        self.parameters = [p for p in model.parameters()]
        self.model = model

    def loss(self, target, logits_fn, loss_fn="symmetric-kl", *wargs, **kwargs):
        """
        Calculate the adversarial loss based on the given logits fucntion and loss function.
        Inputs:
        `target`: the logits from original inputs.
        `logits_fn`: the function that produces logits based on perturbated inputs. E.g.,
        ```python
        def logits_fn(model, *wargs, **kwargs):
        logits = model(*wargs, **kwargs)
        return logits
        ```
        `loss_fn`: the function that caclulate the loss from perturbated logits and target logits.
        - If it's a string, it can be pre-built loss functions, i.e. kl, symmetric_kl, mse.
        - If it's a function, it will be called to calculate the loss, the signature of the function will be,
        ```python
        def loss_fn(source_logits, target_logits):
            # Calculate the loss
            return loss
        ```
        `*wargs`: the positional arguments that will be passed to the model
        `**kwargs`: the key-word arguments that will be passed to the model
        Outputs:
        The loss based on pertubated inputs.
        """
        self.prepare()
        if isinstance(loss_fn, str):
            loss_fn = perturbation_loss_fns[loss_fn]
        pert_logits = logits_fn(*wargs, **kwargs)
        pert_loss = loss_fn(pert_logits, target.detach()).sum()
        pert_loss.backward()
        for m in self.adv_modules:
            ok = m.update_delta(True)

        for r, p in zip(self.prev, self.parameters):
            p.requires_grad_(r)
        pert_logits = logits_fn(*wargs, **kwargs)
        pert_loss = symmetric_kl(pert_logits, target)

        self.cleanup()
        return pert_loss.mean()

    def prepare(self):
        self.prev = [p.requires_grad for p in self.parameters]
        for p in self.parameters:
            p.requires_grad_(False)
        for m in self.adv_modules:
            m.adversarial_(True)

    def cleanup(self):
        for r, p in zip(self.prev, self.parameters):
            p.requires_grad_(r)

        for m in self.adv_modules:
            m.adversarial_(False)


def symmetric_kl(logits, target):
    logit_stu = logits.view(-1, logits.size(-1)).float()
    logit_tea = target.view(-1, target.size(-1)).float()
    logprob_stu = F.log_softmax(logit_stu, -1)
    logprob_tea = F.log_softmax(logit_tea, -1)
    prob_tea = logprob_tea.exp().detach()
    prob_stu = logprob_stu.exp().detach()
    floss = (prob_tea * (-logprob_stu)).sum(-1)  # Cross Entropy
    bloss = (prob_stu * (-logprob_tea)).sum(-1)  # Cross Entropy
    loss = floss + bloss
    return loss


def kl(logits, target):
    logit_stu = logits.view(-1, logits.size(-1)).float()
    logit_tea = target.view(-1, target.size(-1)).float()
    logprob_stu = F.log_softmax(logit_stu, -1)
    logprob_tea = F.log_softmax(logit_tea.detach(), -1)
    prob_tea = logprob_tea.exp()
    loss = (prob_tea * (-logprob_stu)).sum(-1)  # Cross Entropy
    return loss


def mse(logits, target):
    logit_stu = logits.view(-1, logits.size(-1)).float()
    logit_tea = target.view(-1, target.size(-1)).float()
    return F.mse_loss(logit_stu.view(-1), logit_tea.view(-1))


perturbation_loss_fns = {"symmetric-kl": symmetric_kl, "kl": kl, "mse": mse}
