import warnings
from typing import (
    TYPE_CHECKING,
    Callable,
    List,
    Literal,
    Optional,
    Sequence,
    Tuple,
    Type,
    TypeAlias,
    Union,
)

import numpy as np
import torch
import torch.nn.functional as F

# PyTorch Geometric imports
try:
    from torch_geometric.data import Data  # type: ignore
    from torch_geometric.utils import dense_to_sparse, remove_self_loops  # type: ignore

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

TorchTensor: TypeAlias = 'torch.Tensor'

# ============================================================================
# 工具类
# ============================================================================


class AdjacencyUtils:
    """邻接矩阵工具箱

    提供邻接矩阵的归一化、阈值化、拉普拉斯计算、KNN 图构建、距离图构建和稀疏转换功能。
    """

    @staticmethod
    def _ensure_TorchTensor(x: Union[np.ndarray, TorchTensor]) -> TorchTensor:
        """内部辅助：统一转换为 Float TorchTensor"""
        if isinstance(x, np.ndarray):
            return torch.from_numpy(x).float()
        return x.float()

    @staticmethod
    def normalize(
        adj: Union[np.ndarray, TorchTensor],
        mode: Literal['symmetric', 'rw', 'none'] = 'symmetric',
        add_self_loop: bool = True,
    ) -> TorchTensor:
        """归一化邻接矩阵

        Parameters
        ----------
        - `adj` : `Union[np.ndarray, TorchTensor]`, 邻接矩阵
        - `mode` : `Literal['symmetric', 'rw', 'none']`, 归一化模式，'symmetric': D^{-1/2} A D^{-1/2}, 'rw': D^{-1} A, 'none': 不归一化, by default 'symmetric'
        - `add_self_loop` : `bool`, 是否添加自环, by default True

        Returns
        ----------
        - `TorchTensor`, 归一化后的邻接矩阵
        """
        adj = AdjacencyUtils._ensure_TorchTensor(adj)

        if add_self_loop:
            adj = adj + torch.eye(adj.size(0), device=adj.device)

        deg = adj.sum(dim=1)

        if mode == 'none':
            return adj

        elif mode == 'rw':  # Random Walk: D^{-1} A (v2中的 'row' 模式)
            deg_inv = deg.pow(-1)
            deg_inv[deg_inv == float('inf')] = 0
            return deg_inv.unsqueeze(1) * adj

        elif mode == 'symmetric':  # D^{-1/2} A D^{-1/2}
            deg_inv_sqrt = deg.pow(-0.5)
            deg_inv_sqrt[deg_inv_sqrt == float('inf')] = 0
            return deg_inv_sqrt.unsqueeze(1) * adj * deg_inv_sqrt.unsqueeze(0)

        else:
            raise ValueError(f'Unknown normalization mode: {mode}')

    @staticmethod
    def threshold(
        adj: Union[np.ndarray, TorchTensor],
        value: float,
        mode: Literal['absolute', 'percentile', 'topk'] = 'absolute',
        binary: bool = False,
    ) -> TorchTensor:
        """邻接矩阵阈值化

        Parameters
        ----------
        - `adj` : `Union[np.ndarray, TorchTensor]`, 邻接矩阵
        - `value` : `float`, 阈值，含义取决于 mode
        - `mode` : `Literal['absolute', 'percentile', 'topk']`, 阈值模式，'absolute': 绝对值阈值, 'percentile': 百分位阈值, 'topk': 每个节点保留 K 个最强邻居, by default 'absolute'
        - `binary` : `bool`, 是否二值化, by default False

        Returns
        ----------
        - `TorchTensor`, 阈值化后的邻接矩阵
        """
        adj = AdjacencyUtils._ensure_TorchTensor(adj).clone()

        if mode == 'absolute':
            mask = adj.abs() >= value
        elif mode == 'percentile':
            # 保留前 p% 的强连接
            cutoff = torch.quantile(adj.abs().flatten(), value / 100.0)
            mask = adj.abs() >= cutoff
        elif mode == 'topk':
            # 每个节点保留 K 个最强邻居
            k = int(value)
            _, indices = adj.abs().topk(k, dim=1)
            mask = torch.zeros_like(adj, dtype=torch.bool)
            mask.scatter_(1, indices, True)
            # 保证对称性 (可选)
            mask = mask | mask.t()
        else:
            raise ValueError(f'Unknown threshold mode: {mode}')

        result = adj * mask.float()
        if binary:
            result = (result != 0).float()

        return result

    @staticmethod
    def laplacian(
        adj: Union[np.ndarray, TorchTensor], normalized: bool = True
    ) -> TorchTensor:
        """计算拉普拉斯矩阵

        Parameters
        ----------
        - `adj` : `Union[np.ndarray, TorchTensor]`, 邻接矩阵
        - `normalized` : `bool`, 是否归一化， True: L_norm = I - D^{-1/2} A D^{-1/2}, False: L = D - A, by default True

        Returns
        ----------
        - `TorchTensor`, 拉普拉斯矩阵
        """
        adj = AdjacencyUtils._ensure_TorchTensor(adj)

        if normalized:
            # L_norm = I - D^{-1/2} A D^{-1/2}
            adj_norm = AdjacencyUtils.normalize(
                adj, mode='symmetric', add_self_loop=False
            )
            return torch.eye(adj.size(0), device=adj.device) - adj_norm
        else:
            # L = D - A
            deg = adj.sum(dim=1)
            return torch.diag(deg) - adj

    @staticmethod
    def knn_graph(
        features: Union[np.ndarray, TorchTensor],
        k: int,
        metric: Literal['euclidean', 'cosine'] = 'cosine',
        symmetric: bool = True,
    ) -> TorchTensor:
        """基于特征构建 KNN 图

        Parameters
        ----------
        - `features` : `Union[np.ndarray, TorchTensor]`, 节点特征
        - `k` : `int`, 近邻数量
        - `metric` : `Literal['euclidean', 'cosine']`, 距离度量, by default 'cosine'
        - `symmetric` : `bool`, 是否对称化, by default True

        Returns
        ----------
        - `TorchTensor`, 邻接矩阵
        """
        features = AdjacencyUtils._ensure_TorchTensor(features)

        if metric == 'cosine':
            features_norm = F.normalize(features, p=2, dim=1)
            sim = features_norm @ features_norm.t()
        else:
            sim = -torch.cdist(features, features)

        sim.fill_diagonal_(float('-inf'))
        _, indices = sim.topk(k, dim=1)

        adj = torch.zeros_like(sim)
        adj.scatter_(1, indices, 1.0)

        if symmetric:
            adj = ((adj + adj.t()) > 0).float()

        return adj

    @staticmethod
    def distance_graph(
        positions: Union[np.ndarray, TorchTensor],
        threshold: Optional[float] = None,
        sigma: Optional[float] = None,
    ) -> TorchTensor:
        """基于物理距离构建图

        Parameters
        ----------
        - `positions` : `Union[np.ndarray, TorchTensor]`, 节点位置
        - `threshold` : `Optional[float]`, 距离阈值, by default None
        - `sigma` : `Optional[float]`, 高斯核参数，如果设置则使用高斯核转换, by default None

        Returns
        ----------
        - `TorchTensor`, 邻接矩阵
        """
        positions = AdjacencyUtils._ensure_TorchTensor(positions)
        dist = torch.cdist(positions, positions)

        if sigma is not None:
            # 高斯核转换
            adj = torch.exp(-(dist**2) / (2 * sigma**2))
            if threshold is not None:
                adj[adj < threshold] = 0
        elif threshold is not None:
            # 硬截断
            adj = (dist <= threshold).float()
        else:
            adj = dist.clone()

        adj.fill_diagonal_(0)
        return adj

    @staticmethod
    def to_sparse(adj: TorchTensor) -> Tuple[TorchTensor, TorchTensor]:
        """转为 PyG 需要的 edge_index 格式

        Parameters
        ----------
        - `adj` : `TorchTensor`, 邻接矩阵

        Returns
        ----------
        - `Tuple[TorchTensor, TorchTensor]`, `(edge_index, edge_weight)`

        Note
        ----------
        应先调用 threshold() 处理后再转稀疏。
        """
        # 注意：应先调用 threshold() 处理后再转稀疏
        edge_index = adj.nonzero(as_tuple=False).t().contiguous()
        edge_weight = adj[edge_index[0], edge_index[1]]
        return edge_index, edge_weight


class GraphBuilder:
    """统一图构建器

    提供 PyG Data 构建、特征处理、信号相关性计算功能。
    """

    @staticmethod
    def from_adjacency(
        adj: Union[np.ndarray, TorchTensor],
        node_features: Union[np.ndarray, TorchTensor],
        label: Optional[int] = None,
        threshold_val: Optional[float] = None,
        threshold_mode: Literal['absolute', 'percentile', 'topk'] = 'absolute',
        keep_self_loops: bool = False,
        edge_attr: bool = True,
    ) -> Data:
        """从邻接矩阵构建 PyG Data

        Parameters
        ----------
        - `adj` : `Union[np.ndarray, TorchTensor]`, 邻接矩阵
        - `node_features` : `Union[np.ndarray, TorchTensor]`, 节点特征
        - `label` : `Optional[int]`, 图标签, by default None
        - `threshold_val` : `Optional[float]`, 阈值, by default None
        - `threshold_mode` : `Literal['absolute', 'percentile', 'topk']`, 阈值模式, by default 'absolute'
        - `keep_self_loops` : `bool`, 是否保留自环, by default False
        - `edge_attr` : `bool`, 是否保存边权重为边特征, by default True

        Returns
        ----------
        - `Data`, PyG Data 对象
        """
        adj = AdjacencyUtils._ensure_TorchTensor(adj)
        x = AdjacencyUtils._ensure_TorchTensor(node_features)

        if x.dim() == 1:
            x = x.unsqueeze(-1)

        # 1. 阈值化处理
        if threshold_val is not None and threshold_val > 0:
            adj = AdjacencyUtils.threshold(adj, threshold_val, mode=threshold_mode)

        # 2. 转换为稀疏格式
        edge_index, edge_weight = dense_to_sparse(adj)

        # 3. 自环处理
        if not keep_self_loops:
            edge_index, edge_weight = remove_self_loops(edge_index, edge_weight)

        # 4. 组装 Data
        data = Data(x=x, edge_index=edge_index, num_nodes=x.size(0))

        if edge_attr and edge_weight is not None:
            data.edge_attr = (
                edge_weight.unsqueeze(-1) if edge_weight.dim() == 1 else edge_weight
            )

        if label is not None:
            # data.y = torch.Tensor([label], dtype=torch.long)
            data.y = torch.tensor(label, dtype=torch.long).view(-1)

        return data

    @staticmethod
    def from_features(
        features: Union[np.ndarray, TorchTensor],
        adjacency: Optional[Union[np.ndarray, TorchTensor]] = None,
        label: Optional[int] = None,
        threshold: float = 0.0,
    ) -> Data:
        """从特征构建 PyG Data

        如果不提供邻接矩阵，则构建全连接图。

        Parameters
        ----------
        - `features` : `Union[np.ndarray, TorchTensor]`, 节点特征
        - `adjacency` : `Optional[Union[np.ndarray, TorchTensor]]`, 邻接矩阵, by default None
        - `label` : `Optional[int]`, 图标签, by default None
        - `threshold` : `float`, 阈值, by default 0.0

        Returns
        ----------
        - `Data`, PyG Data 对象
        """
        features = AdjacencyUtils._ensure_TorchTensor(features)
        num_nodes = features.size(0)

        if adjacency is None:
            # 全连接图构建 (优化版)
            row = torch.arange(num_nodes).repeat_interleave(num_nodes)
            col = torch.arange(num_nodes).repeat(num_nodes)
            edge_index = torch.stack([row, col], dim=0)
            # 移除自环
            mask = row != col
            edge_index = edge_index[:, mask]

            data = Data(x=features, edge_index=edge_index, num_nodes=num_nodes)
            if label is not None:
                # data.y = torch.Tensor([label], dtype=torch.long)
                data.y = torch.tensor(label, dtype=torch.long).view(-1)
            return data
        else:
            return GraphBuilder.from_adjacency(
                adj=adjacency,
                node_features=features,
                label=label,
                threshold_val=threshold,
                threshold_mode='absolute',  # 默认简单模式
            )

    @staticmethod
    def from_correlation(
        signals: Union[np.ndarray, TorchTensor],
        label: Optional[int] = None,
        threshold: float = 0.3,
        use_abs: bool = True,
    ) -> Data:
        """从信号相关性构建 PyG Data

        计算 Pearson 相关系数作为功能连接，并将原始信号作为节点特征。

        Parameters
        ----------
        - `signals` : `Union[np.ndarray, TorchTensor]`, 信号矩阵, shape: `(Num_Channels, Time_Points)`
        - `label` : `Optional[int]`, 图标签, by default None
        - `threshold` : `float`, 相关性阈值, by default 0.3
        - `use_abs` : `bool`, 是否使用绝对值, by default True

        Returns
        ----------
        - `Data`, PyG Data 对象
        """
        signals = AdjacencyUtils._ensure_TorchTensor(signals)

        # 1. 归一化并计算 Pearson 相关系数
        mean = signals.mean(dim=1, keepdim=True)
        std = signals.std(dim=1, keepdim=True) + 1e-8
        signals_norm = (signals - mean) / std

        corr = torch.mm(signals_norm, signals_norm.t()) / signals.size(1)

        if use_abs:
            corr = corr.abs()

        # 2. 移除自相关 (对角线)
        corr.fill_diagonal_(0)

        # 3. 复用通用构建逻辑
        return GraphBuilder.from_adjacency(
            adj=corr,
            node_features=signals,  # 将原始信号作为节点特征
            label=label,
            threshold_val=threshold,
        )

    @staticmethod
    def from_connectivity(
        signals: Union[np.ndarray, TorchTensor],
        method: Union[str, Callable] = 'pearson',
        label: Optional[int] = None,
        threshold: float = 0.3,
        use_abs: bool = True,
        **method_kwargs,
    ) -> Data:
        """从信号连接性构建 PyG Data

        支持多种连接性度量方法，包括 Pearson 相关、余弦相似度等。
        也可以传入自定义函数计算连接性矩阵。

        Parameters
        ----------
        - `signals` : `Union[np.ndarray, TorchTensor]`, 信号矩阵, shape: `(Num_Channels, Time_Points)`
        - `method` : `Union[str, Callable]`, 连接性度量方法，可选 'pearson', 'cosine'，或自定义函数, by default 'pearson'
        - `label` : `Optional[int]`, 图标签, by default None
        - `threshold` : `float`, 连接性阈值, by default 0.3
        - `use_abs` : `bool`, 是否使用绝对值, by default True
        - `**method_kwargs`, 传递给自定义函数的额外参数

        Returns
        ----------
        - `Data`, PyG Data 对象

        Examples
        ----------
        ```python
        # 使用内置方法
        data = GraphBuilder.from_connectivity(signals, method='pearson', threshold=0.3)

        # 使用自定义函数
        def my_connectivity(signals):
            # 计算自定义连接性矩阵
            return connectivity_matrix

        data = GraphBuilder.from_connectivity(signals, method=my_connectivity)
        ```

        Note
        ----------
        对于更复杂的 EEG 连接性度量（PLV、PLI、coherence 等），建议先在外部计算邻接矩阵，
        然后使用 `from_adjacency` 构建图。
        """
        signals = AdjacencyUtils._ensure_TorchTensor(signals)

        if callable(method):
            # 自定义函数
            connectivity = method(signals, **method_kwargs)
        elif method == 'pearson':
            # Pearson 相关系数
            mean = signals.mean(dim=1, keepdim=True)
            std = signals.std(dim=1, keepdim=True) + 1e-8
            signals_norm = (signals - mean) / std
            connectivity = torch.mm(signals_norm, signals_norm.t()) / signals.size(1)
        elif method == 'cosine':
            # 余弦相似度
            signals_norm = F.normalize(signals, p=2, dim=1)
            connectivity = torch.mm(signals_norm, signals_norm.t())
        else:
            raise ValueError(
                f"Unknown method: {method}. Supported: 'pearson', 'cosine', or callable."
            )

        if use_abs:
            connectivity = connectivity.abs()

        # 移除自连接
        connectivity.fill_diagonal_(0)

        # 构建图
        return GraphBuilder.from_adjacency(
            adj=connectivity, node_features=signals, label=label, threshold_val=threshold
        )

    @classmethod
    def build_dataset(
        cls,
        adjacencies: Sequence[Union[np.ndarray, TorchTensor]],
        features: Sequence[Union[np.ndarray, TorchTensor]],
        labels: Sequence[int],
        **kwargs,
    ) -> List[Data]:
        """批量构建数据集

        Parameters
        ----------
        - `adjacencies` : `Sequence[Union[np.ndarray, TorchTensor]]`, 邻接矩阵列表
        - `features` : `Sequence[Union[np.ndarray, TorchTensor]]`, 节点特征列表
        - `labels` : `Sequence[int]`, 图标签列表
        - `**kwargs`, 传递给 from_adjacency 的额外参数

        Returns
        ----------
        - `List[Data]`, PyG Data 对象列表
        """
        assert len(adjacencies) == len(features) == len(labels)
        return [
            cls.from_adjacency(adj, feat, label, **kwargs)
            for adj, feat, label in zip(adjacencies, features, labels)
        ]
