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

class Model(nn.Module):
    """
    修正后的PyTorch ContrastiveLoss实现
    使用标准公式：0.5 * [y * d² + (1-y) * max(0, margin - d)²]
    其中 d = ||anchor - sample||₂ (实际欧氏距离)
    """
    def __init__(self, margin=2.0):
        super(Model, self).__init__()
        self.margin = margin
    
    def forward(self, anchor: torch.Tensor, sample: torch.Tensor, label: torch.Tensor) -> torch.Tensor:
        """
        使用PyTorch标准实现ContrastiveLoss
        
        Args:
            anchor (torch.Tensor): 锚点样本 [batch_size, feature_dim]
            sample (torch.Tensor): 对比样本 [batch_size, feature_dim] 
            label (torch.Tensor): 标签 [batch_size] (0=相似, 1=不相似)

        Returns:
            torch.Tensor: Contrastive Loss标量值
        """
        # 使用PyTorch高度优化的pairwise_distance计算实际欧氏距离
        distances = F.pairwise_distance(anchor, sample, p=2)
        
        # 标准ContrastiveLoss公式
        # y=1(不相似): 0.5 * max(0, margin - distance)²
        # y=0(相似): 0.5 * distance²
        losses = 0.5 * (label.float() * distances.pow(2) + 
                       (1 - label).float() * F.relu(self.margin - distances).pow(2))
        
        return torch.sum(losses)

batch_size = 128
feature_dim = 512
margin = 2.0

def get_inputs():
    """
    生成合理的测试数据
    """
    anchor = torch.randn(batch_size, feature_dim)
    sample = torch.randn(batch_size, feature_dim)
    label = torch.randint(0, 2, (batch_size,)).float()
    return [anchor, sample, label]

def get_init_inputs():
    return [margin]  # margin参数
