import numpy as np
import time
import matplotlib.pyplot as plt
import random

def calculate_total_distance(coord, path):
    """Calculate the total distance for a given path"""
    total_distance = 0
    for i in range(len(path) - 1):
        current_city = path[i]
        next_city = path[i+1]
        dx = coord[current_city, 0] - coord[next_city, 0]
        dy = coord[current_city, 1] - coord[next_city, 1]
        total_distance += np.sqrt(dx*dx + dy*dy)
    return total_distance

def generate_neighbor_path(path):
    """Generate a neighbor path by swapping two random cities (excluding start and end)"""
    # Keep start and end cities fixed
    new_path = path.copy()
    
    # Select two random positions (excluding first and last)
    if len(new_path) > 3:
        pos1, pos2 = random.sample(range(1, len(new_path)-1), 2)
        new_path[pos1], new_path[pos2] = new_path[pos2], new_path[pos1]
    
    return new_path

def simulated_annealing(coord, initial_temp=1000, cooling_rate=0.995, min_temp=1e-8, max_iterations=10000):
    """
    Find the shortest path for a traveling salesman problem using simulated annealing.
    Starts at the first city and ends at the last city.
    
    Parameters:
    coord: Nx2 array where each row represents (x, y) coordinates of a city
    initial_temp: Initial temperature for annealing
    cooling_rate: Rate at which temperature decreases
    min_temp: Minimum temperature threshold
    max_iterations: Maximum number of iterations
    
    Returns:
    Array listing cities in order of visits
    """
    # Convert to numpy array if not already
    coord = np.array(coord)
    
    # Number of cities
    n = len(coord)
    
    # Initialize with a random path that starts at city 0 and ends at city n-1
    path = list(range(n))  # [0, 1, 2, ..., n-1]
    
    # Keep start and end fixed
    # We'll keep 0 at the beginning and n-1 at the end
    middle_cities = list(range(1, n-1))
    random.shuffle(middle_cities)
    path = [0] + middle_cities + [n-1]
    
    # Calculate initial distance
    current_distance = calculate_total_distance(coord, path)
    best_path = path.copy()
    best_distance = current_distance
    
    # Simulated annealing parameters
    temperature = initial_temp
    iteration = 0
    
    # Time the search
    start_time = time.time()
    
    while temperature > min_temp and iteration < max_iterations:
        # Generate neighbor path
        new_path = generate_neighbor_path(path)
        new_distance = calculate_total_distance(coord, new_path)
        
        # Calculate energy difference
        delta_distance = new_distance - current_distance
        
        # Accept or reject the new path
        if delta_distance < 0 or random.random() < np.exp(-delta_distance / temperature):
            path = new_path
            current_distance = new_distance
            
            # Update best path if improved
            if current_distance < best_distance:
                best_path = path.copy()
                best_distance = current_distance
        
        # Cool down the temperature
        temperature *= cooling_rate
        iteration += 1
    
    end_time = time.time()
    search_time = end_time - start_time
    
    print(f"Path searching time: {search_time:.6f} seconds")
    print(f"Best distance: {best_distance:.6f}")
    print(f"Iterations: {iteration}")
    
    return best_path

# Example with 10 cities
if __name__ == "__main__":
    # Generate random coordinates for 10 cities
    np.random.seed(42)  # For reproducible results
    cities = np.random.rand(10, 2) * 100  # 10 cities with coordinates between 0 and 100
    
    print("Cities coordinates:")
    for i, (x, y) in enumerate(cities):
        print(f"City {i}: ({x:.2f}, {y:.2f})")
    
    # Find the best path using simulated annealing
    best_path = simulated_annealing(cities)
    
    print(f"\nBest path: {best_path}")
    
    # Visualize the result
    plt.figure(figsize=(10, 8))
    
    # Plot cities
    x_coords = cities[:, 0]
    y_coords = cities[:, 1]
    
    # Plot the path
    path_x = [cities[i, 0] for i in best_path]
    path_y = [cities[i, 1] for i in best_path]
    
    # Plot path
    plt.plot(path_x, path_y, 'b-', linewidth=2, marker='o', markersize=8)
    
    # Mark start and end cities differently
    start_x, start_y = cities[best_path[0], 0], cities[best_path[0], 1]
    end_x, end_y = cities[best_path[-1], 0], cities[best_path[-1], 1]
    
    plt.scatter(start_x, start_y, c='green', s=200, marker='o', label='Start City', zorder=5)
    plt.scatter(end_x, end_y, c='red', s=200, marker='o', label='End City', zorder=5)
    
    # Mark all cities
    plt.scatter(x_coords, y_coords, c='blue', s=100, marker='o', label='Cities')
    
    # Add labels for cities
    for i in range(len(cities)):
        plt.annotate(str(i), (cities[i, 0], cities[i, 1]), xytext=(5, 5), textcoords='offset points')
    
    plt.title('Optimal Traveling Salesman Path (Simulated Annealing)')
    plt.xlabel('X Coordinate')
    plt.ylabel('Y Coordinate')
    plt.legend()
    plt.grid(True)
    plt.axis('equal')
    
    # Save the plot
    plt.savefig('salesman_annealing_path.png')
    plt.show()