import torch
from torch import nn


class TimeAwareLoss(nn.Module):
    def __init__(self, decay_rate=0.1, reward_factor=0.1):
        super(TimeAwareLoss, self).__init__()
        self.bce = nn.BCELoss(reduction='none')
        self.decay_rate = decay_rate
        self.reward_factor = reward_factor

    def forward(self, outcome_pred, outcome_true, los_true):
        los_weights = torch.exp(-self.decay_rate * los_true)  # Exponential decay
        loss_unreduced = self.bce(outcome_pred, outcome_true)

        reward_term = (los_true * torch.abs(outcome_true - outcome_pred)).mean()  # Reward term
        loss = (loss_unreduced * los_weights).mean()-self.reward_factor * reward_term  # Weighted loss
        
        return torch.clamp(loss, min=0)

def get_time_aware_loss(outcome_pred, outcome_true, los_true):
    time_aware_loss = TimeAwareLoss()
    return time_aware_loss(outcome_pred, outcome_true, los_true)

if __name__ == "__main__":
    outcome_pred = torch.tensor([0.1])
    outcome_true = torch.tensor([1.])
    los_true = torch.tensor([-4.0])
    print(get_time_aware_loss(outcome_pred, outcome_true, los_true))