import math
from abc import ABC, abstractmethod
from typing import Callable, 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 SelfSupervisedBase(BaseModule, ABC):
    """自监督学习基类

    提供自监督学习算法的基本接口和通用功能。

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `projector` : `nn.Module`, optional
        - 投影头网络, by default None
    - `temperature` : `float`, optional
        - 温度参数, by default 0.07
    """

    def __init__(
        self,
        encoder: nn.Module,
        projector: Optional[nn.Module] = None,
        temperature: float = 0.07,
    ):
        super().__init__()
        self.encoder = encoder
        self.projector = projector
        self.temperature = temperature

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

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

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

        output = {'features': features}

        if self.projector is not None:
            projections = self.projector(features)
            output['projections'] = projections

        return output

    @abstractmethod
    def create_pretext_task(self, data: Tensor) -> Dict[str, Tensor]:
        """创建预训练任务

        Parameters
        ----------
        - `data` : `Tensor`
            - 原始数据

        Returns
        ----------
        - `Dict[str, 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 SwAV(SelfSupervisedBase):
    """SwAV (Swapping Assignments between Views)

    实现 SwAV 自监督学习算法，通过交换不同视图间的聚类分配来学习表征。

    References
    ----------
    Caron, M., Misra, I., Mairal, J., Goyal, P., Bojanowski, P., & Joulin, A. (2020).
    Unsupervised learning of visual features by contrasting cluster assignments.
    Advances in Neural Information Processing Systems, 33, 9912-9924.

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `projector` : `nn.Module`
        - 投影头网络
    - `prototypes` : `nn.Module`
        - 原型层
    - `temperature` : `float`, optional
        - 温度参数, by default 0.1
    - `sinkhorn_iterations` : `int`, optional
        - Sinkhorn 算法迭代次数, by default 3
    - `epsilon` : `float`, optional
        - Sinkhorn 正则化参数, by default 0.05

    Examples
    ----------
    ```python
    # 创建 SwAV 模型
    model = SwAV(
        encoder=encoder,
        projector=projector,
        prototypes=prototypes
    )
    ```
    """

    def __init__(
        self,
        encoder: nn.Module,
        projector: nn.Module,
        prototypes: nn.Module,
        temperature: float = 0.1,
        sinkhorn_iterations: int = 3,
        epsilon: float = 0.05,
    ):
        super().__init__(encoder, projector, temperature)
        self.prototypes = prototypes
        self.sinkhorn_iterations = sinkhorn_iterations
        self.epsilon = epsilon

    def create_pretext_task(self, data: Tensor) -> Dict[str, Tensor]:
        """创建多视图数据

        Parameters
        ----------
        - `data` : `Tensor`
            - 原始图像数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含多个视图的字典
        """
        # 这里应该实现数据增强，创建多个视图
        # 简化实现，假设已经有多个视图
        return {
            'view1': data,
            'view2': data,  # 实际应用中应该是不同的增强版本
        }

    def sinkhorn(self, Q: Tensor) -> Tensor:
        """Sinkhorn 算法进行最优传输

        Parameters
        ----------
        - `Q` : `Tensor`
            - 输入矩阵

        Returns
        ----------
        - `Tensor`
            - 归一化后的分配矩阵
        """
        Q = torch.exp(
            Q / self.epsilon
        ).t()  # Q is K-by-B for consistency with notations from our paper
        B = Q.shape[1]  # number of samples to assign
        K = Q.shape[0]  # how many prototypes

        # make the matrix sums to 1
        sum_Q = torch.sum(Q)
        Q /= sum_Q

        for it in range(self.sinkhorn_iterations):
            # normalize each row: total weight per prototype must be 1/K
            sum_of_rows = torch.sum(Q, dim=1, keepdim=True)
            Q /= sum_of_rows
            Q /= K

            # normalize each column: total weight per sample must be 1/B
            Q /= torch.sum(Q, dim=0, keepdim=True)
            Q /= B

        Q *= B  # the columns must sum to 1 so that Q is an assignment
        return Q.t()

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

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含多视图数据的字典

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

        # 前向传播
        output1 = self.forward(view1)
        output2 = self.forward(view2)

        z1 = F.normalize(output1['projections'], dim=1)
        z2 = F.normalize(output2['projections'], dim=1)

        # 计算原型分数
        p1 = self.prototypes(z1)
        p2 = self.prototypes(z2)

        # Sinkhorn 算法计算分配
        with torch.no_grad():
            q1 = self.sinkhorn(p1)
            q2 = self.sinkhorn(p2)

        # 计算交叉熵损失
        loss = -torch.mean(
            torch.sum(q1 * F.log_softmax(p2 / self.temperature, dim=1), dim=1)
        ) - torch.mean(torch.sum(q2 * F.log_softmax(p1 / self.temperature, dim=1), dim=1))

        return {'swav_loss': loss, 'total_loss': loss}


class DINO(SelfSupervisedBase):
    """DINO (Self-Distillation with No Labels)

    实现 DINO 自监督学习算法，通过自蒸馏方式学习视觉表征。

    References
    ----------
    Caron, M., Touvron, H., Misra, I., Jégou, H., Mairal, J., Bojanowski, P., & Joulin, A. (2021).
    Emerging properties in self-supervised vision transformers.
    In Proceedings of the IEEE/CVF International Conference on Computer Vision (pp. 9650-9660).

    Parameters
    ----------
    - `student_encoder` : `nn.Module`
        - 学生网络编码器
    - `teacher_encoder` : `nn.Module`
        - 教师网络编码器
    - `projector` : `nn.Module`
        - 投影头网络
    - `temperature_student` : `float`, optional
        - 学生网络温度参数, by default 0.1
    - `temperature_teacher` : `float`, optional
        - 教师网络温度参数, by default 0.04
    - `momentum` : `float`, optional
        - 教师网络动量更新参数, by default 0.996
    - `center_momentum` : `float`, optional
        - 中心动量参数, by default 0.9

    Examples
    ----------
    ```python
    # 创建学生和教师编码器
    student_encoder = create_encoder()
    teacher_encoder = create_encoder()

    projector = nn.Sequential(
        nn.Linear(feature_dim, 256),
        nn.GELU(),
        nn.Linear(256, 256),
        nn.GELU(),
        nn.Linear(256, out_dim)
    )

    # 创建 DINO 模型
    model = DINO(
        student_encoder=student_encoder,
        teacher_encoder=teacher_encoder,
        projector=projector
    )
    ```
    """

    def __init__(
        self,
        student_encoder: nn.Module,
        teacher_encoder: nn.Module,
        projector: nn.Module,
        temperature_student: float = 0.1,
        temperature_teacher: float = 0.04,
        momentum: float = 0.996,
        center_momentum: float = 0.9,
    ):
        super().__init__(student_encoder, projector, temperature_student)

        self.student_encoder = student_encoder
        self.teacher_encoder = teacher_encoder
        self.temperature_student = temperature_student
        self.temperature_teacher = temperature_teacher
        self.momentum = momentum
        self.center_momentum = center_momentum

        # 初始化教师网络
        for param_s, param_t in zip(
            self.student_encoder.parameters(), self.teacher_encoder.parameters()
        ):
            param_t.data.copy_(param_s.data)
            param_t.requires_grad = False

        # 中心向量
        self.register_buffer('center', torch.zeros(1, projector[-1].out_features))

    def create_pretext_task(self, data: Tensor) -> Dict[str, Tensor]:
        """创建全局和局部视图

        Parameters
        ----------
        - `data` : `Tensor`
            - 原始图像数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含全局和局部视图的字典
        """
        # 简化实现，实际应该包含不同尺度的裁剪
        return {
            'global_views': [data, data],  # 两个全局视图
            'local_views': [data, data, data, data],  # 四个局部视图
        }

    def update_teacher(self):
        """使用动量更新教师网络"""
        with torch.no_grad():
            for param_s, param_t in zip(
                self.student_encoder.parameters(), self.teacher_encoder.parameters()
            ):
                param_t.data = param_t.data * self.momentum + param_s.data * (
                    1 - self.momentum
                )

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

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含全局和局部视图的字典

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含 DINO 损失的字典
        """
        global_views = batch_data['global_views']
        local_views = batch_data['local_views']

        # 学生网络处理所有视图
        student_outputs = []
        for view in global_views + local_views:
            features = self.student_encoder(view)
            projections = self.projector(features)
            student_outputs.append(projections)

        # 教师网络只处理全局视图
        teacher_outputs = []
        with torch.no_grad():
            for view in global_views:
                features = self.teacher_encoder(view)
                projections = self.projector(features)
                teacher_outputs.append(projections)

        # 计算损失
        total_loss = 0
        n_loss_terms = 0

        for i, teacher_output in enumerate(teacher_outputs):
            teacher_out = F.softmax(
                (teacher_output - self.center) / self.temperature_teacher, dim=-1
            )
            teacher_out = teacher_out.detach()

            for j, student_output in enumerate(student_outputs):
                if j == i:  # 跳过相同的视图
                    continue

                student_out = F.log_softmax(
                    student_output / self.temperature_student, dim=-1
                )
                loss = torch.sum(-teacher_out * student_out, dim=-1).mean()
                total_loss += loss
                n_loss_terms += 1

        total_loss /= n_loss_terms

        # 更新中心
        with torch.no_grad():
            batch_center = torch.cat(teacher_outputs).mean(dim=0, keepdim=True)
            self.center = self.center * self.center_momentum + batch_center * (
                1 - self.center_momentum
            )

        return {'dino_loss': total_loss, 'total_loss': total_loss}


class BarlowTwins(SelfSupervisedBase):
    """Barlow Twins

    实现 Barlow Twins 自监督学习算法，通过最大化不同视图间的互相关矩阵对角线元素，
    最小化非对角线元素来学习表征。

    References
    ----------
    Zbontar, J., Jing, L., Misra, I., LeCun, Y., & Deny, S. (2021).
    Barlow twins: Self-supervised learning via redundancy reduction.
    In International Conference on Machine Learning (pp. 12310-12320). PMLR.

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `projector` : `nn.Module`
        - 投影头网络
    - `lambda_param` : `float`, optional
        - 正则化参数, by default 0.0051

    Examples
    ----------
    ```python
    # 创建编码器和投影头
    encoder = create_resnet_encoder()
    projector = nn.Sequential(
        nn.Linear(2048, 8192),
        nn.BatchNorm1d(8192),
        nn.ReLU(),
        nn.Linear(8192, 8192),
        nn.BatchNorm1d(8192),
        nn.ReLU(),
        nn.Linear(8192, 8192)
    )

    # 创建 Barlow Twins 模型
    model = BarlowTwins(
        encoder=encoder,
        projector=projector
    )
    ```
    """

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

    def create_pretext_task(self, data: Tensor) -> Dict[str, Tensor]:
        """创建两个增强视图

        Parameters
        ----------
        - `data` : `Tensor`
            - 原始数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含两个视图的字典
        """
        return {
            'view1': data,
            'view2': data,  # 实际应该是不同的增强版本
        }

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

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含两个视图的字典

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

        # 前向传播
        output1 = self.forward(view1)
        output2 = self.forward(view2)

        z1 = output1['projections']
        z2 = output2['projections']

        # 标准化
        z1_norm = (z1 - z1.mean(0)) / z1.std(0)
        z2_norm = (z2 - z2.mean(0)) / z2.std(0)

        # 计算互相关矩阵
        batch_size = z1.size(0)
        c = torch.mm(z1_norm.T, z2_norm) / batch_size

        # 计算损失
        on_diag = torch.diagonal(c).add_(-1).pow_(2).sum()
        off_diag = self.off_diagonal(c).pow_(2).sum()

        loss = on_diag + self.lambda_param * off_diag

        return {'barlow_twins_loss': loss, 'total_loss': loss}

    def off_diagonal(self, x: Tensor) -> Tensor:
        """获取矩阵的非对角线元素"""
        n, m = x.shape
        assert n == m
        return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()


class MAE(SelfSupervisedBase):
    """Masked Autoencoder (MAE)

    实现掩码自编码器，通过重建被掩码的图像块来学习视觉表征。

    References
    ----------
    He, K., Chen, X., Xie, S., Li, Y., Dollár, P., & Girshick, R. (2022).
    Masked autoencoders are scalable vision learners.
    In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 16000-16009).

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `decoder` : `nn.Module`
        - 解码器网络
    - `mask_ratio` : `float`, optional
        - 掩码比例, by default 0.75

    Examples
    ----------
    ```python
    # 创建 MAE 模型
    model = MAE(
        encoder=encoder,
        decoder=decoder,
        mask_ratio=0.75
    )

    # 训练
    loss_dict = model.compute_loss(batch_data)
    ```
    """

    def __init__(
        self,
        encoder: nn.Module,
        decoder: nn.Module,
        mask_ratio: float = 0.75,
        patch_size: int = 16,
    ):
        super().__init__(encoder)
        self.decoder = decoder
        self.mask_ratio = mask_ratio
        self.patch_size = patch_size

    def create_pretext_task(self, data: Tensor) -> Dict[str, Tensor]:
        """创建掩码重建任务

        Parameters
        ----------
        - `data` : `Tensor`
            - 原始图像数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含掩码信息的字典
        """
        batch_size, channels, height, width = data.shape

        # 将图像分割成块
        num_patches_h = height // self.patch_size
        num_patches_w = width // self.patch_size
        num_patches = num_patches_h * num_patches_w

        # 生成随机掩码
        num_masked = int(num_patches * self.mask_ratio)
        mask = torch.zeros(batch_size, num_patches, dtype=torch.bool)

        for i in range(batch_size):
            masked_indices = torch.randperm(num_patches)[:num_masked]
            mask[i, masked_indices] = True

        return {'images': data, 'mask': mask}

    def patchify(self, imgs: Tensor) -> Tensor:
        """将图像分割成块"""
        p = self.patch_size
        batch_size, channels, height, width = imgs.shape

        if height % p != 0 or width % p != 0:
            raise ValueError(
                f'Image dimensions ({height}x{width}) must be divisible by patch size ({p}).'
            )

        num_patches_h = height // p
        num_patches_w = width // p

        patches = imgs.view(batch_size, channels, num_patches_h, p, num_patches_w, p)
        patches = patches.permute(0, 2, 4, 1, 3, 5).contiguous()
        patches = patches.view(batch_size, num_patches_h * num_patches_w, -1)

        return patches

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

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含图像和掩码的字典

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含重建损失的字典
        """
        images = batch_data['images']
        mask = batch_data['mask']

        # 将图像转换为块
        patches = self.patchify(images)

        # 编码器只处理可见的块
        visible_patches = patches[~mask].view(images.size(0), -1, patches.size(-1))
        encoded = self.encoder(visible_patches)

        # 解码器重建所有块
        reconstructed = self.decoder(encoded, mask)

        # 计算重建损失（只在掩码区域）
        target = patches
        loss = F.mse_loss(reconstructed[mask], target[mask])

        return {'reconstruction_loss': loss, 'total_loss': loss}


class SimSiam(SelfSupervisedBase):
    """SimSiam (Simple Siamese)

    实现 SimSiam 自监督学习算法，通过简单的孪生网络结构学习表征，
    使用停止梯度操作防止表征坍塌。

    References
    ----------
    Chen, X., & He, K. (2021).
    Exploring simple siamese representation learning.
    In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (pp. 15750-15758).

    Parameters
    ----------
    - `encoder` : `nn.Module`
        - 编码器网络
    - `projector` : `nn.Module`
        - 投影头网络
    - `predictor` : `nn.Module`
        - 预测头网络

    Examples
    ----------
    ```python
    # 创建网络组件
    encoder = create_resnet_encoder()
    projector = create_mlp_projector()
    predictor = create_mlp_predictor()

    # 创建 SimSiam 模型
    model = SimSiam(
        encoder=encoder,
        projector=projector,
        predictor=predictor
    )
    ```
    """

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

    def create_pretext_task(self, data: Tensor) -> Dict[str, Tensor]:
        """创建两个增强视图

        Parameters
        ----------
        - `data` : `Tensor`
            - 原始数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含两个视图的字典
        """
        return {
            'view1': data,
            'view2': data,  # 实际应该是不同的增强版本
        }

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

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含两个视图的字典

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

        # 前向传播
        output1 = self.forward(view1)
        output2 = self.forward(view2)

        z1 = output1['projections']
        z2 = output2['projections']

        # 预测
        p1 = self.predictor(z1)
        p2 = self.predictor(z2)

        # 计算负余弦相似度损失
        loss1 = self.negative_cosine_similarity(p1, z2.detach())
        loss2 = self.negative_cosine_similarity(p2, z1.detach())

        loss = (loss1 + loss2) * 0.5

        return {'simsiam_loss': loss, 'total_loss': loss}

    def negative_cosine_similarity(self, p: Tensor, z: Tensor) -> Tensor:
        """计算负余弦相似度"""
        p = F.normalize(p, dim=1)
        z = F.normalize(z, dim=1)
        return -(p * z).sum(dim=1).mean()
