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

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 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[start_nodes].to_numpy()
    ) / 2.0
    y_vals = (node_y[end_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 find_new_center_pipe(
    node_x, node_y, candidate_pipe, pipe_start_node, pipe_end_node, record_center
):
    new_candidate_pipe = list(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,
):
    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, neighbors in G1.adjacency():
        w_temp = []
        n_t = [node_dict[node_name]]
        for neighbor_name, edge_data in neighbors.items():
            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)
    (edgecuts, parts) = pymetis.part_graph(
        nparts=group_num, adjncy=final_adjacency_list, xadj=xadj, eweights=w_f
    )
    # (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_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 = list(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,
                )

                # 更新
                new_center.append(center_t)
                new_group.append(candidate_pipe)
                new_all_node.append(nodeset)
                all_grouped_pipe = all_grouped_pipe + candidate_pipe
        else:
            for c in sub_graphs:
                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 = list(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,
                    )
                    # 更新
                    new_center.append(center_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,
                    record_center,
                )
            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


def visualize_metis_partition(
    G, center_pipes, pipe_groups, node_x, node_y, pipe_start_node_all, pipe_end_node_all
):
    """
    可视化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)
    """
    plt.figure(figsize=(9, 10))

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

    # --- 绘制背景管网（灰色半透明） ---
    for edge in G.edges():
        start_node, end_node = edge
        plt.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 = plt.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]
            plt.plot(
                [node_x[start], node_x[end]],
                [node_y[start], node_y[end]],
                color="red",
                linewidth=4,
                linestyle="--",
                dash_capstyle="round",
                zorder=3,  # 确保中心管道在最顶层
            )

    # --- 添加图例和标注 ---
    # 分组图例
    group_labels = [f"Group {i + 1}" for i in range(len(pipe_groups))]
    plt.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
            plt.text(
                x,
                y,
                f"C{i + 1}",
                color="red",
                fontsize=10,
                ha="center",
                va="center",
                bbox=dict(facecolor="white", alpha=0.8, edgecolor="none"),
            )

    # --- 图形美化 ---
    plt.title("Water Network Partitioning Overview", fontsize=14, pad=20)
    plt.xlabel("X Coordinate", fontsize=10)
    plt.ylabel("Y Coordinate", fontsize=10)
    plt.grid(True, alpha=0.2, linestyle=":")
    plt.tight_layout()

    # 显示图形
    plt.show()


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


class NetworkPartitioner:
    @staticmethod
    def metis_grouping_pipe_weight(*args, **kwargs):
        return metis_grouping_pipe_weight(*args, **kwargs)

    @staticmethod
    def visualize_metis_partition(*args, **kwargs):
        return visualize_metis_partition(*args, **kwargs)

    @staticmethod
    def generate_adjlist_with_all_edges(*args, **kwargs):
        return generate_adjlist_with_all_edges(*args, **kwargs)

    @staticmethod
    def pick_center_pipe(*args, **kwargs):
        return pick_center_pipe(*args, **kwargs)

    @staticmethod
    def find_new_center_pipe(*args, **kwargs):
        return find_new_center_pipe(*args, **kwargs)

    @staticmethod
    def cal_area_node_linked_pipe(*args, **kwargs):
        return cal_area_node_linked_pipe(*args, **kwargs)

    @staticmethod
    def cal_group_num(*args, **kwargs):
        return cal_group_num(*args, **kwargs)