import torch
import torch.nn as nn
import torch.nn.functional as F

BATCH_SIZE = 2048
NUM_CLASSES = 4096
SHAPE = (BATCH_SIZE, NUM_CLASSES)

BETA = 0.8
REDUCTION = 'none'

class HardBootstrappingLoss(nn.Module):
    """
    Hard Bootstrapping Loss (Reed et al., 2015)
    L = beta * CE(y_true) + (1-beta) * CE(y_pred_max)
    """
    def __init__(self, beta=0.8, reduction='mean'):
        super(HardBootstrappingLoss, self).__init__()
        self.beta = beta
        self.reduction = reduction

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        # logits: (N, C)
        # targets: (N)
        ce_loss = F.cross_entropy(logits, targets, reduction='none')
        
        pred_labels = torch.argmax(logits, dim=1)
        
        ce_pred = F.cross_entropy(logits, pred_labels, reduction='none')
        
        loss = self.beta * ce_loss + (1.0 - self.beta) * ce_pred
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, beta=0.8, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = HardBootstrappingLoss(beta=beta, reduction=reduction)
    
    def forward(self, logits, targets):
        return self.loss_fn(logits, targets)

def get_inputs():
    logits = torch.randn(SHAPE, dtype=torch.float32)
    targets = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
    return [logits.contiguous(), targets.contiguous()]

def get_init_inputs():
    return [BETA, REDUCTION]