import torch
import torch.nn as nn

class Model(nn.Module):
    """
    Hinge Loss implementation - commonly used for maximum-margin classification (SVM).
    Computes hinge loss: max(0, margin - y_true * y_pred)
    """
    def __init__(self, margin=1.0):
        super(Model, self).__init__()
        self.margin = margin
    
    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
        """
        Compute hinge loss between input predictions and target labels.

        Args:
            input (torch.Tensor): Predicted values [batch_size, ...]
            target (torch.Tensor): Target labels [batch_size, ...] (should be -1 or 1)

        Returns:
            torch.Tensor: Scalar hinge loss value
        """
        # Hinge Loss: max(0, margin - y_true * y_pred)
        # 这里假设target是-1或1的标签
        hinge_loss = torch.clamp(self.margin - target * input, min=0)
        
        return torch.sum(hinge_loss)

batch_size = 256
feature_dim = 512
margin = 1.0

def get_inputs():
    # input: 预测值，可以是任意实数
    input = torch.randn(batch_size, feature_dim)
    # target: 标签，应该是-1或1，这里随机生成
    target = torch.randint(-1, 2, (batch_size, feature_dim)).float()
    # 确保没有0值，因为hinge loss通常用-1和1
    target[target == 0] = 1
    return [input, target]

def get_init_inputs():
    return [margin]  # margin parameter
