import math
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

from zyplib.nn.modules import BaseModule


class MetricLearningBase(BaseModule, ABC):
    """度量学习基类

    提供度量学习算法的基本接口和通用功能。

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络，用于提取特征
    - `embedding_dim` : `int`
        - 嵌入向量维度
    - `margin` : `float`, optional
        - 边界参数, by default 0.2
    """

    def __init__(self, encoder: nn.Module, embedding_dim: int, margin: float = 0.2):
        super().__init__()
        self.encoder = encoder
        self.embedding_dim = embedding_dim
        self.margin = margin

    def forward(self, x: Tensor) -> Dict[str, Tensor]:
        """前向传播

        Parameters
        ----------
        - `x` : `Tensor`
            - 输入数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含嵌入向量的字典
        """
        embeddings = self.encoder(x)
        embeddings = F.normalize(embeddings, p=2, dim=1)

        return {'embeddings': embeddings}

    @abstractmethod
    def compute_distance(self, embeddings1: Tensor, embeddings2: Tensor) -> Tensor:
        """计算距离

        Parameters
        ----------
        - `embeddings1` : `Tensor`
            - 第一组嵌入向量
        - `embeddings2` : `Tensor`
            - 第二组嵌入向量

        Returns
        ----------
        - `Tensor`
            - 距离矩阵
        """
        pass

    @abstractmethod
    def mine_pairs(self, embeddings: Tensor, labels: Tensor) -> Tuple[Tensor, Tensor]:
        """挖掘正负样本对

        Parameters
        ----------
        - `embeddings` : `Tensor`
            - 嵌入向量
        - `labels` : `Tensor`
            - 标签

        Returns
        ----------
        - `Tuple[Tensor, Tensor]`
            - 正样本对和负样本对的索引
        """
        pass

    @abstractmethod
    def compute_loss(self, batch_data: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """计算损失函数

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 批次数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含各种损失项的字典
        """
        pass


class TripletLoss(MetricLearningBase):
    """三元组损失

    实现三元组损失函数，通过最小化锚点与正样本的距离，
    最大化锚点与负样本的距离来学习嵌入。

    References
    ----------
    Schroff, F., Kalenichenko, D., & Philbin, J. (2015).
    FaceNet: A unified embedding for face recognition and clustering.
    In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 815-823).

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `embedding_dim` : `int`
        - 嵌入向量维度
    - `margin` : `float`, optional
        - 三元组损失边界, by default 0.2
    - `mining_strategy` : `str`, optional
        - 难样本挖掘策略 ('hard', 'semi_hard', 'all'), by default 'hard'
    - `distance_metric` : `str`, optional
        - 距离度量 ('euclidean', 'cosine'), by default 'euclidean'
    - `max_triplets` : `int`, optional
        - 最大三元组数量（用于内存保护）, by default 10000

    Examples
    ----------
    ```python
    # 创建编码器
    encoder = nn.Sequential(
        nn.Linear(784, 256),
        nn.ReLU(),
        nn.Linear(256, 128),
        nn.ReLU(),
        nn.Linear(128, 64)
    )

    # 创建三元组损失模型
    model = TripletLoss(
        encoder=encoder,
        embedding_dim=64,
        margin=0.2,
        max_triplets=5000  # 限制内存使用
    )

    # 训练数据
    batch_data = {
        'data': torch.randn(32, 784),
        'labels': torch.randint(0, 10, (32,))
    }

    loss_dict = model.compute_loss(batch_data)
    ```
    """

    def __init__(
        self,
        encoder: nn.Module,
        embedding_dim: int,
        margin: float = 0.2,
        mining_strategy: str = 'hard',
        distance_metric: str = 'euclidean',
        max_triplets: int = 10000,
    ):
        super().__init__(encoder, embedding_dim, margin)
        self.mining_strategy = mining_strategy
        self.distance_metric = distance_metric
        self.max_triplets = max_triplets

    def compute_distance(self, embeddings1: Tensor, embeddings2: Tensor) -> Tensor:
        """计算距离矩阵

        Parameters
        ----------
        - `embeddings1` : `Tensor`
            - 第一组嵌入向量，形状为 (N, D)
        - `embeddings2` : `Tensor`
            - 第二组嵌入向量，形状为 (M, D)

        Returns
        ----------
        - `Tensor`
            - 距离矩阵，形状为 (N, M)
        """
        if self.distance_metric == 'euclidean':
            # 使用 torch.cdist 进行高效的欧氏距离计算
            return torch.cdist(embeddings1, embeddings2, p=2)
        elif self.distance_metric == 'cosine':
            # 优化的余弦距离计算，确保数值稳定性
            # 先进行 L2 归一化，然后计算余弦相似度
            norm1 = F.normalize(embeddings1, p=2, dim=1)
            norm2 = F.normalize(embeddings2, p=2, dim=1)

            # 余弦相似度
            similarity = torch.mm(norm1, norm2.t())

            # 裁剪到 [-1, 1] 范围以确保数值稳定性
            similarity = torch.clamp(similarity, -1.0, 1.0)

            # 余弦距离 = 1 - 余弦相似度
            return 1 - similarity
        else:
            raise ValueError(f'Unsupported distance metric: {self.distance_metric}')

    def mine_pairs(
        self, embeddings: Tensor, labels: Tensor
    ) -> Tuple[Tensor, Tensor, Tensor]:
        """挖掘三元组

        Parameters
        ----------
        - `embeddings` : `Tensor`
            - 嵌入向量
        - `labels` : `Tensor`
            - 标签

        Returns
        ----------
        - `Tuple[Tensor, Tensor, Tensor]`
            - 锚点、正样本、负样本的索引
        """
        device = embeddings.device
        n = len(labels)

        if n < 3:
            # 样本数量不足，返回空张量
            empty_tensor = torch.tensor([], dtype=torch.long, device=device)
            return empty_tensor, empty_tensor, empty_tensor

        distance_matrix = self.compute_distance(embeddings, embeddings)

        # 创建标签掩码 - 向量化操作
        labels_equal = labels.unsqueeze(0) == labels.unsqueeze(1)
        labels_not_equal = ~labels_equal

        if self.mining_strategy == 'hard':
            # 向量化难样本挖掘
            return self._mine_hard_triplets_vectorized(
                distance_matrix, labels_equal, labels_not_equal, device
            )
        elif self.mining_strategy == 'all':
            # 向量化全样本挖掘（带内存保护）
            return self._mine_all_triplets_vectorized(
                distance_matrix, labels_equal, labels_not_equal, device
            )
        else:
            raise ValueError(f'Unsupported mining strategy: {self.mining_strategy}')

    def _mine_hard_triplets_vectorized(
        self,
        distance_matrix: Tensor,
        labels_equal: Tensor,
        labels_not_equal: Tensor,
        device: torch.device,
    ) -> Tuple[Tensor, Tensor, Tensor]:
        """向量化难样本三元组挖掘"""
        n = distance_matrix.size(0)

        # 创建自身掩码（排除对角线）
        eye_mask = torch.eye(n, dtype=torch.bool, device=device)
        positive_mask = labels_equal & ~eye_mask
        negative_mask = labels_not_equal

        # 检查每个样本是否有有效的正负样本
        has_positive = positive_mask.any(dim=1)
        has_negative = negative_mask.any(dim=1)
        valid_anchors = has_positive & has_negative

        if not valid_anchors.any():
            empty_tensor = torch.tensor([], dtype=torch.long, device=device)
            return empty_tensor, empty_tensor, empty_tensor

        # 只处理有效的锚点
        valid_indices = valid_anchors.nonzero().squeeze(1)

        # 为无效位置设置极值，避免影响 argmax/argmin
        masked_distance = distance_matrix.clone()
        masked_distance[~positive_mask] = -float('inf')  # 正样本用最大距离
        neg_masked_distance = distance_matrix.clone()
        neg_masked_distance[~negative_mask] = float('inf')  # 负样本用最小距离

        # 向量化选择最难正样本（距离最远）
        hardest_positive_indices = masked_distance[valid_indices].argmax(dim=1)

        # 向量化选择最难负样本（距离最近）
        hardest_negative_indices = neg_masked_distance[valid_indices].argmin(dim=1)

        return valid_indices, hardest_positive_indices, hardest_negative_indices

    def _mine_all_triplets_vectorized(
        self,
        distance_matrix: Tensor,
        labels_equal: Tensor,
        labels_not_equal: Tensor,
        device: torch.device,
    ) -> Tuple[Tensor, Tensor, Tensor]:
        """向量化全样本三元组挖掘（带内存保护）"""
        n = distance_matrix.size(0)

        # 内存保护：如果样本数过大，限制三元组数量
        max_triplets = self.max_triplets  # 可配置的最大三元组数

        # 创建所有可能的三元组索引
        anchors, positives, negatives = [], [], []

        # 使用更高效的索引生成
        eye_mask = torch.eye(n, dtype=torch.bool, device=device)

        for i in range(n):
            # 找到有效的正样本（相同标签，但不是自己）
            positive_indices = (labels_equal[i] & ~eye_mask[i]).nonzero().squeeze(1)
            # 找到有效的负样本
            negative_indices = labels_not_equal[i].nonzero().squeeze(1)

            if len(positive_indices) == 0 or len(negative_indices) == 0:
                continue

            # 生成所有组合（使用广播避免显式循环）
            pos_expanded = positive_indices.unsqueeze(1).expand(-1, len(negative_indices))
            neg_expanded = negative_indices.unsqueeze(0).expand(len(positive_indices), -1)

            # 展平并添加到列表
            num_combinations = len(positive_indices) * len(negative_indices)
            anchors.extend([i] * num_combinations)
            positives.extend(pos_expanded.flatten().tolist())
            negatives.extend(neg_expanded.flatten().tolist())

            # 内存保护：限制三元组数量
            if len(anchors) > max_triplets:
                break

        if len(anchors) == 0:
            empty_tensor = torch.tensor([], dtype=torch.long, device=device)
            return empty_tensor, empty_tensor, empty_tensor

        # 截断到最大数量
        if len(anchors) > max_triplets:
            anchors = anchors[:max_triplets]
            positives = positives[:max_triplets]
            negatives = negatives[:max_triplets]

        return (
            torch.tensor(anchors, device=device),
            torch.tensor(positives, device=device),
            torch.tensor(negatives, device=device),
        )

    def compute_loss(self, batch_data: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """计算三元组损失

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含 'data' 和 'labels' 的字典

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含三元组损失的字典
        """
        data = batch_data['data']
        labels = batch_data['labels']

        # 前向传播
        output = self.forward(data)
        embeddings = output['embeddings']

        # 挖掘三元组
        anchors, positives, negatives = self.mine_pairs(embeddings, labels)

        if len(anchors) == 0:
            return {
                'triplet_loss': torch.tensor(0.0, requires_grad=True),
                'total_loss': torch.tensor(0.0, requires_grad=True),
            }

        # 计算距离
        anchor_embeddings = embeddings[anchors]
        positive_embeddings = embeddings[positives]
        negative_embeddings = embeddings[negatives]

        positive_distances = F.pairwise_distance(anchor_embeddings, positive_embeddings)
        negative_distances = F.pairwise_distance(anchor_embeddings, negative_embeddings)

        # 计算三元组损失
        triplet_loss = F.relu(
            positive_distances - negative_distances + self.margin
        ).mean()

        return {'triplet_loss': triplet_loss, 'total_loss': triplet_loss}


class ContrastiveLoss(MetricLearningBase):
    """对比损失

    实现对比损失函数，通过拉近相似样本、推远不相似样本来学习嵌入。

    References
    ----------
    Hadsell, R., Chopra, S., & LeCun, Y. (2006).
    Dimensionality reduction by learning an invariant mapping.
    In 2006 IEEE Computer Society Conference on Computer Vision and Pattern Recognition (CVPR'06) (Vol. 2, pp. 1735-1742).

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `embedding_dim` : `int`
        - 嵌入向量维度
    - `margin` : `float`, optional
        - 对比损失边界, by default 1.0
    - `distance_metric` : `str`, optional
        - 距离度量, by default 'euclidean'

    Examples
    ----------
    ```python
    # 创建对比损失模型
    model = ContrastiveLoss(
        encoder=encoder,
        embedding_dim=64,
        margin=1.0
    )

    # 训练数据（成对数据）
    batch_data = {
        'data1': torch.randn(32, 784),
        'data2': torch.randn(32, 784),
        'labels': torch.randint(0, 2, (32,))  # 0表示不同类，1表示同类
    }

    loss_dict = model.compute_loss(batch_data)
    ```
    """

    def __init__(
        self,
        encoder: nn.Module,
        embedding_dim: int,
        margin: float = 1.0,
        distance_metric: str = 'euclidean',
    ):
        super().__init__(encoder, embedding_dim, margin)
        self.distance_metric = distance_metric

    def compute_distance(self, embeddings1: Tensor, embeddings2: Tensor) -> Tensor:
        """计算成对距离

        Parameters
        ----------
        - `embeddings1` : `Tensor`
            - 第一组嵌入向量
        - `embeddings2` : `Tensor`
            - 第二组嵌入向量

        Returns
        ----------
        - `Tensor`
            - 成对距离
        """
        if self.distance_metric == 'euclidean':
            # 使用 F.pairwise_distance 进行高效的成对欧氏距离计算
            return F.pairwise_distance(embeddings1, embeddings2, p=2, eps=1e-8)
        elif self.distance_metric == 'cosine':
            # 提升余弦距离的数值稳定性
            cosine_sim = F.cosine_similarity(embeddings1, embeddings2, eps=1e-8)
            # 裁剪到 [-1, 1] 范围
            cosine_sim = torch.clamp(cosine_sim, -1.0, 1.0)
            return 1 - cosine_sim
        else:
            raise ValueError(f'Unsupported distance metric: {self.distance_metric}')

    def mine_pairs(self, embeddings: Tensor, labels: Tensor) -> Tuple[Tensor, Tensor]:
        """生成所有可能的样本对

        Parameters
        ----------
        - `embeddings` : `Tensor`
            - 嵌入向量
        - `labels` : `Tensor`
            - 标签

        Returns
        ----------
        - `Tuple[Tensor, Tensor]`
            - 样本对索引和对应的标签
        """
        device = embeddings.device
        n = len(labels)

        if n < 2:
            # 样本数量不足，返回空张量
            empty_pairs = torch.tensor([], dtype=torch.long, device=device).reshape(0, 2)
            empty_labels = torch.tensor([], dtype=torch.long, device=device)
            return empty_pairs, empty_labels

        # 向量化生成所有样本对索引 - 使用上三角矩阵索引
        i_indices, j_indices = torch.triu_indices(n, n, offset=1, device=device)
        pairs = torch.stack([i_indices, j_indices], dim=1)

        # 向量化计算对标签 - 避免循环
        pair_labels = (labels[i_indices] == labels[j_indices]).long()

        return pairs, pair_labels

    def compute_loss(self, batch_data: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """计算对比损失

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含成对数据的字典

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含对比损失的字典
        """
        if 'data1' in batch_data and 'data2' in batch_data:
            # 成对数据格式
            data1 = batch_data['data1']
            data2 = batch_data['data2']
            labels = batch_data['labels']  # 1表示同类，0表示不同类

            # 前向传播
            output1 = self.forward(data1)
            output2 = self.forward(data2)

            embeddings1 = output1['embeddings']
            embeddings2 = output2['embeddings']

            # 计算距离
            distances = self.compute_distance(embeddings1, embeddings2)

        else:
            # 批次数据格式
            data = batch_data['data']
            labels = batch_data['labels']

            # 前向传播
            output = self.forward(data)
            embeddings = output['embeddings']

            # 挖掘样本对
            pairs, pair_labels = self.mine_pairs(embeddings, labels)

            if len(pairs) == 0:
                return {
                    'contrastive_loss': torch.tensor(0.0, requires_grad=True),
                    'total_loss': torch.tensor(0.0, requires_grad=True),
                }

            # 计算距离
            embeddings1 = embeddings[pairs[:, 0]]
            embeddings2 = embeddings[pairs[:, 1]]
            distances = self.compute_distance(embeddings1, embeddings2)
            labels = pair_labels.float()

        # 计算对比损失
        positive_loss = labels * distances.pow(2)
        negative_loss = (1 - labels) * F.relu(self.margin - distances).pow(2)

        contrastive_loss = (positive_loss + negative_loss).mean()

        return {'contrastive_loss': contrastive_loss, 'total_loss': contrastive_loss}


class CenterLoss(MetricLearningBase):
    """中心损失

    实现中心损失函数，通过最小化每个类别内部样本到类别中心的距离来学习嵌入。

    References
    ----------
    Wen, Y., Zhang, K., Li, Z., & Qiao, Y. (2016).
    A discriminative feature learning approach for deep face recognition.
    In European conference on computer vision (pp. 499-515). Springer.

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `num_classes` : `int`
        - 类别数量
    - `embedding_dim` : `int`
        - 嵌入向量维度
    - `alpha` : `float`, optional
        - 中心更新率, by default 0.5

    Examples
    ----------
    ```python
    # 创建中心损失模型
    model = CenterLoss(
        encoder=encoder,
        num_classes=10,
        embedding_dim=64,
        alpha=0.5
    )

    # 训练数据
    batch_data = {
        'data': torch.randn(32, 784),
        'labels': torch.randint(0, 10, (32,))
    }

    loss_dict = model.compute_loss(batch_data)
    ```
    """

    def __init__(
        self, encoder: nn.Module, num_classes: int, embedding_dim: int, alpha: float = 0.5
    ):
        super().__init__(encoder, embedding_dim)
        self.num_classes = num_classes
        self.alpha = alpha

        # 初始化类别中心
        self.centers = nn.Parameter(torch.randn(num_classes, embedding_dim))

    def compute_distance(self, embeddings1: Tensor, embeddings2: Tensor) -> Tensor:
        """计算欧氏距离"""
        return torch.cdist(embeddings1, embeddings2, p=2)

    def mine_pairs(self, embeddings: Tensor, labels: Tensor) -> Tuple[Tensor, Tensor]:
        """不需要挖掘样本对"""
        return torch.tensor([]), torch.tensor([])

    def compute_loss(self, batch_data: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """计算中心损失

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含 'data' 和 'labels' 的字典

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含中心损失的字典
        """
        data = batch_data['data']
        labels = batch_data['labels']

        # 前向传播
        output = self.forward(data)
        embeddings = output['embeddings']

        # 计算每个样本到对应类别中心的距离
        centers_batch = self.centers[labels]
        center_loss = F.mse_loss(embeddings, centers_batch)

        # 向量化更新类别中心
        if self.training:
            with torch.no_grad():
                # 获取批次中存在的唯一类别
                unique_labels = torch.unique(labels)

                # 向量化计算每个类别的均值和更新
                for label in unique_labels:
                    mask = labels == label
                    if mask.sum() > 0:
                        # 计算该类别样本的均值
                        class_embeddings = embeddings[mask]
                        class_mean = class_embeddings.mean(0)

                        # 更新中心：使用指数移动平均
                        delta = class_mean - self.centers[label]
                        self.centers[label] += self.alpha * delta

        return {'center_loss': center_loss, 'total_loss': center_loss}


class ArcFace(MetricLearningBase):
    """ArcFace 损失

    实现 ArcFace 损失函数，通过在角度空间添加边界来增强类间分离度。

    References
    ----------
    Deng, J., Guo, J., Xue, N., & Zafeiriou, S. (2019).
    ArcFace: Additive angular margin loss for deep face recognition.
    In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 4690-4699).

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `embedding_dim` : `int`
        - 嵌入向量维度
    - `num_classes` : `int`
        - 类别数量
    - `scale` : `float`, optional
        - 缩放参数, by default 64.0
    - `margin` : `float`, optional
        - 角度边界, by default 0.5

    Examples
    ----------
    ```python
    # 创建 ArcFace 损失模型
    model = ArcFace(
        encoder=encoder,
        embedding_dim=512,
        num_classes=10000,
        scale=64.0,
        margin=0.5
    )

    # 训练数据
    batch_data = {
        'data': torch.randn(32, 784),
        'labels': torch.randint(0, 10000, (32,))
    }

    loss_dict = model.compute_loss(batch_data)
    ```
    """

    def __init__(
        self,
        encoder: nn.Module,
        embedding_dim: int,
        num_classes: int,
        scale: float = 64.0,
        margin: float = 0.5,
    ):
        super().__init__(encoder, embedding_dim, margin)
        self.num_classes = num_classes
        self.scale = scale

        # 权重矩阵
        self.weight = nn.Parameter(torch.randn(num_classes, embedding_dim))
        nn.init.xavier_uniform_(self.weight)

    def compute_distance(self, embeddings1: Tensor, embeddings2: Tensor) -> Tensor:
        """计算余弦距离"""
        return 1 - F.cosine_similarity(embeddings1, embeddings2.unsqueeze(1), dim=2)

    def mine_pairs(self, embeddings: Tensor, labels: Tensor) -> Tuple[Tensor, Tensor]:
        """不需要挖掘样本对"""
        return torch.tensor([]), torch.tensor([])

    def compute_loss(self, batch_data: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """计算 ArcFace 损失

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含 'data' 和 'labels' 的字典

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含 ArcFace 损失的字典
        """
        data = batch_data['data']
        labels = batch_data['labels']

        # 前向传播
        output = self.forward(data)
        embeddings = output['embeddings']

        # 归一化权重和特征
        normalized_weight = F.normalize(self.weight, p=2, dim=1)
        normalized_embeddings = F.normalize(embeddings, p=2, dim=1)

        # 计算余弦相似度
        cosine = F.linear(normalized_embeddings, normalized_weight)

        # 计算角度
        sine = torch.sqrt(1.0 - torch.pow(cosine, 2))
        phi = cosine * math.cos(self.margin) - sine * math.sin(self.margin)

        # 创建 one-hot 标签
        one_hot = torch.zeros(cosine.size()).to(cosine.device)
        one_hot.scatter_(1, labels.view(-1, 1).long(), 1)

        # 应用边界
        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
        output *= self.scale

        # 计算交叉熵损失
        arcface_loss = F.cross_entropy(output, labels)

        return {'arcface_loss': arcface_loss, 'total_loss': arcface_loss}


class ProxyNCA(MetricLearningBase):
    """Proxy-NCA 损失

    实现 Proxy-NCA 损失函数，使用代理向量来近似类别分布。

    References
    ----------
    Movshovitz-Attias, Y., Toshev, A., Leung, T. K., Ioffe, S., & Singh, S. (2017).
    No fuss distance metric learning using proxies.
    In Proceedings of the IEEE International Conference on Computer Vision (pp. 360-368).

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `embedding_dim` : `int`
        - 嵌入向量维度
    - `num_classes` : `int`
        - 类别数量
    - `temperature` : `float`, optional
        - 温度参数, by default 0.1

    Examples
    ----------
    ```python
    # 创建 Proxy-NCA 损失模型
    model = ProxyNCA(
        encoder=encoder,
        embedding_dim=128,
        num_classes=100
    )

    # 训练数据
    batch_data = {
        'data': torch.randn(32, 784),
        'labels': torch.randint(0, 100, (32,))
    }

    loss_dict = model.compute_loss(batch_data)
    ```
    """

    def __init__(
        self,
        encoder: nn.Module,
        embedding_dim: int,
        num_classes: int,
        temperature: float = 0.1,
    ):
        super().__init__(encoder, embedding_dim)
        self.num_classes = num_classes
        self.temperature = temperature

        # 代理向量
        self.proxies = nn.Parameter(torch.randn(num_classes, embedding_dim))
        nn.init.xavier_uniform_(self.proxies)

    def compute_distance(self, embeddings1: Tensor, embeddings2: Tensor) -> Tensor:
        """计算欧氏距离"""
        return torch.cdist(embeddings1, embeddings2, p=2)

    def mine_pairs(self, embeddings: Tensor, labels: Tensor) -> Tuple[Tensor, Tensor]:
        """不需要挖掘样本对"""
        return torch.tensor([]), torch.tensor([])

    def compute_loss(self, batch_data: Dict[str, Tensor]) -> Dict[str, Tensor]:
        """计算 Proxy-NCA 损失

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含 'data' 和 'labels' 的字典

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含 Proxy-NCA 损失的字典
        """
        data = batch_data['data']
        labels = batch_data['labels']

        # 前向传播
        output = self.forward(data)
        embeddings = output['embeddings']

        # 归一化
        normalized_embeddings = F.normalize(embeddings, p=2, dim=1)
        normalized_proxies = F.normalize(self.proxies, p=2, dim=1)

        # 计算距离
        distances = self.compute_distance(normalized_embeddings, normalized_proxies)

        # 转换为相似度
        similarities = -distances / self.temperature

        # 计算 NCA 损失
        exp_similarities = torch.exp(similarities)

        # 正样本概率
        positive_similarities = similarities[torch.arange(len(labels)), labels]
        positive_exp = torch.exp(positive_similarities)

        # 总概率
        sum_exp = exp_similarities.sum(dim=1)

        # NCA 损失
        proxy_nca_loss = -torch.log(positive_exp / sum_exp).mean()

        return {'proxy_nca_loss': proxy_nca_loss, 'total_loss': proxy_nca_loss}
