import matplotlib.pyplot as plt
import numpy as np
import sklearn.cluster
import wntr


class KMeansPlacement:
    def __init__(self, wn, num_monitors: int, min_diameter_mm: float):
        self.cluster_num = num_monitors
        self.wn = wn
        self.monitor_nodes: list[str] = []
        self.coords: list[tuple[float, float]] = []
        self.candidate_nodes: list[str] = []
        self.min_diameter_mm = min_diameter_mm

    def get_junctions_coordinates(self) -> None:
        eligible_nodes: set[str] = set()
        junction_names = set(self.wn.junction_name_list)
        for pipe_name in self.wn.pipe_name_list:
            pipe = self.wn.get_link(pipe_name)
            if float(pipe.diameter) * 1000 < self.min_diameter_mm:
                continue
            eligible_nodes.update(
                node_id
                for node_id in (pipe.start_node_name, pipe.end_node_name)
                if node_id in junction_names
            )

        for junction_name in self.wn.junction_name_list:
            if junction_name not in eligible_nodes:
                continue
            junction = self.wn.get_node(junction_name)
            self.candidate_nodes.append(junction_name)
            self.coords.append(junction.coordinates)

    def select_monitoring_points(self) -> list[str]:
        if not self.coords:
            self.get_junctions_coordinates()
        if self.cluster_num <= 0:
            raise ValueError("sensor_count must be greater than zero")
        if self.cluster_num > len(self.candidate_nodes):
            raise ValueError("符合最小管径条件的候选节点数量少于请求的监测点数量")
        coords = np.array(self.coords)
        coordinate_span = coords.max(axis=0) - coords.min(axis=0)
        coordinate_span[coordinate_span == 0] = 1.0
        coords_normalized = (coords - coords.min(axis=0)) / coordinate_span
        kmeans = sklearn.cluster.KMeans(n_clusters=self.cluster_num, random_state=42)
        kmeans.fit(coords_normalized)

        selected_indices: set[int] = set()
        for cluster_index, center in enumerate(kmeans.cluster_centers_):
            cluster_indices = np.flatnonzero(kmeans.labels_ == cluster_index)
            available_indices = [
                int(index)
                for index in cluster_indices
                if int(index) not in selected_indices
            ]
            if not available_indices:
                available_indices = [
                    index
                    for index in range(len(self.candidate_nodes))
                    if index not in selected_indices
                ]
            nearest_index = min(
                available_indices,
                key=lambda index: (
                    float(np.sum((coords_normalized[index] - center) ** 2)),
                    index,
                ),
            )
            selected_indices.add(nearest_index)
            nearest_node = self.candidate_nodes[nearest_index]
            self.monitor_nodes.append(nearest_node)

        return self.monitor_nodes

    def visualize_network(self) -> None:
        """Visualize network with monitoring points."""
        wntr.graphics.plot_network(
            self.wn,
            node_attribute=self.monitor_nodes,
            node_size=30,
            title="Optimal sensor",
        )
        plt.show()


def optimize_sensor_placement(
    network_model: wntr.network.WaterNetworkModel,
    sensor_count: int,
    min_diameter_mm: float,
) -> list[str]:
    """Select sensor nodes from an already loaded network model."""

    placement = KMeansPlacement(network_model, sensor_count, min_diameter_mm)
    return placement.select_monitoring_points()