import numpy as np
import pandas as pd
from scipy.spatial import distance_matrix

def design_with_esau_williams(turbines_df, substation_coord, max_capacity_mw):
    """
    使用 Esau-Williams 启发式算法解决容量受限最小生成树 (CMST) 问题。
    
    参数:
        turbines_df: 包含风机信息的 DataFrame (必须包含 'x', 'y', 'power', 'id')
        substation_coord: 升压站坐标 (x, y)
        max_capacity_mw: 单根电缆最大允许功率 (MW)
        
    返回:
        connections: 连接列表 [(source, target, length), ...]
        turbines_with_cluster: 带有 'cluster' 列的 turbines DataFrame (用于兼容性)
    """
    
    # 数据准备
    n_turbines = len(turbines_df)
    coords = turbines_df[['x', 'y']].values
    powers = turbines_df['power'].values
    ids = turbines_df['id'].values
    
    # 升压站坐标
    if substation_coord.ndim > 1:
        sx, sy = substation_coord[0][0], substation_coord[0][1]
    else:
        sx, sy = substation_coord[0], substation_coord[1]
    
    # 1. 计算距离矩阵
    # 风机到风机
    dist_matrix = distance_matrix(coords, coords)
    # 风机到升压站
    dists_to_sub = np.sqrt((coords[:, 0] - sx)**2 + (coords[:, 1] - sy)**2)
    
    # 2. 初始化组件 (Components)
    # 初始状态下，每个风机是一个独立的组件，直接连接到升压站
    # 为了方便查找，我们维护一个 components 字典
    # key: component_root_id (代表该组件的唯一标识)
    # value: {
    #   'members': {node_idx, ...}, 
    #   'total_power': float, 
    #   'gate_node': int (连接到升压站的节点索引), 
    #   'gate_cost': float (gate_node 到升压站的距离)
    # }
    
    components = {}
    for i in range(n_turbines):
        components[i] = {
            'members': {i},
            'total_power': powers[i],
            'gate_node': i,
            'gate_cost': dists_to_sub[i]
        }
        
    # 记录已经建立的连接 (不包括通往升压站的默认连接)
    # 格式: (u, v, length)
    established_edges = []
    # 记录已建立连接的坐标，用于交叉检查: [((x1, y1), (x2, y2)), ...]
    established_lines = []

    def do_intersect(p1, p2, p3, p4):
        """
        检测线段 (p1, p2) 和 (p3, p4) 是否严格相交 (不包括端点接触)
        """
        # 检查是否共享端点
        if (p1[0]==p3[0] and p1[1]==p3[1]) or (p1[0]==p4[0] and p1[1]==p4[1]) or \
           (p2[0]==p3[0] and p2[1]==p3[1]) or (p2[0]==p4[0] and p2[1]==p4[1]):
            return False
            
        def ccw(A, B, C):
            # 向量叉积
            return (C[1]-A[1]) * (B[0]-A[0]) - (B[1]-A[1]) * (C[0]-A[0])

        # 如果跨立实验符号相反，则相交
        d1 = ccw(p1, p2, p3)
        d2 = ccw(p1, p2, p4)
        d3 = ccw(p3, p4, p1)
        d4 = ccw(p3, p4, p2)
        
        # 严格相交判断 (忽略共线重叠的情况，视为不交叉)
        if ((d1 > 1e-9 and d2 < -1e-9) or (d1 < -1e-9 and d2 > 1e-9)) and \
           ((d3 > 1e-9 and d4 < -1e-9) or (d3 < -1e-9 and d4 > 1e-9)):
            return True
            
        return False
    
    # 3. 迭代优化
    while True:
        # 预先收集当前所有组件的 Gate Edges (连接升压站的线段)
        # 格式: {cid: (gate_node_coord, substation_coord)}
        current_gate_lines = {}
        sub_coord_tuple = (sx, sy)
        for cid, data in components.items():
            gate_idx = data['gate_node']
            current_gate_lines[cid] = (coords[gate_idx], sub_coord_tuple)

        # 收集所有候选移动: (tradeoff, u, v, cid_u, cid_v)
        candidates = []
        
        # 建立 node_to_comp_id 映射以便快速查找
        node_to_comp_id = {}
        for cid, data in components.items():
            for member in data['members']:
                node_to_comp_id[member] = cid
                
        # 遍历所有边 (i, j)
        for i in range(n_turbines):
            cid_i = node_to_comp_id[i]
            gate_cost_i = components[cid_i]['gate_cost']
            
            for j in range(n_turbines):
                if i == j: continue
                
                cid_j = node_to_comp_id[j]
                
                # 必须是不同组件
                if cid_i == cid_j:
                    continue
                
                # 检查容量约束
                if components[cid_i]['total_power'] + components[cid_j]['total_power'] > max_capacity_mw:
                    continue
                
                # 计算 Tradeoff
                dist_ij = dist_matrix[i, j]
                tradeoff = dist_ij - gate_cost_i
                
                # 只有当 tradeoff < 0 时，合并才是有益的
                if tradeoff < -1e-9:
                    candidates.append((tradeoff, i, j, cid_i, cid_j))
        
        # 按 tradeoff 排序 (从小到大，越小越好)
        candidates.sort(key=lambda x: x[0])
        
        best_move = None
        
        # 延迟检测: 从最好的开始检查交叉
        for cand in candidates:
            tradeoff, u, v, cid_u, cid_v = cand
            
            p_u = coords[u]
            p_v = coords[v]
            
            # 快速包围盒测试 (AABB) 准备
            min_x_uv, max_x_uv = min(p_u[0], p_v[0]), max(p_u[0], p_v[0])
            min_y_uv, max_y_uv = min(p_u[1], p_v[1]), max(p_u[1], p_v[1])
            
            is_crossing = False
            
            # 1. 检查与已固定的内部边的交叉
            for line in established_lines:
                p_a, p_b = line[0], line[1]
                
                if max(p_a[0], p_b[0]) < min_x_uv or min(p_a[0], p_b[0]) > max_x_uv or \
                   max(p_a[1], p_b[1]) < min_y_uv or min(p_a[1], p_b[1]) > max_y_uv:
                    continue
                
                if do_intersect(p_u, p_v, p_a, p_b):
                    is_crossing = True
                    break
            
            if is_crossing:
                continue
                
            # 2. 检查与所有活跃 Gate Edges 的交叉 (排除被移除的那个)
            # 正在合并 cid_u -> cid_v，意味着 cid_u 的 Gate 将被移除。
            # 但 cid_v 的 Gate 以及其他所有组件的 Gate 仍然存在。
            for cid, gate_line in current_gate_lines.items():
                if cid == cid_u: 
                    continue # 这个 Gate 即将移除，不构成障碍
                
                p_a, p_b = gate_line[0], gate_line[1]
                
                if max(p_a[0], p_b[0]) < min_x_uv or min(p_a[0], p_b[0]) > max_x_uv or \
                   max(p_a[1], p_b[1]) < min_y_uv or min(p_a[1], p_b[1]) > max_y_uv:
                    continue
                    
                if do_intersect(p_u, p_v, p_a, p_b):
                    is_crossing = True
                    break
            
            if not is_crossing:
                best_move = (u, v, cid_u, cid_v)
                break
        
        # 如果没有找到有益的合并，或者所有可行合并都会增加成本，则停止
        if best_move is None:
            break
            
        # 执行合并
        u, v, cid_u, cid_v = best_move
        
        # 将 cid_u 并入 cid_v
        # 1. 记录新边
        established_edges.append((u, v, dist_matrix[u, v]))
        established_lines.append((coords[u], coords[v]))
        
        # 2. 更新组件信息
        comp_u = components[cid_u]
        comp_v = components[cid_v]
        
        # 合并成员
        comp_v['members'].update(comp_u['members'])
        comp_v['total_power'] += comp_u['total_power']
        
        # Gate 节点和 Cost 保持 cid_v 的不变 
        # (因为我们将 U 接到了 V 上，U 的原 gate 被移除，V 的 gate 仍是通往升压站的路径)
        
        # 3. 删除旧组件
        del components[cid_u]
        
    # 4. 构建最终结果
    connections = []
    
    # 添加内部边 (风机间)
    for u, v, length in established_edges:
        source = f'turbine_{u}'
        target = f'turbine_{v}'
        connections.append((source, target, length))
        
    # 添加 Gate 边 (风机到升压站)
    # 此时 components 中剩下的每个组件都有一个 gate_node 连接到 Substation
    cluster_mapping = {} # node_id -> cluster_id (0..N-1)
    
    for idx, (cid, data) in enumerate(components.items()):
        gate_node = data['gate_node']
        gate_cost = data['gate_cost']
        
        connections.append((f'turbine_{gate_node}', 'substation', gate_cost))
        
        # 记录 cluster id
        for member in data['members']:
            cluster_mapping[member] = idx
            
    # 更新 turbines DataFrame
    turbines_with_cluster = turbines_df.copy()
    turbines_with_cluster['cluster'] = turbines_with_cluster['id'].map(cluster_mapping)
    
    return connections, turbines_with_cluster
