tsp_train_solver.py

import csv
import random
from typing import List, Tuple
import numpy as np

class Station:
    def __init__(self, name: str, departures: List[float], arrivals: List[float]):
        self.name = name
        self.departures = departures
        self.arrivals = arrivals

def read_csv(filename: str) -> List[Station]:
    stations = []
    with open(filename, 'r') as file:
        reader = csv.reader(file)
        next(reader)  # Skip header
        for row in reader:
            name = row[0]
            departures = [float(x) for x in row[1::2]]
            arrivals = [float(x) for x in row[2::2]]
            stations.append(Station(name, departures, arrivals))
    return stations

def calculate_transfer_time(from_station: Station, to_station: Station) -> float:
    # Simplified transfer time calculation
    # In a real scenario, this would be more complex, considering schedules and probabilities
    return min(abs(d - a) for d in from_station.departures for a in to_station.arrivals)

def nearest_neighbor(stations: List[Station]) -> List[int]:
    n = len(stations)
    unvisited = set(range(1, n))
    tour = [0]

    while unvisited:
        last = tour[-1]
        next_station = min(unvisited, key=lambda x: calculate_transfer_time(stations[last], stations[x]))
        tour.append(next_station)
        unvisited.remove(next_station)

    return tour

def two_opt(tour: List[int], stations: List[Station]) -> List[int]:
    improvement = True
    best_tour = tour
    while improvement:
        improvement = False
        for i in range(1, len(tour) - 2):
            for j in range(i + 1, len(tour)):
                new_tour = tour[:i] + tour[i:j][::-1] + tour[j:]
                if calculate_total_time(new_tour, stations) < calculate_total_time(best_tour, stations):
                    best_tour = new_tour
                    improvement = True
        tour = best_tour
    return tour

def calculate_total_time(tour: List[int], stations: List[Station]) -> float:
    total_time = 0
    for i in range(len(tour) - 1):
        total_time += calculate_transfer_time(stations[tour[i]], stations[tour[i + 1]])
    return total_time

def solve_tsp(stations: List[Station]) -> List[int]:
    # Start with Nearest Neighbor heuristic
    initial_tour = nearest_neighbor(stations)

    # Improve the tour using 2-opt (a simplified version of Lin-Kernighan)
    improved_tour = two_opt(initial_tour, stations)

    return improved_tour

def main(filename: str):
    stations = read_csv(filename)
    optimal_tour = solve_tsp(stations)

    print("Optimal Itinerary:")
    for i in optimal_tour:
        print(f"Station: {stations[i].name}")

    total_time = calculate_total_time(optimal_tour, stations)
    print(f"\nExpected total time: {total_time:.2f} hours")

if __name__ == "__main__":
    main("foo.csv")