import torch
import torch.nn as nn


def get_entropy_loss(opt):
    return EntropyLoss()


class EntropyLoss(nn.Module):
    def __init__(self):
        super().__init__()
        self.exp = 1e-7
        assert self.exp < 0.5

    def forward(self, item):
        item = item.clamp(min=self.exp, max=1 - self.exp)
        entropy = -item * torch.log(item) - (1 - item) * torch.log(1 - item)
        entropy = entropy.mean()

        return {"loss": entropy}