from abc import ABC, abstractmethod
from typing import Dict, 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


def _validate_encoder_projector(
    encoder: nn.Module,
    projector: nn.Module,
    input_shape: Tuple[int, ...],
    feature_dim: int,
    projection_dim: int,
) -> None:
    """验证编码器和投影头的输入输出维度是否匹配

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `projector` : `nn.Module`
        - 投影头网络
    - `input_shape` : `Tuple[int, ...]`
        - 输入数据的形状 (不包括batch维度)
    - `feature_dim` : `int`
        - 特征维度
    - `projection_dim` : `int`
        - 投影维度

    Raises
    ----------
    - `ValueError`
        - 当模型维度不匹配时抛出
    """
    device = next(encoder.parameters()).device
    batch_size = 2

    with torch.no_grad():
        # 测试编码器
        try:
            test_input = torch.randn(batch_size, *input_shape, device=device)
            features = encoder(test_input)
            expected_feature_shape = (batch_size, feature_dim)
            if features.shape != expected_feature_shape:
                raise ValueError(
                    f'编码器输出形状 {features.shape} 与期望形状 {expected_feature_shape} 不匹配'
                )
        except Exception as e:
            raise ValueError(f'编码器无法处理形状为 {input_shape} 的输入: {e}')

        # 测试投影头
        try:
            projections = projector(features)
            expected_proj_shape = (batch_size, projection_dim)
            if projections.shape != expected_proj_shape:
                raise ValueError(
                    f'投影头输出形状 {projections.shape} 与期望形状 {expected_proj_shape} 不匹配'
                )
        except Exception as e:
            raise ValueError(f'投影头无法处理 {feature_dim} 维特征: {e}')


class ContrastiveBase(BaseModule, ABC):
    """对比学习基类

    提供对比学习的基本接口和通用功能。

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
        - **输入维度**: 应接受形状为 `(batch_size, *input_shape)` 的输入
        - **输出维度**: 应输出形状为 `(batch_size, feature_dim)` 的特征表示
    - `projector` : `nn.Module`
        - 投影头网络
        - **输入维度**: 应接受形状为 `(batch_size, feature_dim)` 的特征
        - **输出维度**: 应输出形状为 `(batch_size, projection_dim)` 的投影表示
    """

    def __init__(self, encoder: nn.Module, projector: nn.Module):
        super().__init__()
        self.encoder = encoder
        self.projector = projector

    def encode(self, x: Tensor) -> Tensor:
        """编码输入数据为特征表示

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

        Returns
        ----------
        - `Tensor`
            - 特征表示
        """
        return self.encoder(x)

    def project(self, features: Tensor) -> Tensor:
        """将特征投影到对比学习空间

        Parameters
        ----------
        - `features` : `Tensor`
            - 特征表示

        Returns
        ----------
        - `Tensor`
            - 投影表示
        """
        return self.projector(features)

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

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

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含特征和投影的字典
        """
        features = self.encode(x)
        projections = self.project(features)

        return {'features': features, 'projections': projections}

    @abstractmethod
    def contrastive_loss(self, projections: Tensor, **kwargs) -> Tensor:
        """计算对比损失

        Parameters
        ----------
        - `projections` : `Tensor`
            - 投影表示
        - `**kwargs`
            - 其他参数

        Returns
        ----------
        - `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 SimCLR(ContrastiveBase):
    """Simple Contrastive Learning (SimCLR)

    实现 SimCLR 对比学习方法。

    References
    ----------
    Chen, T., Kornblith, S., Norouzi, M., & Hinton, G. (2020).
    A simple framework for contrastive learning of visual representations.
    In International conference on machine learning (pp. 1597-1607). PMLR.

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `projector` : `nn.Module`
        - 投影头网络
    - `temperature` : `float`, optional
        - 温度参数, by default 0.1
    - `input_shape` : `Tuple[int, ...]`, optional
        - 输入数据的形状, by default None
    - `feature_dim` : `int`, optional
        - 特征维度, by default None
    - `projection_dim` : `int`, optional
        - 投影维度, by default None
    - `validate_model` : `bool`, optional
        - 是否验证模型, by default True

    Examples
    ----------
    ```python
    # 创建编码器和投影头
    encoder = nn.Sequential(
        nn.Flatten(),
        nn.Linear(784, 512),
        nn.ReLU(),
        nn.Linear(512, 128)  # 输出128维特征
    )
    projector = nn.Sequential(
        nn.Linear(128, 64),
        nn.ReLU(),
        nn.Linear(64, 32)  # 输出32维投影
    )

    # 创建 SimCLR
    simclr = SimCLR(
        encoder, projector,
        temperature=0.1,
        input_shape=(28, 28),
        feature_dim=128,
        projection_dim=32
    )

    # 训练
    # batch_data 应包含 'anchor' 和 'positive' 键
    batch_data = {
        'anchor': torch.randn(32, 28, 28),
        'positive': torch.randn(32, 28, 28)
    }
    loss_dict = simclr.compute_loss(batch_data)
    ```
    """

    def __init__(
        self,
        encoder: nn.Module,
        projector: nn.Module,
        temperature: float = 0.1,
        input_shape: Optional[Tuple[int, ...]] = None,
        feature_dim: Optional[int] = None,
        projection_dim: Optional[int] = None,
        validate_model: bool = True,
    ):
        super().__init__(encoder, projector)
        self.temperature = temperature

        # 模型验证
        if validate_model and all(
            x is not None for x in [input_shape, feature_dim, projection_dim]
        ):
            _validate_encoder_projector(
                encoder, projector, input_shape, feature_dim, projection_dim
            )

    def contrastive_loss(self, projections: Tensor, **kwargs) -> Tensor:
        """计算 SimCLR 对比损失

        Parameters
        ----------
        - `projections` : `Tensor`
            - 投影表示，形状为 (2*batch_size, projection_dim)

        Returns
        ----------
        - `Tensor`
            - 对比损失
        """
        batch_size = projections.shape[0] // 2

        # 归一化投影
        projections = F.normalize(projections, dim=1)

        # 计算相似度矩阵
        similarity_matrix = torch.matmul(projections, projections.T) / self.temperature

        # 创建标签：正样本对的索引
        labels = torch.cat(
            [torch.arange(batch_size) + batch_size, torch.arange(batch_size)]
        ).to(projections.device)

        # 移除自相似度
        mask = torch.eye(2 * batch_size, dtype=torch.bool).to(projections.device)
        similarity_matrix = similarity_matrix.masked_fill(mask, -float('inf'))

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

        return loss

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

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含 'anchor' 和 'positive' 的批次数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含对比损失的字典
        """
        anchor = batch_data['anchor']
        positive = batch_data['positive']

        # 合并 anchor 和 positive
        combined_input = torch.cat([anchor, positive], dim=0)

        # 前向传播
        output = self.forward(combined_input)
        projections = output['projections']

        # 计算对比损失
        contrastive_loss = self.contrastive_loss(projections)

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


class MoCo(ContrastiveBase):
    """Momentum Contrast (MoCo)

    实现 MoCo 对比学习方法，使用动量更新的队列机制。

    References
    ----------
    He, K., Fan, H., Wu, Y., Xie, S., & Girshick, R. (2020).
    Momentum contrast for unsupervised visual representation learning.
    In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition (pp. 9729-9738).

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `projector` : `nn.Module`
        - 投影头网络
    - `temperature` : `float`, optional
        - 温度参数, by default 0.07
    - `momentum` : `float`, optional
        - 动量参数, by default 0.999
    - `queue_size` : `int`, optional
        - 队列大小, by default 65536
    - `input_shape` : `Tuple[int, ...]`, optional
        - 输入数据的形状, by default None
    - `feature_dim` : `int`, optional
        - 特征维度, by default None
    - `projection_dim` : `int`, optional
        - 投影维度, by default None
    - `validate_model` : `bool`, optional
        - 是否验证模型, by default True

    Examples
    ----------
    ```python
    # 创建 MoCo
    moco = MoCo(
        encoder, projector,
        temperature=0.07,
        momentum=0.999,
        queue_size=65536,
        input_shape=(28, 28),
        feature_dim=128,
        projection_dim=32
    )

    # 训练
    batch_data = {
        'query': torch.randn(32, 28, 28),
        'key': torch.randn(32, 28, 28)
    }
    loss_dict = moco.compute_loss(batch_data)
    ```
    """

    def __init__(
        self,
        encoder: nn.Module,
        projector: nn.Module,
        temperature: float = 0.07,
        momentum: float = 0.999,
        queue_size: int = 65536,
        input_shape: Optional[Tuple[int, ...]] = None,
        feature_dim: Optional[int] = None,
        projection_dim: Optional[int] = None,
        validate_model: bool = True,
    ):
        super().__init__(encoder, projector)
        self.temperature = temperature
        self.momentum = momentum
        self.queue_size = queue_size

        # 创建动量编码器和投影头
        self.momentum_encoder = self._create_momentum_copy(encoder)
        self.momentum_projector = self._create_momentum_copy(projector)

        # 初始化队列
        if projection_dim is not None:
            self.register_buffer('queue', torch.randn(projection_dim, queue_size))
            self.queue = F.normalize(self.queue, dim=0)
            self.register_buffer('queue_ptr', torch.zeros(1, dtype=torch.long))

        # 模型验证
        if validate_model and all(
            x is not None for x in [input_shape, feature_dim, projection_dim]
        ):
            _validate_encoder_projector(
                encoder, projector, input_shape, feature_dim, projection_dim
            )

    def _create_momentum_copy(self, model: nn.Module) -> nn.Module:
        """创建模型的动量副本"""
        import copy

        momentum_model = copy.deepcopy(model)
        for param in momentum_model.parameters():
            param.requires_grad = False
        return momentum_model

    @torch.no_grad()
    def _momentum_update(self):
        """动量更新"""
        # 更新动量编码器
        for param_q, param_k in zip(
            self.encoder.parameters(), self.momentum_encoder.parameters()
        ):
            param_k.data = param_k.data * self.momentum + param_q.data * (
                1.0 - self.momentum
            )

        # 更新动量投影头
        for param_q, param_k in zip(
            self.projector.parameters(), self.momentum_projector.parameters()
        ):
            param_k.data = param_k.data * self.momentum + param_q.data * (
                1.0 - self.momentum
            )

    @torch.no_grad()
    def _dequeue_and_enqueue(self, keys: Tensor):
        """更新队列"""
        batch_size = keys.shape[0]

        ptr = int(self.queue_ptr)

        # 替换队列中的键
        if ptr + batch_size <= self.queue_size:
            self.queue[:, ptr : ptr + batch_size] = keys.T
        else:
            # 处理队列溢出
            remaining = self.queue_size - ptr
            self.queue[:, ptr:] = keys[:remaining].T
            self.queue[:, : batch_size - remaining] = keys[remaining:].T

        ptr = (ptr + batch_size) % self.queue_size
        self.queue_ptr[0] = ptr

    def contrastive_loss(self, queries: Tensor, keys: Tensor, **kwargs) -> Tensor:
        """计算 MoCo 对比损失

        Parameters
        ----------
        - `queries` : `Tensor`
            - 查询投影，形状为 (batch_size, projection_dim)
        - `keys` : `Tensor`
            - 键投影，形状为 (batch_size, projection_dim)

        Returns
        ----------
        - `Tensor`
            - 对比损失
        """
        # 归一化
        queries = F.normalize(queries, dim=1)
        keys = F.normalize(keys, dim=1)

        # 正样本：query 和对应的 key
        l_pos = torch.einsum('nc,nc->n', [queries, keys]).unsqueeze(-1)

        # 负样本：query 和队列中的 key
        l_neg = torch.einsum('nc,ck->nk', [queries, self.queue.clone().detach()])

        # 合并正负样本
        logits = torch.cat([l_pos, l_neg], dim=1) / self.temperature

        # 标签：正样本在第0位
        labels = torch.zeros(logits.shape[0], dtype=torch.long).to(logits.device)

        loss = F.cross_entropy(logits, labels)

        return loss

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

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含 'query' 和 'key' 的批次数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含对比损失的字典
        """
        query = batch_data['query']
        key = batch_data['key']

        # Query 前向传播
        query_output = self.forward(query)
        queries = query_output['projections']

        # Key 前向传播（使用动量网络）
        with torch.no_grad():
            self._momentum_update()
            key_features = self.momentum_encoder(key)
            keys = self.momentum_projector(key_features)

        # 计算对比损失
        contrastive_loss = self.contrastive_loss(queries, keys)

        # 更新队列
        self._dequeue_and_enqueue(keys)

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


class BYOL(ContrastiveBase):
    """Bootstrap Your Own Latent (BYOL)

    实现 BYOL 对比学习方法，不使用负样本。

    References
    ----------
    Grill, J. B., Strub, F., Altché, F., Tallec, C., Richemond, P., Buchatskaya, E., ... & Valko, M. (2020).
    Bootstrap your own latent-a new approach to self-supervised learning.
    Advances in neural information processing systems, 33, 21271-21284.

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `projector` : `nn.Module`
        - 投影头网络
    - `predictor` : `nn.Module`
        - 预测头网络
    - `momentum` : `float`, optional
        - 动量参数, by default 0.996
    - `input_shape` : `Tuple[int, ...]`, optional
        - 输入数据的形状, by default None
    - `feature_dim` : `int`, optional
        - 特征维度, by default None
    - `projection_dim` : `int`, optional
        - 投影维度, by default None
    - `validate_model` : `bool`, optional
        - 是否验证模型, by default True

    Examples
    ----------
    ```python
    # 创建预测头
    predictor = nn.Sequential(
        nn.Linear(32, 16),
        nn.ReLU(),
        nn.Linear(16, 32)
    )

    # 创建 BYOL
    byol = BYOL(
        encoder, projector, predictor,
        momentum=0.996,
        input_shape=(28, 28),
        feature_dim=128,
        projection_dim=32
    )

    # 训练
    batch_data = {
        'view1': torch.randn(32, 28, 28),
        'view2': torch.randn(32, 28, 28)
    }
    loss_dict = byol.compute_loss(batch_data)
    ```
    """

    def __init__(
        self,
        encoder: nn.Module,
        projector: nn.Module,
        predictor: nn.Module,
        momentum: float = 0.996,
        input_shape: Optional[Tuple[int, ...]] = None,
        feature_dim: Optional[int] = None,
        projection_dim: Optional[int] = None,
        validate_model: bool = True,
    ):
        super().__init__(encoder, projector)
        self.predictor = predictor
        self.momentum = momentum

        # 创建目标网络
        self.target_encoder = self._create_momentum_copy(encoder)
        self.target_projector = self._create_momentum_copy(projector)

        # 模型验证
        if validate_model and all(
            x is not None for x in [input_shape, feature_dim, projection_dim]
        ):
            _validate_encoder_projector(
                encoder, projector, input_shape, feature_dim, projection_dim
            )

    def _create_momentum_copy(self, model: nn.Module) -> nn.Module:
        """创建模型的动量副本"""
        import copy

        momentum_model = copy.deepcopy(model)
        for param in momentum_model.parameters():
            param.requires_grad = False
        return momentum_model

    @torch.no_grad()
    def _momentum_update(self):
        """动量更新目标网络"""
        # 更新目标编码器
        for param_online, param_target in zip(
            self.encoder.parameters(), self.target_encoder.parameters()
        ):
            param_target.data = param_target.data * self.momentum + param_online.data * (
                1.0 - self.momentum
            )

        # 更新目标投影头
        for param_online, param_target in zip(
            self.projector.parameters(), self.target_projector.parameters()
        ):
            param_target.data = param_target.data * self.momentum + param_online.data * (
                1.0 - self.momentum
            )

    def predict(self, projections: Tensor) -> Tensor:
        """预测目标表示

        Parameters
        ----------
        - `projections` : `Tensor`
            - 投影表示

        Returns
        ----------
        - `Tensor`
            - 预测表示
        """
        return self.predictor(projections)

    def contrastive_loss(self, predictions: Tensor, targets: Tensor, **kwargs) -> Tensor:
        """计算 BYOL 损失（均方误差）

        Parameters
        ----------
        - `predictions` : `Tensor`
            - 预测表示
        - `targets` : `Tensor`
            - 目标表示

        Returns
        ----------
        - `Tensor`
            - BYOL 损失
        """
        # 归一化
        predictions = F.normalize(predictions, dim=1)
        targets = F.normalize(targets, dim=1)

        # 计算均方误差
        loss = 2 - 2 * (predictions * targets).sum(dim=1).mean()

        return loss

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

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含 'view1' 和 'view2' 的批次数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含对比损失的字典
        """
        view1 = batch_data['view1']
        view2 = batch_data['view2']

        # Online 网络前向传播
        online_output1 = self.forward(view1)
        online_output2 = self.forward(view2)

        online_proj1 = online_output1['projections']
        online_proj2 = online_output2['projections']

        # 预测
        online_pred1 = self.predict(online_proj1)
        online_pred2 = self.predict(online_proj2)

        # Target 网络前向传播
        with torch.no_grad():
            self._momentum_update()

            target_features1 = self.target_encoder(view1)
            target_features2 = self.target_encoder(view2)

            target_proj1 = self.target_projector(target_features1)
            target_proj2 = self.target_projector(target_features2)

        # 计算双向损失
        loss1 = self.contrastive_loss(online_pred1, target_proj2)
        loss2 = self.contrastive_loss(online_pred2, target_proj1)

        total_loss = (loss1 + loss2) / 2

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