Untitled

import random

class Train:
    def __init__(self, name, station, direction):
        self.name = name
        self.station = station
        self.direction = direction

class Station:
    def __init__(self, name, delay_prob):
        self.name = name
        self.delay_prob = delay_prob
        self.trains = []
        
    def add_train(self, train):
        self.trains.append(train)
        
    def remove_train(self, train):
        self.trains.remove(train)
        
class Line:
    def __init__(self, name, color, stations, direction):
        self.name = name
        self.color = color
        self.stations = stations
        self.direction = direction
        
    def get_next_station(self, station):
        index = self.stations.index(station)
        if index == 0:
            return self.stations[1]
        elif index == len(self.stations) - 1:
            return self.stations[-2]
        else:
            if self.direction == "north":
                return self.stations[index - 1]
            else:
                return self.stations[index + 1]
        
class Subway:
    def __init__(self, lines):
        self.lines = lines
        self.stations = {}
        for line in lines:
            for station in line.stations:
                if station not in self.stations:
                    self.stations[station] = Station(station, random.uniform(0,1))
                    
    def add_train(self, train):
        station = self.stations[train.station]
        station.add_train(train)
        
    def remove_train(self, train):
        station = self.stations[train.station]
        station.remove_train(train)
        
    def step(self):
        for station in self.stations.values():
            for train in station.trains:
                line = self.get_line(train.name)
                next_station = line.get_next_station(train.station)
                if random.uniform(0,1) < station.delay_prob:
                    print(f"Train {train.name} at {train.station} is delayed")
                else:
                    self.remove_train(train)
                    train.station = next_station
                    self.add_train(train)
    
    def get_train_info(self, train_name):
        for station in self.stations.values():
            for train in station.trains:
                if train.name == train_name:
                    print(f"Train {train_name} is at {train.station} station, heading {train.direction} on line {self.get_line(train_name).name}")
                    return
        print(f"No train found with name {train_name}")
        
    def get_route_info(self, station_name):
        if station_name not in self.stations:
            print(f"No station found with name {station_name}")
            return
        station = self.stations[station_name]
        print(f"Trains at {station_name}:")
        for train in station.trains:
            print(f"- Train {train.name} heading {train.direction} on line {self.get_line(train.name).name}")
    
    def get_line(self, train_name):
        for line in self.lines:
            if line.name in train_name:
                return line
        
    def run(self):
        while True:
            choice = input("Enter 1 to continue simulation, 2 to get train info, 3 to get route info, or 4 to quit: ")
            if choice == "1":
                self.step()
            elif choice == "2":
                train_name = input("Enter train name: ")
                self.get_train_info(train_name)
            elif choice == "3":
                route_name = input("Enter route name: ")
                self.get_route_info
            elif choice == "4":
                print("Thank you for using our train program!")
            else:
                print("Invalid input. Please enter a number from 1-4.")
                
running_subway = Subway(self,lines)
if __name__ == "__main__":
    running_subway.run()