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

BATCH_SIZE = 256
FEATURE_DIM = 512
TEMPERATURE = 0.1
N_NEGATIVES = BATCH_SIZE * 10


class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.temperature = TEMPERATURE

    def forward(self, query: torch.Tensor, positive: torch.Tensor, negatives: torch.Tensor) -> torch.Tensor:
        # InfoNCE Loss实现
        # (B, D) vs (B, D) -> (B,)
        positive_sim = F.cosine_similarity(query, positive, dim=1) / self.temperature

        # (B, D) @ (D, N_NEG) -> (B, N_NEG)
        negative_sims = torch.matmul(query, negatives.t()) / self.temperature

        # 拼接: (B, 1) 和 (B, N_NEG) -> (B, 1 + N_NEG)
        logits = torch.cat([positive_sim.unsqueeze(1), negative_sims], dim=1)

        # 标签总是 0，因为正样本总是在索引 0
        labels = torch.zeros(query.size(0), dtype=torch.long, device=query.device)

        loss = F.cross_entropy(logits, labels)
        return loss


def get_inputs():
    query = F.normalize(torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32), p=2, dim=1)
    positive = F.normalize(torch.randn(BATCH_SIZE, FEATURE_DIM, dtype=torch.float32), p=2, dim=1)
    negatives = F.normalize(torch.randn(N_NEGATIVES, FEATURE_DIM, dtype=torch.float32), p=2, dim=1)
    return [query, positive, negatives]


def get_init_inputs():
    return []