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

# 尝试导入 torch_geometric 以获得最佳性能
try:
    from torch_geometric.nn import global_add_pool, global_max_pool, global_mean_pool

    HAS_TORCH_GEOMETRIC = True
except ImportError:
    HAS_TORCH_GEOMETRIC = False


def _validate_graph_model(
    model: nn.Module,
    node_features: int,
    edge_features: Optional[int] = None,
    output_dim: Optional[int] = None,
) -> None:
    """验证图神经网络模型的输入输出维度是否匹配

    Parameters
    ----------
    - `model` : `nn.Module`
        - 图神经网络模型
    - `node_features` : `int`
        - 节点特征维度
    - `edge_features` : `int`, optional
        - 边特征维度, by default None
    - `output_dim` : `int`, optional
        - 输出维度, by default None

    Raises
    ----------
    - `ValueError`
        - 当模型维度不匹配时抛出
    """
    device = next(model.parameters()).device
    num_nodes = 10
    num_edges = 20

    with torch.no_grad():
        try:
            # 创建测试数据
            x = torch.randn(num_nodes, node_features, device=device)
            edge_index = torch.randint(0, num_nodes, (2, num_edges), device=device)

            batch_data = {'x': x, 'edge_index': edge_index}

            if edge_features is not None:
                edge_attr = torch.randn(num_edges, edge_features, device=device)
                batch_data['edge_attr'] = edge_attr

            # 测试模型
            output = model(batch_data)

            if output_dim is not None:
                expected_shape = (num_nodes, output_dim)
                if output.shape != expected_shape:
                    raise ValueError(
                        f'模型输出形状 {output.shape} 与期望形状 {expected_shape} 不匹配'
                    )
        except Exception as e:
            raise ValueError(f'图神经网络模型测试失败: {e}')


class GraphBase(BaseModule, ABC):
    """图神经网络基类

    提供图神经网络的基本接口和通用功能。

    Parameters
    ----------
    - `model` : `nn.Module`
        - 图神经网络模型
        - **输入**: 应接受包含 'x', 'edge_index' 等键的字典
        - **输出**: 应输出节点表示张量
    """

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

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

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

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含节点表示的字典
        """
        node_embeddings = self.model(batch_data)

        return {'node_embeddings': node_embeddings}

    @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 NodeClassification(GraphBase):
    """节点分类范式

    实现图神经网络的节点分类任务。

    References
    ----------
    Kipf, T. N., & Welling, M. (2016).
    Semi-supervised classification with graph convolutional networks.
    arXiv preprint arXiv:1609.02907.

    Parameters
    ----------
    - `model` : `nn.Module`
        - 图神经网络模型
    - `num_classes` : `int`
        - 分类类别数
    - `node_features` : `int`, optional
        - 节点特征维度, by default None
    - `edge_features` : `int`, optional
        - 边特征维度, by default None
    - `validate_model` : `bool`, optional
        - 是否验证模型, by default True

    Examples
    ----------
    ```python
    # 创建 GCN 模型
    model = GCN(
        in_channels=128,
        hidden_channels=64,
        out_channels=32,
        num_layers=2
    )

    # 创建节点分类范式
    node_clf = NodeClassification(
        model,
        num_classes=7,
        node_features=128
    )

    # 训练
    batch_data = {
        'x': torch.randn(100, 128),  # 节点特征
        'edge_index': torch.randint(0, 100, (2, 200)),  # 边索引
        'y': torch.randint(0, 7, (100,))  # 节点标签
    }
    loss_dict = node_clf.compute_loss(batch_data)
    ```
    """

    def __init__(
        self,
        model: nn.Module,
        num_classes: int,
        node_features: Optional[int] = None,
        edge_features: Optional[int] = None,
        validate_model: bool = True,
    ):
        super().__init__(model)
        self.num_classes = num_classes

        # 模型验证
        if validate_model and node_features is not None:
            _validate_graph_model(model, node_features, edge_features, num_classes)

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

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含 'x', 'edge_index', 'y' 的批次数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含分类损失的字典
        """
        # 前向传播
        output = self.forward(batch_data)
        logits = output['node_embeddings']

        # 获取标签
        labels = batch_data['y']

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

        return {
            'classification_loss': classification_loss,
            'total_loss': classification_loss,
        }


class GraphClassification(GraphBase):
    """图分类范式

    实现图神经网络的图分类任务。

    References
    ----------
    Ying, Z., You, J., Morris, C., Ren, X., Hamilton, W., & Leskovec, J. (2018).
    Hierarchical graph representation learning with differentiable pooling.
    Advances in neural information processing systems, 31, 4800-4810.

    Parameters
    ----------
    - `model` : `nn.Module`
        - 图神经网络模型
    - `pooling` : `str`, optional
        - 图池化方法 ('mean', 'max', 'sum'), by default 'mean'
    - `num_classes` : `int`
        - 分类类别数
    - `node_features` : `int`, optional
        - 节点特征维度, by default None
    - `edge_features` : `int`, optional
        - 边特征维度, by default None
    - `validate_model` : `bool`, optional
        - 是否验证模型, by default True

    Examples
    ----------
    ```python
    # 创建图分类范式
    graph_clf = GraphClassification(
        model,
        pooling='mean',
        num_classes=2,
        node_features=128
    )

    # 训练
    batch_data = {
        'x': torch.randn(100, 128),  # 节点特征
        'edge_index': torch.randint(0, 100, (2, 200)),  # 边索引
        'batch': torch.repeat_interleave(torch.arange(4), 25),  # 批次索引
        'y': torch.randint(0, 2, (4,))  # 图标签
    }
    loss_dict = graph_clf.compute_loss(batch_data)
    ```
    """

    def __init__(
        self,
        model: nn.Module,
        pooling: str = 'mean',
        num_classes: int = 2,
        node_features: Optional[int] = None,
        edge_features: Optional[int] = None,
        validate_model: bool = True,
    ):
        super().__init__(model)
        self.pooling = pooling
        self.num_classes = num_classes

        # 模型验证
        if validate_model and node_features is not None:
            _validate_graph_model(model, node_features, edge_features)

    def global_pooling(self, x: Tensor, batch: Tensor) -> Tensor:
        """全局池化

        使用 torch_geometric 的优化实现（如果可用），否则使用向量化 fallback 版本。

        Parameters
        ----------
        - `x` : `Tensor`
            - 节点表示
        - `batch` : `Tensor`
            - 批次索引

        Returns
        ----------
        - `Tensor`
            - 图表示
        """
        if HAS_TORCH_GEOMETRIC:
            # 使用 torch_geometric 的高度优化实现
            if self.pooling == 'mean':
                return global_mean_pool(x, batch)
            elif self.pooling == 'max':
                return global_max_pool(x, batch)
            elif self.pooling == 'sum':
                return global_add_pool(x, batch)
            else:
                raise ValueError(f'不支持的池化方法: {self.pooling}')
        else:
            # Fallback: 使用向量化的 torch.scatter 实现
            batch_size = int(batch.max()) + 1

            if self.pooling == 'mean':
                # 计算每个图的节点数量
                counts = torch.zeros(batch_size, device=x.device, dtype=torch.float)
                counts.scatter_add_(0, batch, torch.ones_like(batch, dtype=torch.float))

                # 求和然后除以节点数量
                sums = torch.zeros(batch_size, x.size(1), device=x.device)
                sums.scatter_add_(0, batch.unsqueeze(1).expand(-1, x.size(1)), x)
                return sums / counts.unsqueeze(1).clamp(min=1)

            elif self.pooling == 'sum':
                result = torch.zeros(batch_size, x.size(1), device=x.device)
                result.scatter_add_(0, batch.unsqueeze(1).expand(-1, x.size(1)), x)
                return result

            elif self.pooling == 'max':
                # 初始化为负无穷
                result = torch.full(
                    (batch_size, x.size(1)), float('-inf'), device=x.device
                )
                result.scatter_reduce_(
                    0, batch.unsqueeze(1).expand(-1, x.size(1)), x, reduce='amax'
                )
                return result
            else:
                raise ValueError(f'不支持的池化方法: {self.pooling}')

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

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含 'x', 'edge_index', 'batch', 'y' 的批次数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含分类损失的字典
        """
        # 前向传播
        output = self.forward(batch_data)
        node_embeddings = output['node_embeddings']

        # 全局池化
        batch_index = batch_data['batch']
        graph_embeddings = self.global_pooling(node_embeddings, batch_index)

        # 获取标签
        labels = batch_data['y']

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

        return {
            'classification_loss': classification_loss,
            'total_loss': classification_loss,
        }


class LinkPrediction(GraphBase):
    """链接预测范式

    实现图神经网络的链接预测任务。

    References
    ----------
    Zhang, M., & Chen, Y. (2018).
    Link prediction based on graph neural networks.
    Advances in neural information processing systems, 31, 5165-5175.

    Parameters
    ----------
    - `model` : `nn.Module`
        - 图神经网络模型
    - `decoder` : `str`, optional
        - 解码器类型 ('dot', 'mlp'), by default 'dot'
    - `node_features` : `int`, optional
        - 节点特征维度, by default None
    - `edge_features` : `int`, optional
        - 边特征维度, by default None
    - `validate_model` : `bool`, optional
        - 是否验证模型, by default True

    Examples
    ----------
    ```python
    # 创建链接预测范式
    link_pred = LinkPrediction(
        model,
        decoder='dot',
        node_features=128
    )

    # 训练
    batch_data = {
        'x': torch.randn(100, 128),  # 节点特征
        'edge_index': torch.randint(0, 100, (2, 200)),  # 边索引
        'pos_edge_index': torch.randint(0, 100, (2, 50)),  # 正样本边
        'neg_edge_index': torch.randint(0, 100, (2, 50))   # 负样本边
    }
    loss_dict = link_pred.compute_loss(batch_data)
    ```
    """

    def __init__(
        self,
        model: nn.Module,
        decoder: str = 'dot',
        node_features: Optional[int] = None,
        edge_features: Optional[int] = None,
        validate_model: bool = True,
    ):
        super().__init__(model)
        self.decoder = decoder

        # 模型验证
        if validate_model and node_features is not None:
            _validate_graph_model(model, node_features, edge_features)

    def decode_edges(self, z: Tensor, edge_index: Tensor) -> Tensor:
        """解码边的存在概率

        Parameters
        ----------
        - `z` : `Tensor`
            - 节点表示
        - `edge_index` : `Tensor`
            - 边索引

        Returns
        ----------
        - `Tensor`
            - 边存在概率
        """
        row, col = edge_index

        if self.decoder == 'dot':
            # 点积解码器
            return (z[row] * z[col]).sum(dim=1)
        elif self.decoder == 'mlp':
            # MLP解码器（需要预先定义）
            edge_features = torch.cat([z[row], z[col]], dim=1)
            return self.edge_decoder(edge_features).squeeze()
        else:
            raise ValueError(f'不支持的解码器类型: {self.decoder}')

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

        Parameters
        ----------
        - `batch_data` : `Dict[str, Tensor]`
            - 包含 'x', 'edge_index', 'pos_edge_index', 'neg_edge_index' 的批次数据

        Returns
        ----------
        - `Dict[str, Tensor]`
            - 包含链接预测损失的字典
        """
        # 前向传播
        output = self.forward(batch_data)
        node_embeddings = output['node_embeddings']

        # 获取正负样本边
        pos_edge_index = batch_data['pos_edge_index']
        neg_edge_index = batch_data['neg_edge_index']

        # 计算边的存在概率
        pos_scores = self.decode_edges(node_embeddings, pos_edge_index)
        neg_scores = self.decode_edges(node_embeddings, neg_edge_index)

        # 创建标签
        pos_labels = torch.ones_like(pos_scores)
        neg_labels = torch.zeros_like(neg_scores)

        # 合并正负样本
        scores = torch.cat([pos_scores, neg_scores])
        labels = torch.cat([pos_labels, neg_labels])

        # 计算二元交叉熵损失
        link_loss = F.binary_cross_entropy_with_logits(scores, labels)

        return {'link_loss': link_loss, 'total_loss': link_loss}
