import torch.nn as nn
import torch
from torch.autograd import Variable


class SRLHLoss (nn.Module):
    '''
    except Tr(BLF)
    '''

    def __init__(self, sparsity, code_length):
        super (SRLHLoss, self).__init__ ()
        self.sparsity = sparsity
        self.code_length = code_length

    def forward(self, F_batch, actions, rewards):
        batch_size = B_batch.size (1)
        # mean_F = F_batch.mean(1).view(self.code_length,1)
        ee = 1e-7
        reg_loss = (B_batch - F_batch) ** 2
        FF = F_batch.mm (F_batch.t ())
        F_batch01 = (F_batch + 1) / 2
        L1_loss = - (F_batch01 * torch.log (F_batch01 + ee) + (1 - F_batch01) * torch.log (1 - F_batch01 + ee))

        mean_F = F_batch01.mean (1).view (self.code_length, 1)

        bla_loss = (mean_F * torch.log (mean_F + ee) + (1 - mean_F) * torch.log (1 - mean_F + ee))

        oth_loss = self.lambda_1 * ((FF ** 2).sum () - 2 * batch_size * torch.trace (FF))

        loss = 0.5 * (reg_loss.sum ()) / (batch_size * self.code_length) + self.lambda_3 * L1_loss.sum () / (
                    batch_size * self.code_length) + self.lambda_2 * bla_loss.sum () / self.code_length + 0.25 * oth_loss / (
                           batch_size * self.code_length)
        return loss

