# Import necessary libraries
import numpy as np
from sklearn.cluster import KMeans

# Class to represent a Particle in Particle Swarm Optimization
class Particle:
    def __init__(self, num_cities, num_trucks, distances, time_windows, infinite_distance_value=1e6):
        # Initialize the particle's position, velocity, and best position
        self.position = np.random.permutation(range(1, num_cities+1))
        self.velocity = np.zeros_like(self.position)
        self.best_position = self.position.copy()
        self.best_cost = float('inf')
        self.num_trucks = num_trucks
        self.distances = distances
        self.time_windows = time_windows
        self.infinite_distance_value = infinite_distance_value

        # Cluster the cities based on their distances
        self.city_cluster = self.cluster_cities()

    # Function to cluster cities based on their distances
    def cluster_cities(self):
        # Replace infinite distances with a large finite value
        clustering_distances = np.where(self.distances == float('inf'), self.infinite_distance_value, self.distances)
        kmeans = KMeans(n_clusters=self.num_trucks, n_init=10)
        clusters = kmeans.fit_predict(clustering_distances)
        return clusters

    # Function to evaluate the cost of the particle's current position
    def evaluate_cost(self):
        # Split cities among trucks
        trucks = [[] for _ in range(self.num_trucks)]
        for i in range(len(self.position)):
            current_city = int(self.position[i])
            truck = self.city_cluster[current_city - 1]
            trucks[truck].append(current_city)

        # Calculate the total cost based on the schedule of each truck
        total_cost = 0
        for truck in trucks:
            if len(truck) > 0:
                departure_time = self.time_windows[truck[0] - 1][0]
                for i in range(len(truck) - 1):
                    u, v = int(truck[i]), int(truck[i+1])
                    distance = self.distances[u-1, v-1] if self.distances[u-1, v-1] != float('inf') else self.infinite_distance_value
                    total_cost += distance
                    arrival_time = departure_time + distance
                    time_window = self.time_windows[v-1]
                    if arrival_time < time_window[0]:
                        total_cost += time_window[0] - arrival_time
                        departure_time = time_window[0]
                    elif arrival_time > time_window[1]:
                        total_cost += arrival_time - time_window[1]
                        departure_time = arrival_time
                    else:
                        departure_time = arrival_time
        return total_cost

    # Function to update the particle's velocity
    def update_velocity(self, global_best_position, inertia_weight, cognitive_weight, social_weight):
        r1 = np.random.random()
        r2 = np.random.random()
        cognitive_component = cognitive_weight * r1 * (self.best_position - self.position)
        social_component = social_weight * r2 * (global_best_position - self.position)
        self.velocity = inertia_weight * self.velocity + cognitive_component + social_component

    # Function to update the particle's position based on its velocity
    def update_position(self):
        indices = np.random.choice(range(len(self.position)), 2, replace=False)
        self.position[indices[0]], self.position[indices[1]] = self.position[indices[1]], self.position[indices[0]]
        current_cost = self.evaluate_cost()
        if current_cost < self.best_cost:
            self.best_cost = current_cost
            self.best_position = self.position.copy()