import numpy as np
import networkx as nx
import pandas as pd
from sklearn.decomposition import PCA
from scipy.spatial.distance import cdist
from scipy import sparse
import matplotlib.pyplot as plt
from collections import defaultdict, Counter

from sklearn.neighbors import NearestNeighbors, KDTree
from scipy.sparse import csr_matrix, lil_matrix, diags
from scipy.sparse.linalg import eigs
from scipy.sparse.csgraph import dijkstra
import heapq
from typing import Tuple, List, Dict, Optional, Union
import warnings


class ClusterNetworkBuilder:
    """
    基于KNN的聚类中心网络构建器
    
    算法步骤：
    1. 建立聚类中心c到数据点x的knn网络H
    2. 基于H建立聚类中心c的网络G
       - 如果两个聚类中心c1可以经过某个数据点到达c2，则c1和c2建立边
       - 边的权重由中介数据点的数量决定
       - 中介数据点的比重越大，边的权重越大
    """
    
    def __init__(self, 
                 k_data_to_centers=5,      # 每个数据点连接到多少个聚类中心
                 k_centers_to_data=10,     # 每个聚类中心连接到多少个数据点
                 weight_method='count',     # 权重计算方法: 'count', 'proportion', 'weighted'
                 min_edge_weight=0.01,     # 最小边权重
                 normalize_weights=True,   # 是否归一化权重
                 symmetric=True,           # 是否创建对称边
                 use_jaccard=False):       # 是否使用Jaccard相似度
        """
        参数:
        k_data_to_centers: 每个数据点连接到k个最近的聚类中心
        k_centers_to_data: 每个聚类中心连接到k个最近的数据点
        weight_method: 权重计算方法
            - 'count': 使用中介数据点的数量
            - 'proportion': 使用中介数据点的比例
            - 'weighted': 使用加权计数
        """
        self.k_data_to_centers = k_data_to_centers
        self.k_centers_to_data = k_centers_to_data
        self.weight_method = weight_method
        self.min_edge_weight = min_edge_weight
        self.normalize_weights = normalize_weights
        self.symmetric = symmetric
        self.use_jaccard = use_jaccard
        
    def build_network(self, x, c, metric='euclidean', return_details=False):
        """
        构建聚类中心网络
        
        参数:
        x: 数据点 (n_samples, n_features)
        c: 聚类中心 (n_clusters, n_features)
        metric: 距离度量
        return_details: 是否返回详细中间结果
        
        返回:
        G: 聚类中心网络
        如果return_details=True，还返回(H, edge_weights_dict, data_assignments)
        """
        n_samples = x.shape[0]
        n_clusters = c.shape[0]
        
        # 步骤1: 构建KNN网络H（二部图）
        H, data_assignments = self._build_knn_bipartite(x, c, metric)
        
        # 步骤2: 基于H构建聚类中心网络G
        G, edge_weights_dict = self._build_cluster_network_from_bipartite(
            H, n_clusters, n_samples
        )
        
        # 为节点添加属性
        self._add_node_attributes(G, c, data_assignments, n_samples)
        
        if return_details:
            return G, H, edge_weights_dict, data_assignments
        return G
    
    def _build_knn_bipartite(self, x, c, metric):
        """构建数据点和聚类中心的KNN二部图H"""
        n_samples = x.shape[0]
        n_clusters = c.shape[0]
        
        # 创建二部图
        H = nx.Graph()
        
        # 添加节点
        # 数据点节点: 前缀'd_'，如'd_0', 'd_1', ...
        for i in range(n_samples):
            H.add_node(f'd_{i}', type='data', index=i, features=x[i])
        
        # 聚类中心节点: 前缀'c_'，如'c_0', 'c_1', ...
        for j in range(n_clusters):
            H.add_node(f'c_{j}', type='cluster', index=j, center=c[j])
        
        # 1. 数据点到聚类中心的连接
        if self.k_data_to_centers > 0:
            # 计算每个数据点到所有聚类中心的距离
            data_to_center_dists = cdist(x, c, metric=metric)
            
            # 为每个数据点找到k个最近的聚类中心
            k1 = min(self.k_data_to_centers, n_clusters)
            for i in range(n_samples):
                # 找到最近的k1个聚类中心
                if k1 < n_clusters:
                    # 使用argpartition提高效率
                    nearest = np.argpartition(data_to_center_dists[i], k1)[:k1]
                else:
                    nearest = np.arange(n_clusters)
                
                for j in nearest:
                    distance = data_to_center_dists[i, j]
                    H.add_edge(f'd_{i}', f'c_{j}', 
                              #weight=1.0/(1.0 + distance),
                              weight=np.exp(-distance),
                              distance=distance,
                              type='data_to_center')
        
        # 2. 聚类中心到数据点的连接
        if self.k_centers_to_data > 0:
            # 计算每个聚类中心到所有数据点的距离
            center_to_data_dists = cdist(c, x, metric=metric)
            
            # 为每个聚类中心找到k个最近的数据点
            k2 = min(self.k_centers_to_data, n_samples)
            for j in range(n_clusters):
                if k2 < n_samples:
                    nearest = np.argpartition(center_to_data_dists[j], k2)[:k2]
                else:
                    nearest = np.arange(n_samples)
                
                for i in nearest:
                    distance = center_to_data_dists[j, i]
                    H.add_edge(f'c_{j}', f'd_{i}',
                              #weight=1.0/(1.0 + distance),
                              weight=np.exp(-distance),
                              distance=distance,
                              type='center_to_data')
        
        # 创建数据点分配矩阵
        data_assignments = self._create_data_assignments(H, n_samples, n_clusters)
        
        return H, data_assignments
    
    def _create_data_assignments(self, H, n_samples, n_clusters):
        """从二部图创建数据点分配矩阵"""
        assignments = np.zeros((n_samples, n_clusters))
        
        for i in range(n_samples):
            data_node = f'd_{i}'
            if data_node in H:
                # 获取连接到这个数据点的所有聚类中心
                cluster_neighbors = [n for n in H.neighbors(data_node) 
                                   if n.startswith('c_')]
                
                for cluster_node in cluster_neighbors:
                    j = int(cluster_node.split('_')[1])
                    # 使用边的权重作为分配权重
                    weight = H[data_node][cluster_node].get('weight', 1.0)
                    assignments[i, j] = weight
        
        # 归一化每行的分配（使每行和为1）
        row_sums = assignments.sum(axis=1, keepdims=True)
        row_sums[row_sums == 0] = 1  # 避免除零
        assignments = assignments / row_sums
        
        return assignments
    
    def _build_cluster_network_from_bipartite(self, H, n_clusters, n_samples):
        """从二部图构建聚类中心网络"""
        G = nx.Graph()
        
        # 添加聚类中心节点
        for j in range(n_clusters):
            G.add_node(j, index=j)
        
        # 计算边权重
        edge_weights = defaultdict(float)
        
        # 方法1: 通过数据点连接聚类中心
        for i in range(n_samples):
            data_node = f'd_{i}'
            if data_node in H:
                # 获取连接到这个数据点的所有聚类中心
                cluster_neighbors = [n for n in H.neighbors(data_node) 
                                   if n.startswith('c_')]
                cluster_indices = [int(n.split('_')[1]) for n in cluster_neighbors]
                
                # 为每对聚类中心添加权重
                for idx1, c1 in enumerate(cluster_indices):
                    for c2 in cluster_indices[idx1+1:]:
                        # 这个数据点同时连接到c1和c2
                        edge_key = tuple(sorted((c1, c2)))
                        
                        # 计算这个数据点对这条边的贡献
                        weight_c1 = H[data_node][f'c_{c1}'].get('weight', 1.0)
                        weight_c2 = H[data_node][f'c_{c2}'].get('weight', 1.0)
                        
                        if self.weight_method == 'count':
                            # 计数：每个中介数据点贡献1
                            contribution = 1.0
                        elif self.weight_method == 'weighted':
                            # 加权：使用连接权重的乘积
                            contribution = weight_c1 * weight_c2
                        else:  # 'proportion'
                            # 比例：平均权重
                            contribution = (weight_c1 + weight_c2) / 2.0
                        
                        edge_weights[edge_key] += contribution
        
        # 如果使用Jaccard相似度
        if self.use_jaccard:
            edge_weights = self._compute_jaccard_weights(H, n_clusters, edge_weights)
        
        # 归一化权重
        if self.normalize_weights and edge_weights:
            max_weight = max(edge_weights.values())
            if max_weight > 0:
                for key in edge_weights:
                    edge_weights[key] /= max_weight
        
        # 添加边到网络G
        for (c1, c2), weight in edge_weights.items():
            if weight >= self.min_edge_weight:
                G.add_edge(c1, c2, weight=weight, 
                          n_mediators=int(weight * 10))  # 估计的中介点数
        
        return G, dict(edge_weights)
    
    def _compute_jaccard_weights(self, H, n_clusters, edge_weights):
        """计算Jaccard相似度作为边权重"""
        # 计算每个聚类中心连接的数据点集合
        cluster_data_sets = {}
        for j in range(n_clusters):
            cluster_node = f'c_{j}'
            if cluster_node in H:
                data_neighbors = {n for n in H.neighbors(cluster_node) 
                                if n.startswith('d_')}
                cluster_data_sets[j] = data_neighbors
        
        # 重新计算边权重为Jaccard相似度
        jaccard_weights = {}
        clusters = list(cluster_data_sets.keys())
        
        for i in range(len(clusters)):
            for j in range(i+1, len(clusters)):
                c1, c2 = clusters[i], clusters[j]
                set1 = cluster_data_sets.get(c1, set())
                set2 = cluster_data_sets.get(c2, set())
                
                if set1 and set2:
                    intersection = len(set1.intersection(set2))
                    union = len(set1.union(set2))
                    jaccard = intersection / union if union > 0 else 0
                    
                    if jaccard > 0:
                        jaccard_weights[(c1, c2)] = jaccard
        
        return jaccard_weights
    
    def _add_node_attributes(self, G, c, data_assignments, n_samples):
        """为节点添加属性"""
        n_clusters = c.shape[0]
        
        for j in range(n_clusters):
            if j in G:
                # 计算分配给这个聚类的数据点
                cluster_assignments = data_assignments[:, j]
                
                # 硬分配计数（分配权重大于0.5）
                hard_count = np.sum(cluster_assignments > 0.5)
                
                # 软分配权重和
                soft_weight = np.sum(cluster_assignments)
                
                # 更新节点属性
                G.nodes[j].update({
                    'center': c[j],
                    'hard_assigned_count': hard_count,
                    'soft_assigned_weight': soft_weight,
                    'proportion': soft_weight / n_samples,
                    'avg_assignment': np.mean(cluster_assignments[cluster_assignments > 0])
                })
    



class DiffusionClusterNetworkBuilder:
    """
    基于扩散距离的聚类中心网络构建器
    
    算法：
    1. 建立聚类中心c到数据点x的knn网络H
    2. 建立数据x的knn网络M
    3. 基于H建立聚类中心c的网络G，对任何两个中心c1和c2，
       它们之间的距离定义为从其中一个中心出发经过网络M到达另一个中心的扩散距离d，
       c1和c2之间的边的权重定义为exp(-d/t)，t为temperature
    """
    
    def __init__(self, 
                 k_data: int = 5,
                 k_centers: int = 3,
                 temperature: float = 1.0,
                 alpha: float = 0.9,  # 扩散过程的随机游走概率
                 diffusion_time: int = 3,  # 扩散时间步数
                 metric: str = 'euclidean',
                 algorithm: str = 'auto',
                 directed: bool = False,
                 weight_strategy: str = 'heat',  # 'heat', 'commute', 'normalized'
                 use_sparse: bool = True,
                 normalize: bool = True,
                 symmetrize: bool = True):
        """
        参数:
        k_data: 数据点KNN的k值
        k_centers: 聚类中心到数据点的KNN的k值
        temperature: 距离到权重的转换温度
        alpha: 随机游走重启概率
        diffusion_time: 扩散时间步数
        metric: 距离度量
        algorithm: KNN算法
        directed: 是否构建有向图
        weight_strategy: 权重策略 ('heat', 'commute', 'normalized')
        use_sparse: 是否使用稀疏矩阵
        normalize: 是否归一化扩散核
        symmetrize: 是否对称化距离
        """
        self.k_data = k_data
        self.k_centers = k_centers
        self.temperature = temperature
        self.alpha = alpha
        self.diffusion_time = diffusion_time
        self.metric = metric
        self.algorithm = algorithm
        self.directed = directed
        self.weight_strategy = weight_strategy
        self.use_sparse = use_sparse
        self.normalize = normalize
        self.symmetrize = symmetrize
        
        # 缓存
        self._H = None
        self._M = None
        self._G = None
        self._diffusion_matrix = None
        self._diffusion_distances = None
        
    def build_network(self, 
                     x: np.ndarray, 
                     c: np.ndarray,
                     return_intermediate: bool = False) -> Union[nx.Graph, Tuple]:
        """
        构建聚类中心网络
        
        参数:
        x: 数据点，形状 (n_samples, n_features)
        c: 聚类中心，形状 (n_clusters, n_features)
        return_intermediate: 是否返回中间图
        
        返回:
        G: 聚类中心网络
        如果return_intermediate为True，也返回(H, M, diffusion_matrix)
        """        
        # 1. 构建数据点的KNN图 M
        M, M_adj = self._build_data_knn_graph(x)
        
        # 2. 构建聚类中心到数据的KNN图 H
        H = self._build_centers_to_data_knn_graph(x, c)
        
        # 3. 计算扩散核和扩散距离
        diffusion_matrix = self._compute_diffusion_matrix(M_adj)
        
        # 4. 构建最终的聚类中心网络 G
        G = self._build_final_network(x, c, H, M, diffusion_matrix)
        
        # 保存中间结果
        self._M = M
        self._H = H
        self._G = G
        self._diffusion_matrix = diffusion_matrix
        
        if return_intermediate:
            return G, H, M, diffusion_matrix
        return G
    
    def _build_data_knn_graph(self, x: np.ndarray) -> Tuple[nx.Graph, csr_matrix]:
        """构建数据点的KNN图 M 和邻接矩阵"""
        n_samples = x.shape[0]
        
        # 构建KNN
        k_actual = min(self.k_data, n_samples - 1)
        if k_actual < 1:
            # 如果样本太少，构建完全图
            M = nx.complete_graph(n_samples)
            distances = cdist(x, x, self.metric)
            for i in range(n_samples):
                for j in range(i+1, n_samples):
                    M[i][j]['weight'] = distances[i, j]
            # 创建邻接矩阵
            adj_matrix = distances.copy()
            np.fill_diagonal(adj_matrix, 0)
            return M, csr_matrix(adj_matrix)
        
        # 使用sklearn的KNN
        nbrs = NearestNeighbors(
            n_neighbors=k_actual + 1,  # +1 因为包含自己
            algorithm=self.algorithm,
            metric=self.metric
        ).fit(x)
        
        distances, indices = nbrs.kneighbors(x)
        
        # 构建图和邻接矩阵
        if self.directed:
            M = nx.DiGraph()
        else:
            M = nx.Graph()
        
        M.add_nodes_from(range(n_samples))
        
        # 创建稀疏邻接矩阵
        if self.use_sparse:
            adj_matrix = lil_matrix((n_samples, n_samples))
        else:
            adj_matrix = np.zeros((n_samples, n_samples))
        
        for i in range(n_samples):
            neighbors = indices[i, 1:]  # 第0个是自己
            neighbor_distances = distances[i, 1:]
            
            for j, dist in zip(neighbors, neighbor_distances):
                if not self.directed and i > j:
                    if M.has_edge(i, j):
                        continue
                
                # 计算权重
                if self.weight_strategy == 'heat':
                    weight = np.exp(-dist**2 / (2 * self.temperature**2))
                elif self.weight_strategy == 'normalized':
                    weight = 1.0 / (1.0 + dist)
                else:  # 'distance'
                    weight = dist
                
                M.add_edge(i, j, weight=weight, distance=dist)
                
                # 更新邻接矩阵
                if self.use_sparse:
                    adj_matrix[i, j] = weight
                    if not self.directed:
                        adj_matrix[j, i] = weight
                else:
                    adj_matrix[i, j] = weight
                    if not self.directed:
                        adj_matrix[j, i] = weight
        
        if self.use_sparse:
            adj_matrix = adj_matrix.tocsr()
        
        return M, adj_matrix
    
    def _build_centers_to_data_knn_graph(self, x: np.ndarray, c: np.ndarray) -> nx.Graph:
        """构建聚类中心到数据的KNN图 H"""
        n_samples = x.shape[0]
        n_clusters = c.shape[0]
        
        k_actual = min(self.k_centers, n_samples)
        
        nbrs = NearestNeighbors(
            n_neighbors=k_actual,
            algorithm=self.algorithm,
            metric=self.metric
        ).fit(x)
        
        distances, indices = nbrs.kneighbors(c)
        
        # 构建二分图 H
        H = nx.Graph()
        
        # 添加节点
        H.add_nodes_from(range(n_clusters + n_samples))
        
        # 标记节点类型
        for i in range(n_clusters):
            H.nodes[i]['type'] = 'center'
            H.nodes[i]['center_id'] = i
        
        for i in range(n_samples):
            node_id = n_clusters + i
            H.nodes[node_id]['type'] = 'data'
            H.nodes[node_id]['data_id'] = i
        
        # 添加边：中心到数据
        for center_idx in range(n_clusters):
            data_indices = indices[center_idx]
            data_distances = distances[center_idx]
            
            for data_idx, dist in zip(data_indices, data_distances):
                data_node_id = n_clusters + data_idx
                
                # 计算边权重
                if self.weight_strategy == 'heat':
                    weight = np.exp(-dist**2 / (2 * self.temperature**2))
                else:
                    weight = 1.0 / (1.0 + dist)
                
                H.add_edge(center_idx, data_node_id, 
                          weight=weight, 
                          distance=dist,
                          center_id=center_idx,
                          data_id=data_idx)
        
        return H
    
    def _compute_diffusion_matrix(self, adj_matrix: Union[np.ndarray, csr_matrix]) -> csr_matrix:
        """计算扩散核矩阵"""
        n = adj_matrix.shape[0]
        
        print(f"   计算扩散核 (n={n}, time={self.diffusion_time})...")
        
        if self.use_sparse and isinstance(adj_matrix, csr_matrix):
            return self._compute_sparse_diffusion_matrix(adj_matrix)
        else:
            return self._compute_dense_diffusion_matrix(adj_matrix)
    
    def _compute_sparse_diffusion_matrix(self, adj_matrix: csr_matrix) -> csr_matrix:
        """计算稀疏扩散矩阵"""
        n = adj_matrix.shape[0]
        
        # 1. 计算度矩阵
        degrees = adj_matrix.sum(axis=1).A1
        D_inv_sqrt = diags(1.0 / np.sqrt(degrees + 1e-10))
        
        # 2. 计算归一化拉普拉斯
        I = diags(np.ones(n))
        
        if self.weight_strategy == 'normalized':
            # 对称归一化拉普拉斯: L = I - D^{-1/2} A D^{-1/2}
            L = I - D_inv_sqrt @ adj_matrix @ D_inv_sqrt
        else:
            # 随机游走拉普拉斯
            D_inv = diags(1.0 / (degrees + 1e-10))
            L = I - D_inv @ adj_matrix
        
        # 3. 计算扩散核: exp(-tL) 或 (I - αL)^t
        if self.alpha < 1.0:
            # 使用带重启的随机游走
            P = self.alpha * (I - L) + (1 - self.alpha) * I / n
            diffusion_matrix = P
            for _ in range(1, self.diffusion_time):
                diffusion_matrix = diffusion_matrix @ P
        else:
            # 使用热核
            try:
                # 对于大矩阵，使用近似
                if n > 1000:
                    # 使用低秩近似
                    k_eigen = min(50, n // 10)
                    eigenvalues, eigenvectors = eigs(L, k=k_eigen, which='SR')
                    
                    # 重构扩散核
                    exp_eigenvalues = np.exp(-self.diffusion_time * np.real(eigenvalues))
                    diffusion_matrix = eigenvectors @ diags(exp_eigenvalues) @ eigenvectors.T
                else:
                    # 精确计算
                    L_dense = L.toarray()
                    eigenvalues, eigenvectors = np.linalg.eigh(L_dense)
                    exp_eigenvalues = np.exp(-self.diffusion_time * eigenvalues)
                    diffusion_matrix = eigenvectors @ np.diag(exp_eigenvalues) @ eigenvectors.T
                    diffusion_matrix = csr_matrix(diffusion_matrix)
                    
            except:
                # 回退到简单方法
                warnings.warn("特征分解失败，使用简单扩散")
                P = 0.5 * (I - L) + 0.5 * I / n
                diffusion_matrix = P
                for _ in range(1, self.diffusion_time):
                    diffusion_matrix = diffusion_matrix @ P
        
        # 确保对称
        if self.symmetrize and not self.directed:
            diffusion_matrix = 0.5 * (diffusion_matrix + diffusion_matrix.T)
        
        return diffusion_matrix
    
    def _compute_dense_diffusion_matrix(self, adj_matrix: np.ndarray) -> np.ndarray:
        """计算密集扩散矩阵"""
        n = adj_matrix.shape[0]
        
        # 1. 计算度矩阵
        degrees = adj_matrix.sum(axis=1)
        D_inv_sqrt = np.diag(1.0 / np.sqrt(degrees + 1e-10))
        
        # 2. 计算归一化拉普拉斯
        I = np.eye(n)
        
        if self.weight_strategy == 'normalized':
            L = I - D_inv_sqrt @ adj_matrix @ D_inv_sqrt
        else:
            D_inv = np.diag(1.0 / (degrees + 1e-10))
            L = I - D_inv @ adj_matrix
        
        # 3. 计算扩散核
        if self.alpha < 1.0:
            P = self.alpha * (I - L) + (1 - self.alpha) * I / n
            diffusion_matrix = np.linalg.matrix_power(P, self.diffusion_time)
        else:
            # 热核扩散
            try:
                eigenvalues, eigenvectors = np.linalg.eigh(L)
                exp_eigenvalues = np.exp(-self.diffusion_time * eigenvalues)
                diffusion_matrix = eigenvectors @ np.diag(exp_eigenvalues) @ eigenvectors.T
            except:
                P = 0.5 * (I - L) + 0.5 * I / n
                diffusion_matrix = np.linalg.matrix_power(P, self.diffusion_time)
        
        if self.symmetrize and not self.directed:
            diffusion_matrix = 0.5 * (diffusion_matrix + diffusion_matrix.T)
        
        return diffusion_matrix
    
    def _compute_diffusion_distance(self, 
                                  diffusion_matrix: Union[np.ndarray, csr_matrix],
                                  i: int, 
                                  j: int) -> float:
        """计算两个点之间的扩散距离"""
        if isinstance(diffusion_matrix, csr_matrix):
            # 稀疏矩阵
            row_i = diffusion_matrix[i].toarray().flatten()
            row_j = diffusion_matrix[j].toarray().flatten()
        else:
            # 密集矩阵
            row_i = diffusion_matrix[i]
            row_j = diffusion_matrix[j]
        
        # 扩散距离定义为行向量的L2距离
        distance = np.sqrt(np.sum((row_i - row_j) ** 2))
        
        return distance
    
    def _build_final_network(self, 
                           x: np.ndarray, 
                           c: np.ndarray,
                           H: nx.Graph,
                           M: nx.Graph,
                           diffusion_matrix: Union[np.ndarray, csr_matrix]) -> nx.Graph:
        """构建最终的聚类中心网络 G"""
        n_clusters = c.shape[0]
        n_samples = x.shape[0]
        
        # 创建最终的网络
        G = nx.Graph()
        
        # 添加聚类中心节点
        for i in range(n_clusters):
            G.add_node(i, center=c[i], index=i)
        
        # 获取H中每个中心连接的数据点
        center_to_data = {}
        center_data_weights = {}
        
        for center_idx in range(n_clusters):
            connected_data = []
            data_weights = []
            
            for neighbor in H.neighbors(center_idx):
                if H.nodes[neighbor]['type'] == 'data':
                    data_id = H.nodes[neighbor]['data_id']
                    weight = H[center_idx][neighbor]['weight']
                    connected_data.append(data_id)
                    data_weights.append(weight)
            
            center_to_data[center_idx] = connected_data
            center_data_weights[center_idx] = np.array(data_weights)
        
        # 计算中心之间的扩散距离
        
        # 预计算数据点之间的扩散距离
        n_data_pairs = min(1000, n_samples * (n_samples - 1) // 2)
        data_diffusion_distances = {}
        
        # 采样计算，避免O(n^2)复杂度
        if n_samples > 1000:
            sample_size = min(1000, n_samples)
            sample_indices = np.random.choice(n_samples, sample_size, replace=False)
            
            for idx_i, i in enumerate(sample_indices[:-1]):
                for j in sample_indices[idx_i+1:]:
                    dist = self._compute_diffusion_distance(diffusion_matrix, i, j)
                    data_diffusion_distances[(i, j)] = dist
                    data_diffusion_distances[(j, i)] = dist
        else:
            # 计算所有对
            for i in range(n_samples):
                for j in range(i+1, n_samples):
                    dist = self._compute_diffusion_distance(diffusion_matrix, i, j)
                    data_diffusion_distances[(i, j)] = dist
                    data_diffusion_distances[(j, i)] = dist
        
        # 计算中心之间的扩散距离
        center_distances = np.full((n_clusters, n_clusters), np.inf)
        np.fill_diagonal(center_distances, 0)
        
        edges_added = 0
        
        for i in range(n_clusters):
            data_i = center_to_data[i]
            weights_i = center_data_weights[i]
            
            if len(data_i) == 0:
                continue
                
            for j in range(i+1, n_clusters):
                data_j = center_to_data[j]
                weights_j = center_data_weights[j]
                
                if len(data_j) == 0:
                    continue
                
                # 计算加权扩散距离
                min_distance = np.inf
                
                for idx_i, d_i in enumerate(data_i):
                    for idx_j, d_j in enumerate(data_j):
                        if d_i == d_j:
                            dist = 0
                        else:
                            dist = data_diffusion_distances.get((d_i, d_j), np.inf)
                        
                        if dist < np.inf:
                            # 加权距离
                            total_dist = (weights_i[idx_i] + weights_j[idx_j]) * dist
                            if total_dist < min_distance:
                                min_distance = total_dist
                
                if min_distance < np.inf:
                    center_distances[i, j] = min_distance
                    center_distances[j, i] = min_distance
                    
                    # 计算边权重
                    weight = np.exp(-min_distance / self.temperature)
                    
                    # 添加边
                    G.add_edge(i, j, 
                              weight=weight,
                              diffusion_distance=min_distance,
                              temperature=self.temperature,
                              n_data_i=len(data_i),
                              n_data_j=len(data_j))
                    
                    edges_added += 1
                
        return G
    

    


def visualize_network(G, c=None, node_size_scale=5000, edge_size_scale=30, figsize=(12, 10)):
    """可视化聚类中心网络"""
    fig, ax1 = plt.subplots(1, 1, figsize=figsize)
    
    # 位置布局（如果提供了聚类中心坐标，使用它们；否则使用力导向布局）
    if c is not None and c.shape[1] >= 2:
        # 使用前两个维度作为位置
        pos = {i: (c[i, 0], c[i, 1]) for i in G.nodes()}
        
        # 如果只有1D，添加一个虚拟维度
        if c.shape[1] == 1:
            pos = {i: (c[i, 0], i) for i in G.nodes()}
    else:
        pos = nx.spring_layout(G, seed=42, weight='weight')
    
    # 节点大小基于权重
    node_weights = [G.nodes[node]['proportion'] for node in G.nodes()]
    node_sizes = [node_size_scale * w for w in node_weights]
    
    # 节点颜色基于权重
    node_colors = node_weights
    
    # 边宽度基于权重
    edge_weights = [G[u][v]['weight'] for u, v in G.edges()]
    edge_widths = [edge_size_scale * w for w in edge_weights]
    
    # 绘制网络
    nodes = nx.draw_networkx_nodes(G, pos, 
                                  node_size=node_sizes,
                                  node_color=node_colors,
                                  cmap='viridis',
                                  alpha=0.8,
                                  edgecolors='black',
                                  linewidths=1,
                                  ax=ax1)
    
    nx.draw_networkx_edges(G, pos,
                          width=edge_widths,
                          alpha=0.6,
                          edge_color='gray',
                          ax=ax1)
    
    nx.draw_networkx_labels(G, pos,
                           font_size=10,
                           font_weight='bold',
                           ax=ax1)
    
    # 添加颜色条
    sm = plt.cm.ScalarMappable(cmap='viridis', 
                               norm=plt.Normalize(vmin=min(node_colors), 
                                                  vmax=max(node_colors)))
    sm.set_array([])
    cbar = plt.colorbar(sm, ax=ax1, shrink=0.8)
    cbar.set_label('Node Weight')
    
    ax1.axis('off')
    
    plt.tight_layout()
    plt.show()