#TSP greedy
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 greedy_tsp(cities=None, dist_matrix=None, city_names=None, start_city=None, end_city=None, circular=True):
    """
    Solve TSP using a greedy approach 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.
    - 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

    # Initialize start and end cities if not provided
    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 unvisited cities and the tour
    unvisited = list(range(n))
    path = [unvisited.pop(start_city)]  # Start with the start city
    total_distance = 0

    while unvisited:
        last_city = path[-1]
        next_city = min(unvisited, key=lambda city: calculate_distance(last_city, city))
        unvisited.remove(next_city)
        path.append(next_city)
        total_distance += calculate_distance(last_city, next_city)

    # Close the tour if circular is True
    if circular:
        total_distance += calculate_distance(path[-1], path[0])  # Return to the starting city
        path.append(path[0])

    # Convert path from indices to city names
    city_path = [city_names[city] for city in path]
    return total_distance, city_path

# Example usage for Greedy TSP with City Names:
city_names = ["City A", "City B", "City C", "City D"]
dist_matrix = [
    [0, 2, 4, 1],
    [2, 0, 3, 5],
    [4, 3, 0, 6],
    [1, 5, 6, 0]
]

# Example: Greedy TSP with circular route
best_distance, best_path = greedy_tsp(dist_matrix=dist_matrix, city_names=city_names, circular=True)
print("Greedy Best Distance:", best_distance)
print("Greedy Best Path:", " -> ".join(best_path))

# Example: Greedy TSP with non-circular route
best_distance, best_path = greedy_tsp(dist_matrix=dist_matrix, city_names=city_names, start_city=2, circular=False)
print("Greedy Best Distance (Non-Circular):", best_distance)
print("Greedy Best Path (Non-Circular):", " -> ".join(best_path))