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


class Model(nn.Module):
    def __init__(self, num_bins=10):
        super(Model, self).__init__()
        self.num_bins = num_bins
        self.max_dist = 4.0

    def forward(self, x: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
        x = F.normalize(x, p=2, dim=1)

        dist = torch.cdist(x, x, p=2).pow(2)

        width = self.max_dist / self.num_bins
        centers = torch.linspace(width / 2, self.max_dist - width / 2, self.num_bins, device=x.device)

        d = dist.unsqueeze(-1)
        c = centers.view(1, 1, -1)

        bin_weights = torch.relu(1.0 - torch.abs(d - c) / width)

        eq_labels = labels.unsqueeze(1) == labels.unsqueeze(0)
        diag_mask = torch.eye(x.size(0), device=x.device, dtype=torch.bool)

        pos_mask = eq_labels & (~diag_mask)
        neg_mask = (~eq_labels) & (~diag_mask)

        pos_hist = (bin_weights * pos_mask.unsqueeze(-1)).sum(dim=1)
        neg_hist = (bin_weights * neg_mask.unsqueeze(-1)).sum(dim=1)

        pos_cdf = torch.cumsum(pos_hist, dim=1)
        neg_cdf = torch.cumsum(neg_hist, dim=1)

        precision = pos_cdf / (pos_cdf + neg_cdf + 1e-10)

        total_pos = pos_cdf[:, -1].unsqueeze(1) + 1e-10
        delta_recall = pos_hist / total_pos

        ap = (precision * delta_recall).sum(dim=1)

        return 1.0 - ap.mean()


batch_size = 128
input_dim = 1024
num_classes = 32


def get_inputs():
    x = torch.randn(batch_size, input_dim)
    labels = torch.randint(0, num_classes, (batch_size,))
    return [x, labels]


def get_init_inputs():
    return []