import numpy as np
import networkx as nx
import pandas as pd
from sklearn.neighbors import NearestNeighbors
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

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),
                              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),
                              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 ClusterNetworkBuilder2:
    """聚类中心网络构建器"""
    
    def __init__(self, assignment_method='gaussian', temperature=1.0):
        """
        参数:
        assignment_method: 分配方法 ('softmax', 'hard', 'gaussian', 'knn')
        temperature: softmax的温度参数
        """
        self.assignment_method = assignment_method
        self.temperature = temperature
    
    def build_network(self, x, c, k_neighbors=5, metric='euclidean', edge_threshold=None,
                     min_edge_weight=1e-6, normalize=True):
        """
        构建聚类中心网络
        
        返回:
        G: NetworkX图
        assignments: 数据点分配矩阵
        metrics: 网络指标字典
        """
        n_samples = x.shape[0]
        n_clusters = c.shape[0]
        
        pcs = min(5, x.shape[1])
        pc_model = PCA(n_components=pcs)
        pc_model.fit(x)
        x = pc_model.transform(x)
        c = pc_model.transform(c)
        
        # 1. 计算数据分配
        assignments = self._compute_assignments(x, c, k_neighbors, metric)
        
        # 2. 计算节点权重
        node_weights = assignments.sum(axis=0) / n_samples
        
        # 3. 构建网络
        G = nx.Graph()
        
        # 添加节点
        for i in range(n_clusters):
            G.add_node(i, 
                      weight=node_weights[i],
                      #center=c[i],
                      cluster_size=int(np.sum(assignments[:, i])))
        
        # 4. 计算边权重
        #edges_info = self._compute_edge_weights(c, assignments, k_neighbors, 
        #                                       metric, min_edge_weight)
        edges_info = self._compute_edge_weights(
                                x, c, assignments, 
                                edge_threshold, 
                                min_edge_weight
                            )
        
        # 添加边
        for (i, j), weight, distance, shared_samples in edges_info:
            G.add_edge(i, j, 
                      weight=weight,
                      distance=distance,
                      shared_samples=shared_samples)
        
        return G
    
    def _compute_bandwidths(self, x, k_neighbors, metric):
        n_samples = x.shape[0]
        # 方法2: 每行单独计算，基于局部密度
        bandwidths = np.zeros(n_samples)
            
        for i in range(n_samples):
            # 计算到所有其他点的距离
            if metric == 'euclidean':
                dist_to_all = np.linalg.norm(x - x[i], axis=1)
            else:
                # 对于其他度量，使用cdist
                dist_to_all = cdist(x[i:i+1], x, metric=metric).flatten()
                
            # 排序并取k近邻距离
            sorted_dist = np.sort(dist_to_all)
            k_idx = min(k_neighbors + 1, len(sorted_dist))
            bandwidths[i] = sorted_dist[k_idx] if k_idx < len(sorted_dist) else sorted_dist[-1]
            
        # 平滑处理
        bandwidths = np.maximum(bandwidths, np.percentile(bandwidths, 10))
        return bandwidths.reshape(-1, 1)
    
    def _compute_assignments(self, x, c, k_neighbors, metric):
        """计算数据点分配"""
        n_samples = x.shape[0]
        n_clusters = c.shape[0]
        
        # 计算距离
        distances = cdist(x, c, metric=metric)
        
        if self.assignment_method == 'hard':
            # 硬分配：每个点分配到最近的聚类
            assignments = np.zeros((n_samples, n_clusters))
            nearest = np.argmin(distances, axis=1)
            assignments[np.arange(n_samples), nearest] = 1
        
        elif self.assignment_method == 'softmax':
            # softmax分配
            # 使用负距离（距离越小，概率越大）
            similarities = np.exp(-distances / self.temperature)
            assignments = similarities / similarities.sum(axis=1, keepdims=True)
        
        elif self.assignment_method == 'gaussian':
            # 高斯核分配
            #sigma = np.median(distances, axis=1, keepdims=True)  # 使用中值作为带宽
            sigma = self._compute_bandwidths(x, k_neighbors, metric)
            similarities = np.exp(-distances**2 / (2 * sigma**2))
            assignments = similarities / similarities.sum(axis=1, keepdims=True)
        
        elif self.assignment_method == 'knn':
            # k近邻分配
            k = min(5, n_clusters)
            assignments = np.zeros((n_samples, n_clusters))
            knn_indices = np.argsort(distances, axis=1)[:, :k]
            
            for i in range(n_samples):
                # 平均分配给k个最近邻
                assignments[i, knn_indices[i]] = 1.0 / k
        
        else:
            raise ValueError(f"未知的分配方法: {self.assignment_method}")
        
        return assignments
    
    '''def _compute_edge_weights(self, c, assignments, k_neighbors, metric, min_edge_weight):
        """计算边权重"""
        n_clusters = c.shape[0]
        n_samples = assignments.shape[0]
        
        # 计算聚类中心之间的距离
        center_distances = cdist(c, c, metric=metric)
        
        # 找到每个聚类的k个最近邻
        edges_info = []
        
        for i in range(n_clusters):
            # 获取距离最近的k+1个聚类（包括自己）
            k = min(k_neighbors + 1, n_clusters)
            nearest = np.argsort(center_distances[i])[:k]
            
            for j in nearest:
                if i < j:  # 避免重复
                    distance = center_distances[i, j]
                    
                    # 计算边权重：共享数据点的比重
                    # 方法1: 使用Jaccard相似度
                    mask_i = assignments[:, i] > 0
                    mask_j = assignments[:, j] > 0
                    shared = np.logical_and(mask_i, mask_j)
                    
                    if shared.sum() > 0:
                        # 计算共享数据点的总分配权重
                        #weight_ij = (np.sqrt(assignments[shared, i] * assignments[shared, j])).sum() #/ 2
                        weight_ij = (-assignments[shared, i]*np.log(assignments[shared, i]) - assignments[shared, j]*np.log(assignments[shared, j])).sum() #/ 2
                        weight = weight_ij / n_samples
                        
                        if weight >= min_edge_weight:
                            edges_info.append(((i, j), weight, distance, shared.sum()))
                    
                    # 方法2: 使用最小分配权重（备选）
                    # min_assign = np.minimum(assignments[:, i], assignments[:, j])
                    # weight = min_assign.sum() / n_samples
        
        return edges_info'''
    
    def _compute_edge_weights(self, x, c, edge_threshold, min_edge_weight):
        """计算边权重（核心方法）"""
        n_samples = x.shape[0]
        n_clusters = c.shape[0]
        
        # 计算聚类中心之间的距离
        center_distances = cdist(c, c, self.distance_metric)
        
        edges_info = []
        
        # 确定要连接的边
        if edge_threshold is not None:
            # 使用距离阈值
            connected_pairs = np.where(center_distances < edge_threshold)
            pairs = list(zip(connected_pairs[0], connected_pairs[1]))
        else:
            # 使用k近邻
            pairs = []
            for i in range(n_clusters):
                distances_i = center_distances[i]
                # 排除自身
                distances_i[i] = np.inf
                # 找到k个最近邻
                k_actual = min(self.k_neighbors, n_clusters - 1)
                nearest = np.argsort(distances_i)[:k_actual]
                for j in nearest:
                    if i < j:  # 避免重复
                        pairs.append((i, j))
        
        # 为每对计算权重
        for i, j in pairs:
            if i >= j:  # 确保i < j
                continue
                
            distance = center_distances[i, j]
            
            # 基于距离的拟合权重
            weight, n_fitted, shared_weight = self._compute_distance_based_weight(
                x, c[i], c[j], n_samples
            )
            
            # 应用最小权重阈值
            if weight >= min_edge_weight:
                edges_info.append(((i, j), weight, distance, n_fitted, shared_weight))
        
        return edges_info
    
    def _compute_distance_based_weight(self, x, center_u, center_v, n_samples):
        """计算基于距离的边权重"""
        # 计算数据点到边的距离
        distances = self._distance_to_line_segment(x, center_u, center_v)
        
        # 自适应阈值：使用距离的中位数
        if len(distances) > 0:
            threshold = np.median(distances) * 0.5
        else:
            threshold = 1.0
        
        # 找到拟合的数据点
        fitted_mask = distances < threshold
        n_fitted = np.sum(fitted_mask)
        
        if n_fitted == 0:
            return 0.0, 0, 0.0
        
        # 计算权重：拟合点的平均接近度
        fitted_distances = distances[fitted_mask]
        proximity_scores = 1.0 / (1.0 + fitted_distances)  # 距离越小，接近度越高
        avg_proximity = np.mean(proximity_scores)
        
        # 归一化权重
        weight = (n_fitted / n_samples) * avg_proximity
        shared_weight = weight * n_samples  # 用于记录
        
        return weight, n_fitted, shared_weight
    
    def _distance_to_line_segment(self, points, A, B):
        """计算点到线段的欧氏距离"""
        A = np.array(A)
        B = np.array(B)
        points = np.array(points)
        
        # 向量计算
        AB = B - A
        AP = points - A
        
        # 计算投影
        t = np.dot(AP, AB) / (np.dot(AB, AB) + 1e-10)
        t = np.clip(t, 0, 1)  # 限制在线段内
        
        # 计算投影点
        projection = A + t[:, np.newaxis] * AB
        
        # 计算距离
        distances = np.linalg.norm(points - projection, axis=1)
        
        return distances
    
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()