# Copyright (c) Facebook, Inc. and its affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import math

import torch
from torch.nn.modules.loss import _Loss


class LatentLayersKLLoss(_Loss):
    def __init__(self, args):
        super().__init__()
        self.args = args

    def forward(self, layer_samples, lang_idx, update_num, sample_size):
        prior = self.args.prior
        samples = layer_samples[lang_idx]
        eps = 1e-7
        if prior == "uniform":
            # uniform prior
            kl_loss = (samples * (torch.log(samples + eps) - math.log(0.5))).sum(-1)
        elif prior == "agged_posterior":
            # aggregated posterior
            y_t = torch.stack([x.detach() for x in layer_samples], dim=0)
            agged_q = torch.sum(y_t, dim=0)
            row_norm = agged_q.sum(-1)
            normed_agg_q = agged_q / row_norm
            kl_loss = (
                samples * (torch.log(samples + eps) - torch.log(normed_agg_q + eps))
            ).sum(-1)
        else:
            raise NotImplementedError("The specified prior is not implemented.")

        # normalized by number of layers
        kl_loss /= layer_samples[0].size()[0]
        kl_weight = min(
            self.args.sparsity_weight,
            (update_num - self.args.soft_update)
            * self.args.sparsity_weight
            / self.args.anneal_updates,
        )
        kl_loss *= kl_weight * sample_size
        return kl_loss


class LatentLayersSparsityLoss(_Loss):
    def __init__(self, args):
        super().__init__()
        self.args = args

    def is_valid(self, update_num):
        if self.args.target_layers <= 0:
            return False
        return update_num > (self.args.soft_update + self.args.anneal_updates)

    def forward(self, layer_samples_list, update_num, sample_size):
        batch_loss = 0
        share_loss = 0
        global_sparsity_loss = 0
        layer_samples = torch.stack(layer_samples_list, dim=0)
        if (
            self.args.target_layers > 0 or self.args.share_weight > 0
        ) and update_num > (self.args.soft_update + self.args.anneal_updates):
            # anneal sparsity weight
            if update_num < (self.args.anneal_updates + self.args.soft_update):
                weight_anneal = 0
            elif update_num < (2 * self.args.anneal_updates + self.args.soft_update):
                weight_anneal = (
                    (update_num - self.args.soft_update - self.args.anneal_updates)
                    * self.args.share_weight
                    / self.args.anneal_updates
                )
            else:
                weight_anneal = 1
            # compute ratio among languages
            layer_utilization = torch.sum(layer_samples, dim=0)
            layer_utilization /= layer_samples.size()[0]
            if self.args.share_weight > 0:
                # encouraging sharing across languages
                share_loss = sum(
                    -1.0 * v * math.log(v) for v in layer_utilization if v > 0
                )
                batch_loss += (
                    weight_anneal * self.args.share_weight * sample_size * share_loss
                )
            if self.args.target_layers > 0:
                # computed expected number of layers selected
                expeted_layers = sum(layer_utilization)
                # compute l2 loss wrt target number of layers
                global_sparsity_loss = (expeted_layers - self.args.target_layers) ** 2
                batch_loss += (
                    weight_anneal
                    * self.args.share_weight
                    * sample_size
                    * global_sparsity_loss
                )
        return batch_loss