import inspect
import warnings
from typing import Any, Callable, Dict, Optional, Protocol, Tuple, TypeAlias, Union, cast

import torch
import torch.nn as nn
import torch.nn.functional as F

# PyTorch Geometric imports
try:
    from torch_geometric.data import Batch, Data  # type: ignore
    from torch_geometric.nn import GAT, GCN, GIN, GraphSAGE  # type: ignore
    from torch_geometric.utils import dropout_edge  # type: ignore

    HAS_PYG = True
except ImportError:
    HAS_PYG = False
    warnings.warn(
        'PyTorch Geometric not found. Install with: pip install torch-geometric',
        ImportWarning,
    )

from zyplib.nn.modules import BaseModule
from zyplib.train.common import FabricTrainingContext

from ._base import (
    BaseGraphEncoder,
    BaseGraphReadout,
    BaseGraphStructureLearner,
    GNNType,
    GraphClassifierConfig,
    GraphOutput,
    PoolingType,
)
from .layers import (
    AddReadout,
    AttentionReadout,
    ConcatReadout,
    CosSimStructureLearner,
    DropEdgeStructureLearner,
    EdgeThresholdStructureLearner,
    GNNEncoder,
    IdentityStructureLearner,
    KNNStructureLearner,
    MaxReadout,
    MeanReadout,
)

TorchTensor: TypeAlias = 'torch.Tensor'


class _GraphBatchLike(Protocol):
    x: TorchTensor
    edge_index: TorchTensor
    edge_attr: TorchTensor | None
    batch: TorchTensor | None
    y: TorchTensor


class _OptimizerLike(Protocol):
    def zero_grad(self) -> None: ...

    def step(self) -> None: ...


# ============================================================================
# 图分类器
# ============================================================================


class GraphClassifier(BaseModule):
    """统一的图分类器

    组件化设计，支持灵活配置结构学习器、编码器、读出和分类头。
    forward() 保持简洁，可视化逻辑隔离到 get_attention_weights()，辅助损失通过 aux_loss_fn 回调支持。

    Parameters
    ----------
    - `encoder` : `GraphEncoder`, 图编码器
    - `readout` : `GraphReadout`, 图级池化
    - `config` : `GraphClassifierConfig`, 分类器配置
    - `structure_learner` : `Optional[GraphStructureLearner]`, 结构学习器，默认 IdentityStructureLearner, by default None
    - `fn_main_loss` : `Optional[Callable]`, 主损失函数，签名: `(logits, labels) -> TorchTensor`, by default None
        - 注意如果传入, 则 config 中的 label_smoothing 和 class_weights 不会生效
    - `fn_aux_loss` : `Optional[Callable]`, 辅助损失函数，签名: `(output, data) -> TorchTensor`, by default None

    Examples
    ----------
    ```python
    # 基础用法
    model = GraphClassifier(
        encoder=GNNEncoder(GCN(...), output_dim=128),
        readout=MeanReadout(),
        config=GraphClassifierConfig(num_classes=4, hidden_dim=128),
    )
    output = model(batch)
    loss_dict = model.compute_loss(batch)

    # 高级用法：动态 KNN 图 + 注意力池化 + 对比损失
    def contrastive_loss(output, data):
        # 实现你的对比损失
        return loss

    model = GraphClassifier(
        structure_learner=KNNStructureLearner(k=5),
        encoder=GNNEncoder(GAT(...), output_dim=128),
        readout=AttentionReadout(hidden_dim=128),
        config=GraphClassifierConfig(num_classes=4, hidden_dim=128),
        aux_loss_fn=contrastive_loss,
    )
    ```
    """

    def __init__(
        self,
        encoder: BaseGraphEncoder,
        readout: BaseGraphReadout,
        config: GraphClassifierConfig,
        structure_learner: Optional[BaseGraphStructureLearner] = None,
        fn_main_loss: Optional[Callable[[TorchTensor, TorchTensor], TorchTensor]] = None,
        fn_aux_loss: Optional[
            Callable[[GraphOutput, Union[Data, Batch]], TorchTensor]
        ] = None,
    ):
        super().__init__()

        # 保存配置
        self.config = config

        # 组件
        self.structure_learner = structure_learner or IdentityStructureLearner()
        self.encoder = encoder
        self.readout = readout

        self.fn_main_loss = fn_main_loss
        self.fn_aux_loss = fn_aux_loss

        # 构建分类头
        readout_dim = readout.get_output_dim(encoder.get_output_dim())
        hidden = config.classifier_hidden_dim or config.hidden_dim

        self.classifier = nn.Sequential(
            nn.Linear(readout_dim, hidden),
            nn.ReLU(),
            nn.Dropout(config.classifier_dropout),
            nn.Linear(hidden, config.num_classes),
        )

        # 类别权重 (用于 compute_loss)
        if config.class_weights is not None:
            self.register_buffer('class_weights', config.class_weights)
        else:
            self.class_weights = None

        # Check
        if isinstance(self.structure_learner, DropEdgeStructureLearner):
            if config.drop_edge_prob > 0:
                warnings.warn(
                    'DropEdgeStructureLearner is used along with drop_edge_prob in config. This may lead to double dropping of edges.'
                )

    def forward(
        self, data: Union[Data, Batch], return_node_emb: bool = False
    ) -> GraphOutput:
        """前向传播：结构学习 → 编码 → 池化 → 分类

        Parameters
        ----------
        - `data` : `Union[Data, Batch]`, PyG Data 或 Batch 对象
        - `return_node_emb` : `bool`, 是否返回节点嵌入, by default False

        Returns
        ----------
        - `GraphOutput`, 包含 logits, graph_emb, (可选) node_emb
        """
        # 解包数据
        graph_data = cast(_GraphBatchLike, data)
        x, edge_index = graph_data.x, graph_data.edge_index
        batch = getattr(graph_data, 'batch', None)
        edge_attr = getattr(graph_data, 'edge_attr', None)

        # 处理单图场景（batch=None）
        if batch is None:
            batch = x.new_zeros(x.size(0), dtype=torch.long)

        # 1. 结构学习 (预处理或动态学习)
        x, edge_index, edge_attr = self.structure_learner(x, edge_index, edge_attr, batch)

        # 2. 训练时 DropEdge 正则化 (如果配置了)
        if self.training and self.config.drop_edge_prob > 0:
            edge_index, mask = dropout_edge(
                edge_index, p=self.config.drop_edge_prob, force_undirected=True
            )
            if edge_attr is not None:
                edge_attr = edge_attr[mask]

        # 3. 节点编码
        node_emb = self.encoder(x, edge_index, edge_attr, batch)

        # 4. 图级池化
        graph_emb = self.readout(node_emb, batch)

        # 5. 分类
        logits = self.classifier(graph_emb)

        return GraphOutput(
            logits=logits,
            graph_emb=graph_emb,
            node_emb=node_emb if return_node_emb else None,
        )

    def compute_loss(
        self,
        logits: TorchTensor,
        labels: TorchTensor,
        output: Optional[GraphOutput] = None,
        data: Optional[Union[Data, Batch]] = None,
    ) -> Dict[str, TorchTensor]:
        """计算损失

        支持主损失（交叉熵，带标签平滑和类别权重）和辅助损失（通过 fn_aux_loss 回调）。

        Parameters
        ----------
        - `logits` : `TorchTensor`, 模型输出的 logits, shape: `(batch_size, num_classes)`
        - `labels` : `TorchTensor`, 真实标签, shape: `(batch_size,)`
        - `output` : `Optional[GraphOutput]`, 完整模型输出，用于计算辅助损失, by default None
        - `data` : `Optional[Union[Data, Batch]]`, 原始数据，用于计算辅助损失, by default None

        Returns
        ----------
        - `Dict[str, TorchTensor]`, 损失字典: `{'ce_loss', 'aux_loss' (可选), 'total_loss'}`
        """
        if labels.dim() > 1:
            labels = labels.view(-1)
        if labels.dtype != torch.long:
            labels = labels.long()

        # 主损失
        if self.fn_main_loss is not None:
            main_loss = self.fn_main_loss(logits, labels)
        else:
            main_loss = F.cross_entropy(
                logits,
                labels,
                weight=self.class_weights,
                label_smoothing=self.config.label_smoothing,
            )

        total_loss = main_loss
        loss_dict = {'ce_loss': main_loss}

        # 辅助损失
        if self.fn_aux_loss is not None:
            if output is None or data is None:
                raise ValueError(
                    'fn_aux_loss requires both output and data, but got None'
                )
            aux_loss = self.fn_aux_loss(output, data)
            loss_dict['aux_loss'] = aux_loss
            total_loss = total_loss + aux_loss

        loss_dict['total_loss'] = total_loss
        return loss_dict

    @torch.no_grad()
    def predict(self, data: Union[Data, Batch]) -> TorchTensor:
        """预测类别

        Parameters
        ----------
        - `data` : `Union[Data, Batch]`, PyG Data 或 Batch

        Returns
        ----------
        - `TorchTensor`, 预测类别, shape: `(batch_size,)`
        """
        self.eval()
        return self.forward(data).predictions

    @torch.no_grad()
    def predict_proba(self, data: Union[Data, Batch]) -> TorchTensor:
        """预测概率

        Parameters
        ----------
        - `data` : `Union[Data, Batch]`, PyG Data 或 Batch

        Returns
        ----------
        - `TorchTensor`, 预测概率, shape: `(batch_size, num_classes)`
        """
        self.eval()
        return self.forward(data).probabilities

    @torch.no_grad()
    def get_attention_weights(self, data: Union[Data, Batch]) -> Optional[Any]:
        """获取注意力权重，用于可视化

        专门的方法隔离可视化逻辑，委托给 encoder 处理。

        Parameters
        ----------
        - `data` : `Union[Data, Batch]`, PyG Data 或 Batch

        Returns
        ----------
        - `Optional[Any]`, 注意力权重，格式取决于 encoder 实现。如果 encoder 不支持，返回 None
        """
        self.eval()

        # 首先应用结构学习 (保持与 forward 一致)
        graph_data = cast(_GraphBatchLike, data)
        x, edge_index = graph_data.x, graph_data.edge_index
        edge_attr = getattr(graph_data, 'edge_attr', None)
        batch = getattr(graph_data, 'batch', None)

        # 与 forward 保持同样的 batch 处理逻辑
        if batch is None:
            batch = x.new_zeros(x.size(0), dtype=torch.long)
        x, edge_index, edge_attr = self.structure_learner(x, edge_index, edge_attr, batch)

        # 委托给 encoder（目前 encoder.get_attention_weights 不使用 batch）
        return self.encoder.get_attention_weights(x, edge_index, edge_attr)

    def get_embeddings(self, data: Union[Data, Batch]) -> Tuple[TorchTensor, TorchTensor]:
        """获取节点和图嵌入，用于下游任务

        Parameters
        ----------
        - `data` : `Union[Data, Batch]`, PyG Data 或 Batch

        Returns
        ----------
        - `Tuple[TorchTensor, TorchTensor]`, `(node_embeddings, graph_embeddings)`
        """
        output = self.forward(data, return_node_emb=True)
        return cast(Tuple[TorchTensor, TorchTensor], (output.node_emb, output.graph_emb))


# ============================================================================
# 训练流程
# ============================================================================


class GraphClassificationProcedure:
    """与 zyplib 训练框架集成的训练流程"""

    def train_step(
        self, ctx: 'FabricTrainingContext', batch: Batch, **kwargs
    ) -> Tuple[TorchTensor, Dict[str, float]]:
        model = cast(GraphClassifier, ctx.model)
        optimizer = cast(_OptimizerLike, ctx.optimizer)

        model.train()
        optimizer.zero_grad()

        # Forward 一次
        output = model(batch, return_node_emb=True)
        # 计算损失
        loss_dict = model.compute_loss(
            logits=output.logits,
            labels=cast(_GraphBatchLike, batch).y,
            output=output,
            data=batch,
        )

        if ctx.fabric:
            ctx.fabric.backward(loss_dict['total_loss'])
        else:
            loss_dict['total_loss'].backward()

        optimizer.step()

        return output.logits, {k: v.item() for k, v in loss_dict.items()}

    def val_step(
        self, ctx: 'FabricTrainingContext', batch: Batch, **kwargs
    ) -> Tuple[TorchTensor, Dict[str, float]]:
        model = cast(GraphClassifier, ctx.model)
        model.eval()

        with torch.no_grad():
            # Forward 一次
            output = model(batch, return_node_emb=True)
            # 计算损失
            loss_dict = model.compute_loss(
                logits=output.logits,
                labels=cast(_GraphBatchLike, batch).y,
                output=output,
                data=batch,
            )

        return output.logits, {k: v.item() for k, v in loss_dict.items()}


# ============================================================================
# 工厂函数
# ============================================================================


def create_gnn(
    gnn_type: Union[str, GNNType],
    in_channels: int,
    hidden_channels: int,
    out_channels: int,
    num_layers: int = 3,
    dropout: float = 0.5,
    **kwargs,
) -> nn.Module:
    """创建 PyG GNN 模块

    Parameters
    ----------
    - `gnn_type` : `Union[str, GNNType]`, GNN 类型 ('gcn', 'gat', 'gin', 'sage')
    - `in_channels` : `int`, 输入通道数
    - `hidden_channels` : `int`, 隐藏层通道数
    - `out_channels` : `int`, 输出通道数
    - `num_layers` : `int`, 层数, by default 3
    - `dropout` : `float`, Dropout 概率, by default 0.5
    - `**kwargs`, 传递给具体 GNN 的额外参数

    Returns
    ----------
    - `nn.Module`, PyG GNN 模块
    """
    if not HAS_PYG:
        raise ImportError(
            'PyG is not installed. Please install torch-geometric to use GNN modules.'
        )

    if isinstance(gnn_type, str):
        gnn_type = GNNType(gnn_type.lower())

    gnn_classes = {
        GNNType.GCN: GCN,
        GNNType.GAT: GAT,
        GNNType.GIN: GIN,
        GNNType.SAGE: GraphSAGE,
    }

    gnn_class = gnn_classes.get(gnn_type)
    if gnn_class is None:
        raise ValueError(f'Unknown GNN type: {gnn_type}')

    return gnn_class(
        in_channels=in_channels,
        hidden_channels=hidden_channels,
        out_channels=out_channels,
        num_layers=num_layers,
        dropout=dropout,
        **kwargs,
    )


def create_readout(
    pooling_type: Union[str, PoolingType], hidden_dim: Optional[int] = None, **kwargs
) -> BaseGraphReadout:
    """创建图级池化模块

    Parameters
    ----------
    - `pooling_type` : `Union[str, PoolingType]`, 池化类型
    - `hidden_dim` : `Optional[int]`, 隐藏维度，Attention 池化需要, by default None
    - `**kwargs`, 传递给具体 Readout 的额外参数

    Returns
    ----------
    - `GraphReadout`, GraphReadout 实例
    """
    if isinstance(pooling_type, str):
        pooling_type = PoolingType(pooling_type.lower())

    readout_map = {
        PoolingType.MEAN: MeanReadout,
        PoolingType.MAX: MaxReadout,
        PoolingType.ADD: AddReadout,
        PoolingType.CONCAT: ConcatReadout,
        PoolingType.ATTENTION: AttentionReadout,
    }

    readout_class = readout_map.get(pooling_type)
    if readout_class is None:
        raise ValueError(f'Unknown pooling type: {pooling_type}')

    if pooling_type == PoolingType.ATTENTION:
        if hidden_dim is None:
            raise ValueError('hidden_dim is required for AttentionReadout')
        return readout_class(hidden_dim=hidden_dim, **kwargs)

    return readout_class(**kwargs)


def create_graph_classifier(
    gnn_type: Union[str, GNNType],
    in_channels: int,
    hidden_channels: int,
    num_classes: int,
    num_layers: int = 3,
    pooling: Union[str, PoolingType] = PoolingType.MEAN,
    structure_learner: Optional[BaseGraphStructureLearner] = None,
    aux_loss_fn: Optional[Callable] = None,
    dropout: float = 0.5,
    force_edge_attr: bool = False,
    **kwargs,
) -> GraphClassifier:
    """创建图分类器

    便捷工厂函数，封装了 GNN、Encoder、Readout 和 Config 的创建。

    Parameters
    ----------
    - `gnn_type` : `Union[str, GNNType]`, GNN 类型 ('gcn', 'gat', 'gin', 'sage')
    - `in_channels` : `int`, 输入特征维度
    - `hidden_channels` : `int`, 隐藏层维度
    - `num_classes` : `int`, 分类类别数
    - `num_layers` : `int`, GNN 层数, by default 3
    - `pooling` : `Union[str, PoolingType]`, 池化类型, by default PoolingType.MEAN
    - `structure_learner` : `Optional[GraphStructureLearner]`, 结构学习器, by default None
    - `aux_loss_fn` : `Optional[Callable]`, 辅助损失函数, by default None
    - `dropout` : `float`, Dropout 概率, by default 0.5
    - `**kwargs`, 传递给 GNN 的额外参数

    Returns
    ----------
    - `GraphClassifier`, GraphClassifier 实例

    Examples
    ----------
    ```python
    # 简单用法
    model = create_graph_classifier(
        'gcn', in_channels=64, hidden_channels=128, num_classes=4
    )

    # 高级用法
    model = create_graph_classifier(
        'gat',
        in_channels=64,
        hidden_channels=128,
        num_classes=4,
        pooling='attention',
        structure_learner=KNNStructureLearner(k=5),
        heads=4,  # GAT 特有参数
    )
    ```
    """
    # 创建 GNN
    gnn = create_gnn(
        gnn_type=gnn_type,
        in_channels=in_channels,
        hidden_channels=hidden_channels,
        out_channels=hidden_channels,
        num_layers=num_layers,
        dropout=dropout,
        **kwargs,
    )

    # 创建 Encoder
    encoder = GNNEncoder(gnn, output_dim=hidden_channels, force_edge_attr=force_edge_attr)

    # 创建 Readout
    readout = create_readout(pooling, hidden_dim=hidden_channels)

    # 创建配置
    config = GraphClassifierConfig(
        num_classes=num_classes, hidden_dim=hidden_channels, classifier_dropout=dropout
    )

    return GraphClassifier(
        encoder=encoder,
        readout=readout,
        config=config,
        structure_learner=structure_learner,
        fn_aux_loss=aux_loss_fn,
    )
