import math
import random

# Function to calculate Euclidean distance for cities with coordinates
def calculate_distance_between_coordinates(city1, city2):
    return math.sqrt((city2[0] - city1[0])**2 + (city2[1] - city1[1])**2)

# Function to calculate the distance from a distance matrix
def calculate_distance_between_matrix(i, j, dist_matrix):
    return dist_matrix[i][j]

def simulated_annealing_tsp(cities=None, dist_matrix=None, city_names=None, temp=1000, temp_decay=0.995, max_iter=10000, start_city=None, end_city=None, circular=True):
    """
    Solve TSP using Simulated Annealing while including city names in the result.

    Args:
    - cities (list): List of cities (x, y) coordinates (only if using coordinates).
    - dist_matrix (list): Distance matrix (only if distances between cities are provided directly).
    - city_names (list): List of city names corresponding to the cities.
    - temp (float): Initial temperature for simulated annealing.
    - temp_decay (float): Decay factor for temperature after each iteration.
    - max_iter (int): Maximum number of iterations.
    - start_city (int): Optional starting city index. If None, the best start will be chosen.
    - end_city (int): Optional end city index. If None, the best end will be chosen.
    - circular (bool): If True, the route will return to the starting city. If False, it will not.

    Returns:
    - tuple: The best distance and the best path (list of city names).
    """
    # Choose the appropriate distance calculation based on input format
    if cities is not None:
        def calculate_distance(i, j):
            return calculate_distance_between_coordinates(cities[i], cities[j])
    else:
        def calculate_distance(i, j):
            return calculate_distance_between_matrix(i, j, dist_matrix)

    def calculate_total_distance(path):
        total_distance = 0
        for i in range(len(path) - 1):
            total_distance += calculate_distance(path[i], path[i+1])
        if circular:
            total_distance += calculate_distance(path[-1], path[0])  # Return to the starting city
        return total_distance

    def neighbor(path):
        # Create a neighbor by swapping two cities in the path (except the start and end cities)
        new_path = path[:]
        unvisited = path[1:-1]  # Exclude start and end cities
        i, j = random.sample(range(len(unvisited)), 2)
        unvisited[i], unvisited[j] = unvisited[j], unvisited[i]
        # Rebuild the path including start and end cities
        new_path[1:-1] = unvisited
        return new_path

    # If start_city or end_city are not provided, default to the first city
    n = len(cities) if cities is not None else len(dist_matrix)
    if start_city is None:
        start_city = random.randint(0, n - 1)  # Random start city if not specified
    if end_city is None:
        end_city = start_city  # End city is the same as start if not specified

    # Initialize the path (fixed start and end cities)
    current_path = list(range(n))
    current_path = [start_city] + [i for i in range(n) if i != start_city] + [end_city]  # Ensure start and end cities are fixed
    current_distance = calculate_total_distance(current_path)
    best_path = current_path[:]
    best_distance = current_distance

    # Perform simulated annealing
    for _ in range(max_iter):
        temp *= temp_decay
        if temp <= 0:
            break

        new_path = neighbor(current_path)
        new_distance = calculate_total_distance(new_path)

        # Accept the new solution if it's better, or with probability if it's worse
        if new_distance < current_distance or random.random() < math.exp((current_distance - new_distance) / temp):
            current_path = new_path
            current_distance = new_distance

            if current_distance < best_distance:
                best_path = current_path[:]
                best_distance = current_distance

    # Convert the best path from indices to city names
    city_path = [city_names[city] for city in best_path]
    return best_distance, city_path

# Example usage for Simulated Annealing with City Names:

# City names
city_names = ["City A", "City B", "City C", "City D"]

# Example distance matrix (distances between cities)
dist_matrix = [
    [0, 2, 4, 1],
    [2, 0, 3, 5],
    [4, 3, 0, 6],
    [1, 5, 6, 0]
]

# Example: Simulated Annealing TSP with circular route
best_distance, best_path = simulated_annealing_tsp(dist_matrix=dist_matrix, city_names=city_names, temp=1000, temp_decay=0.995, max_iter=10000, circular=True)
print("Simulated Annealing Best Distance:", best_distance)
print("Simulated Annealing Best Path:", " -> ".join(best_path))

# Example: Simulated Annealing TSP with non-circular route
best_distance, best_path = simulated_annealing_tsp(dist_matrix=dist_matrix, city_names=city_names, temp=1000, temp_decay=0.995, max_iter=10000, circular=False)
print("Simulated Annealing Best Distance (Non-Circular):", best_distance)
print("Simulated Annealing Best Path (Non-Circular):", " -> ".join(best_path))