from pypdevs.DEVS import AtomicDEVS
from pypdevs.simulator import Simulator
import dataclasses

@dataclasses.dataclass
class AutoCounterState:
    count: int
    def __init__(self, count):
        self.count = count

class AutoCounter(AtomicDEVS):
    def __init__(self, start):
        AtomicDEVS.__init__(self, "System")
        self.state = AutoCounterState(count=start)
        self.outport = self.addOutPort("outport")

    def timeAdvance(self):
        return 1.0

    def outputFnc(self):
        print("outputFnc...", self, self.state)
        return {self.outport: self.state.count}

    def intTransition(self):
        return AutoCounterState(self.state.count+1)


# Termination condition:
class StopAt:
    def __init__(self, stop_at):
        self.stop_at = stop_at
    def __call__(self, time, model):
        return model.state.count >= self.stop_at

if __name__ == "__main__":
    ctr = AutoCounter(start=0)
    sim = Simulator(ctr)
    sim.setClassicDEVS()
    sim.setVerbose()
    sim.setTerminationCondition(StopAt(3))
    sim.simulate()

    ctr2 = AutoCounter(start=ctr.state.count + 10) # we restart at 13
    sim2 = Simulator(ctr2)
    sim2.setClassicDEVS()
    sim2.setVerbose()
    sim2.setTerminationCondition(StopAt(15))
    sim2.simulate()