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

class Model(nn.Module):
    """
    CosineDistance + Softmax加权融合实现。
    计算余弦相似度，应用Softmax加权，然后计算加权余弦距离。
    """
    def __init__(self, eps=1e-8, temperature=1.0):
        super(Model, self).__init__()
        self.eps = eps
        self.temperature = temperature
    
    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        """
        Compute CosineDistance + Softmax weighted fusion.

        Args:
            x (torch.Tensor): First set of vectors [batch_size, feature_dim]
            y (torch.Tensor): Second set of vectors [batch_size, feature_dim]

        Returns:
            torch.Tensor: Weighted cosine distances [batch_size]
        """
        # 计算余弦相似度
        cosine_sim = F.cosine_similarity(x, y, dim=1, eps=self.eps)
        
        # 应用Softmax加权（带温度参数）
        softmax_weights = F.softmax(cosine_sim / self.temperature, dim=0)
        
        # 计算余弦距离
        cosine_dist = 1.0 - cosine_sim
        
        # 计算加权距离
        weighted_distances = softmax_weights * cosine_dist
        
        return weighted_distances

batch_size = 256
feature_dim = 512

def get_inputs():
    # Generate two sets of vectors
    x = torch.randn(batch_size, feature_dim)
    y = torch.randn(batch_size, feature_dim)
    return [x, y]

def get_init_inputs():
    return [1e-8, 1.0]  # eps and temperature values
