DEVS2Agents/DEVS/TrafficLight.py

from enum import Enum
from pypdevs.DEVS import AtomicDEVS
from pypdevs.infinity import INFINITY

class Phase(Enum):
    Green = "Green"
    Orange = "Orange"
    Red = "Red"

class TrafficLight(AtomicDEVS):
    def __init__(self, phases: tuple[int, int, int] = (25, 5, 30)):
        """
        :param phases: timings for green yellow/orange red phases of the traffic light
        """
        AtomicDEVS.__init__(self, "TrafficLight")
        self.phases = phases
        self.state = Phase.Red

    def intTransition(self):
        if self.state == Phase.Green:
            return Phase.Orange
        elif self.state == Phase.Orange:
            return Phase.Red
        return Phase.Green

    def timeAdvance(self):
        if self.state == Phase.Green:
            return self.phases[0]
        elif self.state == Phase.Orange:
            return self.phases[1]
        return self.phases[2]