import inspect
import warnings
from typing import (
    Any,
    Callable,
    Literal,
    Optional,
    Protocol,
    Sequence,
    Tuple,
    TypeAlias,
    cast,
)

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

# PyTorch Geometric imports
try:
    from torch_geometric.nn import (  # type: ignore
        global_add_pool,
        global_max_pool,
        global_mean_pool,
        knn_graph,
    )
    from torch_geometric.utils import add_self_loops, 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 ._base import BaseGraphEncoder, BaseGraphReadout, BaseGraphStructureLearner

TorchTensor: TypeAlias = 'torch.Tensor'


class _GraphConvLike(Protocol):
    def __call__(self, *args: Any, **kwargs: Any) -> TorchTensor: ...

    def forward(self, *args: Any, **kwargs: Any) -> TorchTensor: ...


class _GraphModelWithConvs(Protocol):
    convs: Sequence[_GraphConvLike]


# ============================================================================
# 图结构学习器 (Structure Learner)
# ============================================================================


class IdentityStructureLearner(BaseGraphStructureLearner):
    """恒等结构学习器

    不做任何处理，直接返回原始图结构。作为默认选项或消融实验的基线。
    """

    def forward(
        self,
        x: TorchTensor,
        edge_index: TorchTensor,
        edge_attr: Optional[TorchTensor] = None,
        batch: Optional[TorchTensor] = None,
    ) -> Tuple[TorchTensor, TorchTensor, Optional[TorchTensor]]:
        return x, edge_index, edge_attr


class EdgeThresholdStructureLearner(BaseGraphStructureLearner):
    """边权重阈值化结构学习器

    根据边权重进行阈值化，移除权重低于阈值的边。

    Parameters
    ----------
    - `threshold` : `float`, 阈值，含义取决于 mode
    - `mode` : `Literal['absolute', 'percentile']`, 阈值模式，'absolute': 绝对值阈值, 'percentile': 百分位阈值, by default 'absolute'
    """

    def __init__(
        self, threshold: float, mode: Literal['absolute', 'percentile'] = 'absolute'
    ):
        super().__init__()
        self.threshold = threshold
        self.mode = mode

    def forward(
        self,
        x: TorchTensor,
        edge_index: TorchTensor,
        edge_attr: Optional[TorchTensor] = None,
        batch: Optional[TorchTensor] = None,
        *args,
        **kwargs,
    ):
        if edge_attr is None:
            return x, edge_index, edge_attr

        weights = edge_attr.squeeze(-1) if edge_attr.dim() > 1 else edge_attr

        if self.mode == 'absolute':
            mask = weights.abs() >= self.threshold
        else:  # percentile
            cutoff = torch.quantile(weights.abs(), self.threshold / 100.0)
            mask = weights.abs() >= cutoff

        return x, edge_index[:, mask], edge_attr[mask] if edge_attr is not None else None


class CosSimStructureLearner(BaseGraphStructureLearner):
    """余弦相似度结构学习器

    基于节点特征余弦相似度构建图结构，保留相似度超过阈值的边。
    支持 batch 场景，会为每个图分别构建连接。

    Parameters
    ----------
    - `threshold` : `float`, 相似度阈值，范围 [0, 1], by default 0.5
    - `add_self_loop` : `bool`, 是否添加自环, by default True

    Examples
    ----------
    ```python
    learner = CosSimStructureLearner(threshold=0.5)
    x_new, edge_index_new, edge_attr_new = learner(x, edge_index, batch=batch)
    ```
    """

    def __init__(self, threshold: float = 0.5, add_self_loop: bool = True):
        super().__init__()
        if not 0 <= threshold <= 1:
            raise ValueError(f'threshold must be in [0, 1], got {threshold}')
        self.threshold = threshold
        self.add_self_loop = add_self_loop

    def _build_single_graph(self, x: TorchTensor, num_nodes: int) -> TorchTensor:
        """为单个图构建余弦相似度连接"""
        # 归一化特征
        x_norm = F.normalize(x, p=2, dim=-1)

        # 计算余弦相似度矩阵
        sim = torch.mm(x_norm, x_norm.t())

        # 阈值化
        adj = (sim > self.threshold).float()

        # 移除对角线 (自环单独处理)
        adj.fill_diagonal_(0)

        # 转换为稀疏格式
        edge_index = adj.nonzero(as_tuple=False).t().contiguous()

        # 添加自环
        if self.add_self_loop:
            edge_index, _ = add_self_loops(edge_index, num_nodes=num_nodes)

        return edge_index

    def forward(
        self,
        x: TorchTensor,
        edge_index: TorchTensor,
        edge_attr: Optional[TorchTensor] = None,
        batch: Optional[TorchTensor] = None,
    ) -> Tuple[TorchTensor, TorchTensor, Optional[TorchTensor]]:
        if batch is None:
            # 单图场景
            new_edge_index = self._build_single_graph(x, x.size(0))
        else:
            # batch 场景：为每个图分别构建
            edge_list = []
            for b in batch.unique():
                mask = batch == b
                x_sub = x[mask]
                num_nodes_sub = x_sub.size(0)

                # 构建局部连接
                edge_sub = self._build_single_graph(x_sub, num_nodes_sub)

                # 映射到全局索引
                global_nodes = mask.nonzero(as_tuple=True)[0]
                edge_sub = global_nodes[edge_sub]

                edge_list.append(edge_sub)

            new_edge_index = torch.cat(edge_list, dim=1)

        return x, new_edge_index, None


class KNNStructureLearner(BaseGraphStructureLearner):
    """K 近邻结构学习器

    基于 K 近邻构建图结构，在特征空间找每个节点的 K 个最近邻并连接。
    适用于 EEG 电极空间邻接关系建模等需要稀疏局部连接的场景。

    Parameters
    ----------
    - `k` : `int`, 近邻数量, by default 5
    - `loop` : `bool`, 是否包含自环, by default True
    - `flow` : `str`, 边方向，'source_to_target' 或 'target_to_source', by default 'source_to_target'
    - `force_symmetric` : `bool`, 是否强制对称化（添加反向边并去重）, by default False

    Examples
    ----------
    ```python
    learner = KNNStructureLearner(k=5, force_symmetric=True)
    x_new, edge_index_new, _ = learner(x, edge_index, batch=batch)
    ```
    """

    def __init__(
        self,
        k: int = 5,
        loop: bool = True,
        flow: str = 'source_to_target',
        force_symmetric: bool = False,
    ):
        super().__init__()
        if k < 1:
            raise ValueError(f'k must be >= 1, got {k}')
        self.k = k
        self.loop = loop
        self.flow = flow
        self.force_symmetric = force_symmetric

    def forward(
        self,
        x: TorchTensor,
        edge_index: TorchTensor,
        edge_attr: Optional[TorchTensor] = None,
        batch: Optional[TorchTensor] = None,
    ) -> Tuple[TorchTensor, TorchTensor, Optional[TorchTensor]]:
        # 使用 PyG 的 knn_graph
        new_edge_index = knn_graph(
            x, k=self.k, batch=batch, loop=self.loop, flow=self.flow
        )

        # 强制对称化
        if self.force_symmetric:
            from torch_geometric.utils import to_undirected  # pyright: ignore[reportMissingImports]

            new_edge_index = to_undirected(new_edge_index)

        return x, new_edge_index, None


class DropEdgeStructureLearner(BaseGraphStructureLearner):
    """DropEdge 结构学习器

    随机丢弃边，作为图结构的正则化手段。仅在训练模式下丢弃边，评估模式下保持原结构。

    Parameters
    ----------
    - `drop_prob` : `float`, 边丢弃概率, by default 0.1
    - `force_undirected` : `bool`, 是否强制保持无向图, by default True
    """

    def __init__(self, drop_prob: float = 0.1, force_undirected: bool = True):
        super().__init__()
        if not 0 <= drop_prob < 1:
            raise ValueError(f'drop_prob must be in [0, 1), got {drop_prob}')
        self.drop_prob = drop_prob
        self.force_undirected = force_undirected

    def forward(
        self,
        x: TorchTensor,
        edge_index: TorchTensor,
        edge_attr: Optional[TorchTensor] = None,
        *args,
        **kwargs,
    ) -> Tuple[TorchTensor, TorchTensor, Optional[TorchTensor]]:
        if self.training and self.drop_prob > 0:
            edge_index, mask = dropout_edge(
                edge_index, p=self.drop_prob, force_undirected=self.force_undirected
            )
            if edge_attr is not None:
                edge_attr = edge_attr[mask]

        return x, edge_index, edge_attr


# ============================================================================
# 图编码器 (Graph Encoder)
# ============================================================================


class GNNEncoder(BaseGraphEncoder):
    """通用 GNN 编码器包装器

    包装 PyTorch Geometric 的 GNN 模块 (GCN, GAT, GIN, GraphSAGE 等)。

    Parameters
    ----------
    - `gnn` : `nn.Module`, PyG GNN 模块实例
    - `output_dim` : `int`, 输出维度，必须与 gnn 的输出一致

    Examples
    ----------
    ```python
    from torch_geometric.nn import GCN
    gnn = GCN(in_channels=64, hidden_channels=128, out_channels=128, num_layers=3)
    encoder = GNNEncoder(gnn, output_dim=128)
    node_emb = encoder(x, edge_index)
    ```
    """

    def __init__(self, gnn: nn.Module, output_dim: int, force_edge_attr: bool = False):
        super().__init__()
        self.gnn = gnn
        self._sig = inspect.signature(self.gnn.forward)
        self._output_dim = output_dim
        self.force_edge_attr = force_edge_attr

    def forward(
        self,
        x: TorchTensor,
        edge_index: TorchTensor,
        edge_attr: Optional[TorchTensor] = None,
        batch: Optional[TorchTensor] = None,
    ) -> TorchTensor:
        # 尝试传递 edge_attr，如果 GNN 不支持则忽略
        if self.force_edge_attr and edge_attr is None:
            raise RuntimeError('edge_attr is required but not provided.')
        kwargs = {}
        """
        对标准 GCN/GIN：自动走 edge_weight 分支，你的 weighted adjacency 就真的参与 message passing 了；
        对只支持 edge_attr 的模型：也能吃到边特征；
        对完全不支持边特征的 GNN：如果你硬提供 edge_attr，就直接报错（而不是静默忽略），科研上更安全。
        """
        if edge_attr is not None:
            # squeeze 一下，兼容 from_adjacency 里加的 [...,1]
            edge_attr_1d = edge_attr.squeeze(-1) if edge_attr.dim() > 1 else edge_attr

            if 'edge_weight' in self._sig.parameters:
                kwargs['edge_weight'] = edge_attr_1d
            elif 'edge_attr' in self._sig.parameters:
                kwargs['edge_attr'] = edge_attr_1d
            else:
                raise TypeError(
                    f'{type(self.gnn).__name__} does not accept edge_weight/edge_attr, '
                    f'but edge_attr was provided.'
                )

        return self.gnn(x, edge_index, **kwargs)

    def get_output_dim(self) -> int:
        return self._output_dim

    def get_attention_weights(
        self,
        x: TorchTensor,
        edge_index: TorchTensor,
        edge_attr: Optional[TorchTensor] = None,
    ) -> Optional[Tuple[TorchTensor, TorchTensor]]:
        """获取 GAT 等模型的注意力权重

        仅支持具有 convs 属性的 PyG 模型，且最后一层必须支持 return_attention_weights 参数。
        由于重新运行了最后一层，如果层包含 Dropout，结果可能与 forward 略有不同。

        Returns
        ----------
        - `Optional[Tuple[TorchTensor, TorchTensor]]`, `(edge_index, attention_weights)` 或 None
        """
        try:
            gnn = self.gnn
            convs: Sequence[_GraphConvLike] | None = None
            if hasattr(gnn, 'convs'):
                convs_value = cast(_GraphModelWithConvs, gnn).convs
                if isinstance(convs_value, (nn.ModuleList, list, tuple)):
                    convs = convs_value

            # 检查是否有卷积层
            if convs is None or len(convs) == 0:
                return None

            last_conv = convs[-1]

            # 检查最后一层是否支持返回注意力权重
            # 通过检查函数签名判断
            import inspect

            sig = inspect.signature(last_conv.forward)
            if 'return_attention_weights' not in sig.parameters:
                return None

            # 获取中间层输出作为最后一层的输入
            # 这是一个 hack，但对可视化来说足够用
            with torch.no_grad():
                h = x
                for conv in convs[:-1]:
                    try:
                        h = conv(h, edge_index, edge_attr=edge_attr)
                    except TypeError:
                        h = conv(h, edge_index)
                    # 应用激活函数 (假设有 act 属性)
                    act = getattr(gnn, 'act', None)
                    if callable(act):
                        act = cast(Callable[[TorchTensor], TorchTensor], act)
                        h = act(h)

                # 最后一层带注意力权重
                _, att = last_conv(h, edge_index, return_attention_weights=True)
                return (edge_index, att)

        except Exception:
            # 任何错误都返回 None，不中断主流程
            return None


class MLPEncoder(BaseGraphEncoder):
    """MLP 编码器

    不使用图结构，仅对节点特征应用 MLP。用于消融实验，验证图结构的重要性。

    Parameters
    ----------
    - `in_channels` : `int`, 输入特征维度
    - `hidden_channels` : `int`, 隐藏层维度
    - `out_channels` : `int`, 输出维度
    - `num_layers` : `int`, MLP 层数, by default 2
    - `dropout` : `float`, Dropout 概率, by default 0.5
    """

    def __init__(
        self,
        in_channels: int,
        hidden_channels: int,
        out_channels: int,
        num_layers: int = 2,
        dropout: float = 0.5,
    ):
        super().__init__()
        self._output_dim = out_channels

        layers = []
        for i in range(num_layers):
            in_dim = in_channels if i == 0 else hidden_channels
            out_dim = out_channels if i == num_layers - 1 else hidden_channels
            layers.extend(
                [
                    nn.Linear(in_dim, out_dim),
                    nn.ReLU() if i < num_layers - 1 else nn.Identity(),
                    nn.Dropout(dropout) if i < num_layers - 1 else nn.Identity(),
                ]
            )

        self.mlp = nn.Sequential(*layers)

    def forward(
        self,
        x: TorchTensor,
        edge_index: TorchTensor,
        edge_attr: Optional[TorchTensor] = None,
        batch: Optional[TorchTensor] = None,
    ) -> TorchTensor:
        # 完全忽略图结构
        return self.mlp(x)

    def get_output_dim(self) -> int:
        return self._output_dim


# ============================================================================
# 图级池化 (Graph Readout)
# ============================================================================


class MeanReadout(BaseGraphReadout):
    """平均池化

    对图内所有节点特征取平均，最常用的池化方式，对节点数量不敏感。
    """

    def forward(self, x: TorchTensor, batch: Optional[TorchTensor] = None) -> TorchTensor:
        return global_mean_pool(x, batch)

    def get_output_dim(self, input_dim: int) -> int:
        return input_dim


class MaxReadout(BaseGraphReadout):
    """最大池化

    对每个特征维度取最大值，可能捕获最显著的特征。
    """

    def forward(self, x: TorchTensor, batch: Optional[TorchTensor] = None) -> TorchTensor:
        return global_max_pool(x, batch)

    def get_output_dim(self, input_dim: int) -> int:
        return input_dim


class AddReadout(BaseGraphReadout):
    """加和池化

    对图内所有节点特征求和，输出与节点数量成正比，可能需要归一化。
    """

    def forward(self, x: TorchTensor, batch: Optional[TorchTensor] = None) -> TorchTensor:
        return global_add_pool(x, batch)

    def get_output_dim(self, input_dim: int) -> int:
        return input_dim


class ConcatReadout(BaseGraphReadout):
    """拼接池化

    将 mean, max, add 三种池化结果拼接，输出维度是输入的 3 倍，但信息更丰富。

    适用于节点数基本一致的场景；如果节点数差异大，建议对 add 部分做归一化
    """

    def __init__(self, normalize_add: bool = True):
        super().__init__()
        self._normalize_add = normalize_add

    def enable_normalize_add(self, enable: bool = True):
        """启用或禁用对加和池化部分的归一化"""
        self._normalize_add = enable

    def forward(self, x: TorchTensor, batch: Optional[TorchTensor] = None) -> TorchTensor:
        mean_pool = global_mean_pool(x, batch)
        max_pool = global_max_pool(x, batch)
        add_pool = global_add_pool(x, batch)
        if self._normalize_add:
            if batch is None:
                add_pool = add_pool / x.size(0)
            else:
                try:
                    from torch_scatter import scatter_add  # pyright: ignore[reportMissingImports]
                except Exception as E:
                    raise ImportError(
                        'Please install torch-scatter to use ConcatReadout with normalization.'
                    )
                num_nodes = scatter_add(
                    torch.ones(x.size(0), device=x.device), batch, dim=0
                ).unsqueeze(-1)
                add_pool = add_pool / num_nodes

        return torch.cat([mean_pool, max_pool, add_pool], dim=-1)

    def get_output_dim(self, input_dim: int) -> int:
        return input_dim * 3


class AttentionReadout(BaseGraphReadout):
    """注意力池化

    学习节点重要性权重，加权聚合节点特征。
    通过 MLP 计算节点注意力分数，在图内应用 softmax 归一化后加权求和得到图表示。

    Parameters
    ----------
    - `hidden_dim` : `int`, 注意力 MLP 的隐藏层维度
    - `dropout` : `float`, Dropout 概率, by default 0.1
    """

    def __init__(self, hidden_dim: int, dropout: float = 0.1):
        super().__init__()
        self.hidden_dim = hidden_dim

        # 注意力打分网络
        self.attention = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.Tanh(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim // 2, 1),
        )

    def forward(self, x: TorchTensor, batch: Optional[TorchTensor] = None) -> TorchTensor:
        # 计算注意力分数
        attn_scores = self.attention(x).squeeze(-1)  # (num_nodes,)

        # 图内 softmax
        if batch is not None:
            from torch_geometric.utils import softmax  # pyright: ignore[reportMissingImports]

            attn_weights = softmax(attn_scores, batch)  # (num_nodes,)
        else:
            attn_weights = F.softmax(attn_scores, dim=0)

        # 加权聚合
        weighted = x * attn_weights.unsqueeze(-1)  # (num_nodes, hidden_dim)

        if batch is not None:
            return global_add_pool(weighted, batch)
        return weighted.sum(dim=0, keepdim=True)

    def get_output_dim(self, input_dim: int) -> int:
        return input_dim
