"""管网分区模块。"""

import math

import matplotlib.pyplot as plt
import networkx as nx
import networkx as networkx
import numpy as np
import pandas as pd
import pymetis
from scipy.sparse import coo_matrix, csr_matrix
from scipy.sparse.csgraph import connected_components


def _to_metis_edge_weight(edge_weight):
    weight = float(edge_weight)
    if not math.isfinite(weight):
        raise ValueError(f"Invalid non-finite METIS edge weight: {edge_weight}")
    # pymetis expects integer edge weights.
    return max(1, int(round(weight)))


def _dedupe_preserve_order(items):
    seen = set()
    output = []
    for item in items:
        if item in seen:
            continue
        seen.add(item)
        output.append(item)
    return output


def pick_center_pipe(node_x, node_y, candidate_pipe, pipe_start_node, pipe_end_node):
    candidate_pipe_list = list(candidate_pipe)
    start_nodes = pipe_start_node[candidate_pipe_list]
    end_nodes = pipe_end_node[candidate_pipe_list]

    x_vals = (node_x[start_nodes].to_numpy() + node_x[end_nodes].to_numpy()) / 2.0
    y_vals = (node_y[start_nodes].to_numpy() + node_y[end_nodes].to_numpy()) / 2.0
    mean_x = float(np.mean(x_vals))
    mean_y = float(np.mean(y_vals))
    distance = np.abs(x_vals - mean_x) + np.abs(y_vals - mean_y)
    center_idx = int(np.argmin(distance))
    return candidate_pipe_list[center_idx]


def pick_max_diameter_pipe(candidate_pipe, pipe_diameter):
    candidate_pipe_list = list(candidate_pipe)
    diameters = pd.to_numeric(
        pipe_diameter.reindex(candidate_pipe_list), errors="coerce"
    ).dropna()
    if len(diameters) != len(candidate_pipe_list):
        missing = sorted(set(candidate_pipe_list) - set(diameters.index))
        preview = ", ".join(map(str, missing[:10]))
        raise ValueError(f"Missing or invalid diameter for pipes: {preview}")

    max_diameter = float(diameters.max())
    max_diameter_pipes = sorted(
        [pipe for pipe, diameter in diameters.items() if float(diameter) == max_diameter],
        key=str,
    )
    return max_diameter_pipes[0]


def pick_dual_center_pipes(
    node_x, node_y, candidate_pipe, pipe_start_node, pipe_end_node, pipe_diameter
):
    geometric_center = pick_center_pipe(
        node_x, node_y, candidate_pipe, pipe_start_node, pipe_end_node
    )
    diameter_center = pick_max_diameter_pipe(candidate_pipe, pipe_diameter)
    return _dedupe_preserve_order([geometric_center, diameter_center])


def find_new_center_pipe(
    node_x,
    node_y,
    candidate_pipe,
    pipe_start_node,
    pipe_end_node,
    pipe_diameter,
    record_center,
):
    new_candidate_pipe = sorted(set(candidate_pipe) - set(record_center))
    if new_candidate_pipe == []:
        new_candidate_pipe = candidate_pipe
    center_t = pick_center_pipe(
        node_x,
        node_y,
        new_candidate_pipe,
        pipe_start_node,
        pipe_end_node,
    )
    return center_t


def cal_area_node_linked_pipe(nodeset, node_pipe_dic):
    pipeset = []
    for temp_node in nodeset:
        pipeset.extend(node_pipe_dic[temp_node])
    return pipeset


def metis_grouping_pipe_weight(
    G0,
    wn,
    all_node_iter,
    candidate_pipe_input,
    group_num,
    node_x,
    node_y,
    pipe_start_node_all,
    pipe_end_node_all,
    node_pipe_dic,
    all_node_series,
    couple_node_length,
    pipe_diameter,
):
    all_node_iter_series_new = all_node_series[all_node_iter]
    all_node_iter_series_new = all_node_iter_series_new.sort_values(ascending=True)
    all_node_iter_new = list(all_node_iter_series_new.index)

    G1 = G0.subgraph(all_node_iter_new)
    delimiter = " "
    adjacency_list = []
    node_dict = {}
    c_new = 0
    for each_node in all_node_iter_new:
        node_dict[each_node] = c_new
        c_new = c_new + 1
    correspond_dic = {}
    count_node = 0
    w = []
    for node_name in all_node_iter_new:
        neighbors = G1[node_name]
        w_temp = []
        n_t = [node_dict[node_name]]
        for neighbor_name in sorted(neighbors.keys()):
            edge_data = neighbors[neighbor_name]
            edge_key = f"{node_name},{neighbor_name}"
            reverse_edge_key = f"{neighbor_name},{node_name}"
            if edge_key in couple_node_length:
                edge_weight = couple_node_length[edge_key]
            elif reverse_edge_key in couple_node_length:
                edge_weight = couple_node_length[reverse_edge_key]
            elif edge_data.get("weight") is not None:
                edge_weight = float(edge_data["weight"])
            else:
                # Ignore graph edges that are outside candidate pipes and have no usable
                # partition weight (e.g. some non-pipe links in mixed network graphs).
                continue
            w_temp.append(_to_metis_edge_weight(edge_weight))
            n_t.append(node_dict[neighbor_name])
        w.append(w_temp)
        correspond_dic[n_t[0]] = count_node
        count_node = count_node + 1
        # del n_t[0]
        adjacency_list.append(n_t)
    adjacency_list_new = [[] * 1 for i in range(len(adjacency_list))]
    w_new = [[] * 1 for i in range(len(adjacency_list))]
    for i in range(len(adjacency_list)):
        adjacency_list_new[int(adjacency_list[i][0])] = adjacency_list[i]
        w_new[int(adjacency_list[i][0])] = w[i]

    for i in range(len(adjacency_list)):
        del adjacency_list_new[i][0]
    xadj = [0]
    w_f = []
    final_adjacency_list = []
    for i in range(len(adjacency_list_new)):
        final_adjacency_list = final_adjacency_list + adjacency_list_new[i]
        xadj.append(len(final_adjacency_list))
        w_f = w_f + w_new[i]

    # (edgecuts, parts) = pymetis.part_graph(nparts=group_num, adjacency=adjacency_list_new)
    metis_options = pymetis.Options()
    metis_options.seed = 42
    (edgecuts, parts) = pymetis.part_graph(
        nparts=group_num,
        adjncy=final_adjacency_list,
        xadj=xadj,
        eweights=w_f,
        options=metis_options,
    )
    # (edgecuts, parts) = pymetis.part_graph(nparts=group_num, adjacency=adjacency_list_new)
    candidate_group_list = [[] * 1 for i in range(group_num)]
    for i in range(len(all_node_iter_new)):
        candidate_group_list[parts[i]].append(all_node_iter_new[i])

    """parts_new = np.zeros(len(candidate_node_input), dtype=int)
    for i in range(len(candidate_group_list)):
        temp_group = candidate_group_list[i]
        for each_node in temp_group:
            parts_new[node_dict[each_node]] = i
    parts_new = list(parts_new)"""

    new_center = []
    new_group = []
    new_center_candidates = []
    new_all_node = []
    candidate_pipe_set = set(candidate_pipe_input)
    all_grouped_pipe = []
    for i in range(group_num):
        # 构建子图
        G_sub = G0.subgraph(candidate_group_list[i])

        # 计算联通子图
        sub_graphs = networkx.connected_components(G_sub)
        if networkx.number_connected_components(G_sub) == 1:
            # 求交集
            nodeset = G_sub.nodes()
            pipeset_set = set(cal_area_node_linked_pipe(nodeset, node_pipe_dic))
            candidate_pipe = sorted(pipeset_set.intersection(candidate_pipe_set))

            # 判断集合是否保留
            if len(candidate_pipe) > 0:
                # 保留 计算中心
                center_t = pick_center_pipe(
                    node_x,
                    node_y,
                    candidate_pipe,
                    pipe_start_node_all,
                    pipe_end_node_all,
                )
                center_candidates_t = pick_dual_center_pipes(
                    node_x,
                    node_y,
                    candidate_pipe,
                    pipe_start_node_all,
                    pipe_end_node_all,
                    pipe_diameter,
                )

                # 更新
                new_center.append(center_t)
                new_center_candidates.append(center_candidates_t)
                new_group.append(candidate_pipe)
                new_all_node.append(nodeset)
                all_grouped_pipe = all_grouped_pipe + candidate_pipe
        else:
            for c in sorted(sub_graphs, key=lambda c: min(c)):
                G_temp = G0.subgraph(c)
                nodeset = G_temp.nodes()
                pipeset = cal_area_node_linked_pipe(nodeset, node_pipe_dic)
                pipeset_set = set(pipeset)

                # 求交集
                candidate_pipe = sorted(pipeset_set.intersection(candidate_pipe_set))
                # print(len(candidate_node))
                # 判断集合是否保留
                if len(candidate_pipe) > 0:
                    # 保留 计算中心
                    center_t = pick_center_pipe(
                        node_x,
                        node_y,
                        candidate_pipe,
                        pipe_start_node_all,
                        pipe_end_node_all,
                    )
                    center_candidates_t = pick_dual_center_pipes(
                        node_x,
                        node_y,
                        candidate_pipe,
                        pipe_start_node_all,
                        pipe_end_node_all,
                        pipe_diameter,
                    )
                    # 更新
                    new_center.append(center_t)
                    new_center_candidates.append(center_candidates_t)
                    new_group.append(candidate_pipe)
                    new_all_node.append(nodeset)
                    all_grouped_pipe = all_grouped_pipe + candidate_pipe
    record_center = []
    c_g = 0
    for each_group in new_group:
        if len(each_group) < 3:
            record_center.append(new_center[c_g])
        c_g += 1
    c_g = 0
    for each_group in new_group:
        if len(each_group) >= 3:
            if new_center[c_g] in record_center:
                new_center[c_g] = find_new_center_pipe(
                    node_x,
                    node_y,
                    each_group,
                    pipe_start_node_all,
                    pipe_end_node_all,
                    pipe_diameter,
                    record_center,
                )
                new_center_candidates[c_g] = _dedupe_preserve_order(
                    [new_center[c_g]] + list(new_center_candidates[c_g])
                )
            record_center.append(new_center[c_g])
        c_g += 1

    # visualize_metis_partition(
    #     G0, new_center, new_group,
    #     node_x, node_y,
    #     pipe_start_node_all, pipe_end_node_all
    # )
    return new_center, new_group, new_all_node, new_center_candidates


def visualize_metis_partition(
    G,
    center_pipes,
    pipe_groups,
    node_x,
    node_y,
    pipe_start_node_all,
    pipe_end_node_all,
    title: str | None = None,
    block: bool = True,
    pause_seconds: float | None = None,
):
    """
    可视化METIS分区结果（单图模式）
    参数：
        G: 原始管网图(nx.Graph)
        center_pipes: 中心管道列表(list)
        pipe_groups: 分组管道列表(list of lists)
        node_x: 节点X坐标字典(dict)
        node_y: 节点Y坐标字典(dict)
        pipe_start_node_all: 管道起点字典(dict)
        pipe_end_node_all: 管道终点字典(dict)
    """
    fig = plt.figure("metis_partition_convergence", figsize=(22.51, 12.48))
    fig.clf()
    ax = fig.add_subplot(111)
    if not block:
        plt.ion()

    # 生成颜色映射（自动扩展颜色数量）
    colors = plt.cm.tab20(np.linspace(0, 1, len(pipe_groups)))

    # --- 绘制背景管网（灰色半透明） ---
    for edge in G.edges():
        start_node, end_node = edge
        ax.plot(
            [node_x[start_node], node_x[end_node]],
            [node_y[start_node], node_y[end_node]],
            color="lightgray",
            linewidth=0.5,
            alpha=0.3,
            zorder=1,  # 确保背景在底层
        )

    # --- 绘制各分区管道（彩色）---
    legend_handles = []  # 用于图例的句柄
    for i, (group, center) in enumerate(zip(pipe_groups, center_pipes)):
        color = colors[i % len(colors)]  # 循环使用颜色

        # 绘制分组管道
        for pipe in group:
            start = pipe_start_node_all[pipe]
            end = pipe_end_node_all[pipe]
            line = ax.plot(
                [node_x[start], node_x[end]],
                [node_y[start], node_y[end]],
                color=color,
                linewidth=2.5,
                alpha=0.8,
                zorder=2,
            )
            # 只为每个分组的第一个管道添加图例句柄
            if pipe == group[0]:
                legend_handles.append(line[0])

        # 高亮中心管道（红色虚线）
        if center in pipe_start_node_all and center in pipe_end_node_all:
            start = pipe_start_node_all[center]
            end = pipe_end_node_all[center]
            ax.plot(
                [node_x[start], node_x[end]],
                [node_y[start], node_y[end]],
                color="red",
                linewidth=4,
                linestyle="--",
                dash_capstyle="round",
                zorder=3,  # 确保中心管道在最顶层
            )

    # --- 添加图例和标注 ---
    # 分组图例
    if legend_handles:
        group_labels = [f"Group {i + 1}" for i in range(len(pipe_groups))]
        ax.legend(
            legend_handles,
            group_labels,
            loc="upper right",
            title="Partitions",
            fontsize=8,
            title_fontsize=10,
        )

    # 中心管道标注（可选）
    for i, center in enumerate(center_pipes):
        if center in pipe_start_node_all:
            x = (
                node_x[pipe_start_node_all[center]] + node_x[pipe_end_node_all[center]]
            ) / 2
            y = (
                node_y[pipe_start_node_all[center]] + node_y[pipe_end_node_all[center]]
            ) / 2
            ax.text(
                x,
                y,
                f"C{i + 1}",
                color="red",
                fontsize=10,
                ha="center",
                va="center",
                bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"),
            )

    # --- 图形美化 ---
    ax.set_title(title or "Water Network Partitioning Overview", fontsize=14, pad=20)
    ax.set_xlabel("X Coordinate", fontsize=10)
    ax.set_ylabel("Y Coordinate", fontsize=10)
    ax.grid(True, alpha=0.2, linestyle=":")
    fig.tight_layout()

    # 显示图形并强制刷新，避免迭代显示滞后一轮。
    plt.show(block=block)
    if not block:
        fig.canvas.draw_idle()
        fig.canvas.flush_events()
        pause_value = 0.001 if pause_seconds is None else max(0.0, float(pause_seconds))
        plt.pause(max(0.001, pause_value))
    elif pause_seconds is not None:
        plt.pause(max(0.0, float(pause_seconds)))
    return fig


def generate_adjlist_with_all_edges(G, delimiter):
    for s, nbrs in G.adjacency():
        line = str(s) + delimiter
        for t, data in nbrs.items():
            line += str(t) + delimiter
        yield line[: -len(delimiter)]


def cal_group_num(candidate_node_input, cal_group_num):
    candidate_node_num = len(candidate_node_input)
    if candidate_node_num > 100:
        group_num_input = cal_group_num  # 30
    else:
        group_num_input = 10

    return group_num_input