DEVS2Agents/DEVS/PolicedTrafficLight/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"
    TakeOver = "Take Over"

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

        self.takeOver = self.addInPort("takeOver")

    def extTransition(self, inputs):
        if self.takeOver in inputs:
            if inputs[self.takeOver][0]:
                return Phase.TakeOver
            return Phase.Red
        return self.state

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

    def timeAdvance(self):
        if self.state == Phase.Green:
            return self.phases[0]
        elif self.state == Phase.Orange:
            return self.phases[1]
        elif self.state == Phase.Red:
            return self.phases[2]
        return INFINITY # Police Take Over