import pandas as pd
from itertools import product

cities = ["New York", "Boston", "Chicago", "Washington", "Philadelphia"]  #Name of all the cities we are taking
days = ["Monday", "Tuesday", "Wednesday", "Thursday", "Friday", "Saturday", "Sunday"]# Days of the week on which the trains are running
train_types = ["Local", "Express"] # 2 Kinds of trains
departure_times = ["06:00", "09:00", "12:00"] # 3 Possible departure times
distances = {                             # Real world approximated distances between cities
    ("New York", "Boston"): "350",
    ("New York", "Chicago"): "1200",
    ("New York", "Washington"): "400",
    ("New York", "Philadelphia"): "150",
    ("Boston", "Chicago"): "1400",
    ("Boston", "Washington"): "600",
    ("Boston", "Philadelphia"): "500",
    ("Chicago", "Washington"): "1100",
    ("Chicago", "Philadelphia"): "1200",
    ("Washington", "Philadelphia"): "200",
}

for (city1, city2), dist in list(distances.items()): # Distances for the opposite routes will also be the same
    distances[(city2, city1)] = dist

def calculate_arrival_time(departure_time, distance_km):    # Function to calculate arrival time given departure time and distances between stations
    hours, minutes = map(int, departure_time.split(":"))
    departure_hours = hours + minutes / 60.0     #Converted into hours
    distance = float(distance_km)
    duration_hours = distance / 100.0      #Assumed the spped of train to be 100Kmph
    arrival_hours = departure_hours + duration_hours
    
    arrival_hours = arrival_hours % 24
    arrival_hour = int(arrival_hours)
    arrival_minute = int((arrival_hours - arrival_hour) * 60)
    return f"{arrival_hour:02d}:{arrival_minute:02d}"

routes = [(dep, arr) for dep in cities for arr in cities if dep != arr]  # Generate all combinations = 20 routes
combinations = list(product(routes, days, train_types, departure_times))

data = []              #Creating the dataset
for (dep, arr), day, train_type, dep_time in combinations:
    distance = distances[(dep, arr)]
    arr_time = calculate_arrival_time(dep_time, distance)
    data.append({
        "departure_station": dep,
        "arrival_station": arr,
        "departure_time": dep_time,
        "arrival_time": arr_time,
        "day_of_week": day,
        "train_type": train_type,
        "distance_km": distance
    })

df = pd.DataFrame(data) #Made a dataframe
df.to_csv("train_dataset_named.csv", index=False) #Converted into a CSV File