from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, Optional, Tuple, TypeAlias

import torch
import torch.nn.functional as F

from zyplib.nn.modules import BaseModule

TorchTensor: TypeAlias = 'torch.Tensor'

# ============================================================================
# 枚举类型
# ============================================================================


class PoolingType(str, Enum):
    """图级池化类型枚举

    - `MEAN`: 平均池化
    - `MAX`: 最大池化
    - `ADD`: 加和池化
    - `CONCAT`: 拼接池化，将 mean/max/add 结果拼接，输出维度为输入的 3 倍
    - `ATTENTION`: 注意力池化，学习节点重要性权重
    """

    MEAN = 'mean'
    MAX = 'max'
    ADD = 'add'
    CONCAT = 'concat'
    ATTENTION = 'attention'


class GNNType(str, Enum):
    """GNN 模型类型枚举

    - `GCN`: 图卷积网络
    - `GAT`: 图注意力网络
    - `GIN`: 图同构网络
    - `SAGE`: GraphSAGE 网络
    """

    GCN = 'gcn'
    GAT = 'gat'
    GIN = 'gin'
    SAGE = 'sage'


# ============================================================================
# 配置类
# ============================================================================


@dataclass
class GraphClassifierConfig:
    """图分类器配置

    Parameters
    ----------
    - `num_classes` : `int`, 分类类别数
    - `hidden_dim` : `int`, 隐藏层维度，用于分类头, by default 128
    - `classifier_hidden_dim` : `Optional[int]`, 分类头隐藏层维度，默认等于 hidden_dim, by default None
    - `classifier_dropout` : `float`, 分类头 Dropout 概率, by default 0.5
    - `drop_edge_prob` : `float`, 训练时边 Dropout 概率，用于正则化, by default 0.0
    - `label_smoothing` : `float`, 标签平滑系数, by default 0.0
    - `class_weights` : `Optional[TorchTensor]`, 类别权重，处理类别不平衡, by default None

    Examples
    ----------
    ```python
    config = GraphClassifierConfig(
        num_classes=4,
        hidden_dim=128,
        classifier_dropout=0.3,
        label_smoothing=0.1,
    )
    ```
    """

    num_classes: int
    hidden_dim: int = 128
    classifier_hidden_dim: Optional[int] = None
    classifier_dropout: float = 0.5
    drop_edge_prob: float = 0.0
    label_smoothing: float = 0.0
    class_weights: Optional[TorchTensor] = None

    def __post_init__(self):
        """参数验证"""
        if self.num_classes < 2:
            raise ValueError(f'num_classes must be >= 2, got {self.num_classes}')
        if not 0 <= self.classifier_dropout < 1:
            raise ValueError(
                f'classifier_dropout must be in [0, 1), got {self.classifier_dropout}'
            )
        if not 0 <= self.drop_edge_prob < 1:
            raise ValueError(
                f'drop_edge_prob must be in [0, 1), got {self.drop_edge_prob}'
            )
        if not 0 <= self.label_smoothing < 1:
            raise ValueError(
                f'label_smoothing must be in [0, 1), got {self.label_smoothing}'
            )


# ============================================================================
# 输出协议
# ============================================================================


@dataclass
class GraphOutput:
    """图分类器输出

    统一的输出格式，包含分类结果和中间表示。

    Parameters
    ----------
    - `logits` : `TorchTensor`, 分类 logits, shape: `(batch_size, num_classes)`
    - `graph_emb` : `TorchTensor`, 图级嵌入, shape: `(batch_size, emb_dim)`
    - `node_emb` : `Optional[TorchTensor]`, 节点嵌入, shape: `(num_nodes, emb_dim)`, by default None
    - `extra` : `Dict[str, Any]`, 通用扩展字典，可存储任意额外信息, by default {}

    Examples
    ----------
    ```python
    output = model(batch)
    probs = F.softmax(output.logits, dim=-1)
    preds = output.logits.argmax(dim=-1)
    # 访问图嵌入用于下游任务
    embeddings = output.graph_emb
    ```
    """

    logits: TorchTensor
    graph_emb: TorchTensor
    node_emb: Optional[TorchTensor] = None
    extra: Dict[str, Any] = field(default_factory=dict)

    @property
    def predictions(self) -> TorchTensor:
        """获取预测类别"""
        return self.logits.argmax(dim=-1)

    @property
    def probabilities(self) -> TorchTensor:
        """获取预测概率"""
        return F.softmax(self.logits, dim=-1)


# ============================================================================
# Abstract
# ============================================================================


class BaseGraphStructureLearner(BaseModule, ABC):
    """图结构学习器抽象基类

    负责图结构的预处理或动态学习，可以修改边连接、边权重等。
    采用策略模式，不同的结构学习策略实现同一接口，可在运行时灵活切换。

    子类需实现
    ----------
    - `forward()`: 接收原始图结构，返回处理后的图结构
    """

    @abstractmethod
    def forward(
        self,
        x: TorchTensor,
        edge_index: TorchTensor,
        edge_attr: Optional[TorchTensor] = None,
        batch: Optional[TorchTensor] = None,
    ) -> Tuple[TorchTensor, TorchTensor, Optional[TorchTensor]]:
        """处理图结构

        Parameters
        ----------
        - `x` : `TorchTensor`, 节点特征, shape: `(num_nodes, num_features)`
        - `edge_index` : `TorchTensor`, 边索引, shape: `(2, num_edges)`
        - `edge_attr` : `Optional[TorchTensor]`, 边特征, shape: `(num_edges, edge_features)`, by default None
        - `batch` : `Optional[TorchTensor]`, 批次索引, shape: `(num_nodes,)`, by default None

        Returns
        ----------
        - `Tuple[TorchTensor, TorchTensor, Optional[TorchTensor]]`, 处理后的 `(x, edge_index, edge_attr)`
        """
        pass


class BaseGraphEncoder(BaseModule, ABC):
    """图编码器抽象基类

    负责将节点特征编码为节点嵌入。forward() 只返回 TorchTensor，可视化需求通过 get_attention_weights() 方法满足。

    子类需实现
    ----------
    - `forward()`: 返回节点嵌入
    - `get_output_dim()`: 返回输出维度

    可选实现
    ----------
    - `get_attention_weights()`: 返回注意力权重，用于可视化
    """

    @abstractmethod
    def forward(
        self,
        x: TorchTensor,
        edge_index: TorchTensor,
        edge_attr: Optional[TorchTensor] = None,
        batch: Optional[TorchTensor] = None,
    ) -> TorchTensor:
        """编码节点特征

        Parameters
        ----------
        - `x` : `TorchTensor`, 节点特征, shape: `(num_nodes, in_channels)`
        - `edge_index` : `TorchTensor`, 边索引, shape: `(2, num_edges)`
        - `edge_attr` : `Optional[TorchTensor]`, 边特征, by default None
        - `batch` : `Optional[TorchTensor]`, 批次索引, by default None

        Returns
        ----------
        - `TorchTensor`, 节点嵌入, shape: `(num_nodes, out_channels)`
        """
        pass

    @abstractmethod
    def get_output_dim(self) -> int:
        """获取输出维度"""
        pass

    def get_attention_weights(
        self,
        x: TorchTensor,
        edge_index: TorchTensor,
        edge_attr: Optional[TorchTensor] = None,
    ) -> Optional[Any]:
        """获取注意力权重，用于可视化

        默认返回 None，子类可选实现。

        Returns
        ----------
        - `Optional[Any]`, 注意力权重，格式取决于具体实现。对于 GAT: `(edge_index, attention_weights)`
        """
        return None


class BaseGraphReadout(BaseModule, ABC):
    """图级读出抽象基类

    将节点嵌入聚合为图级表示。

    子类需实现
    ----------
    - `forward()`: 执行池化操作
    - `get_output_dim()`: 返回输出维度
    """

    @abstractmethod
    def forward(self, x: TorchTensor, batch: Optional[TorchTensor] = None) -> TorchTensor:
        """执行图级池化

        Parameters
        ----------
        - `x` : `TorchTensor`, 节点嵌入, shape: `(num_nodes, node_dim)`
        - `batch` : `Optional[TorchTensor]`, 批次索引, shape: `(num_nodes,)`, by default None

        Returns
        ----------
        - `TorchTensor`, 图嵌入, shape: `(batch_size, graph_dim)`
        """
        pass

    @abstractmethod
    def get_output_dim(self, input_dim: int) -> int:
        """计算输出维度

        Parameters
        ----------
        - `input_dim` : `int`, 输入节点嵌入维度

        Returns
        ----------
        - `int`, 输出图嵌入维度
        """
        pass
