import wntr
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import sklearn.cluster
import os



class QD_KMeans(object):
    def __init__(self, wn, num_monitors):
        # self.inp = inp
        self.cluster_num = num_monitors # 聚类中心个数，也即测压点个数
        self.wn=wn
        self.monitor_nodes = []
        self.coords = []
        self.junction_nodes = {}  # Added missing initialization


    def get_junctions_coordinates(self):
       
        for junction_name in self.wn.junction_name_list: 
            junction = self.wn.get_node(junction_name)
            self.junction_nodes[junction_name] = junction.coordinates
            self.coords.append(junction.coordinates )

        # print(f"Total junctions: {self.junction_coordinates}")   

    def select_monitoring_points(self):
        if not self.coords:  # Add check if coordinates are collected
            self.get_junctions_coordinates()
        coords = np.array(self.coords)
        coords_normalized = (coords - coords.min(axis=0)) / (coords.max(axis=0) - coords.min(axis=0))
        kmeans = sklearn.cluster.KMeans(n_clusters= self.cluster_num, random_state=42)
        kmeans.fit(coords_normalized) 

        for center in kmeans.cluster_centers_:
            distances = np.sum((coords_normalized - center) ** 2, axis=1)
            nearest_node = self.wn.junction_name_list[np.argmin(distances)]
            self.monitor_nodes.append(nearest_node)  

        return self.monitor_nodes
    

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


    
    
def kmeans_sensor_placement(name: str, sensor_num: int, min_diameter: int) -> list:
    inp_name = f'./db_inp/{name}.db.inp'
    wn= wntr.network.WaterNetworkModel(inp_name)
    wn_cluster=QD_KMeans(wn, sensor_num)

    # Select monitoring pointse
    sensor_ids= wn_cluster.select_monitoring_points()

    # wn_cluster.visualize_network()

    return sensor_ids


     
if __name__ == "__main__":
    #sensorindex = get_ID(name='suzhouhe_2024_cloud_0817', sensor_num=30, min_diameter=500)
    sensorindex = kmeans_sensor_placement(name='szh', sensor_num=50, min_diameter=300)
    print(sensorindex)