from DEVS import *
from infinity import INFINITY

class Generator(AtomicDEVS):
        def __init__(self, name="Generator"):
                AtomicDEVS.__init__(self, name)
                self.state = {'name': "generating"}
                self.my_ports = {"g_out": self.addOutPort("g_out")}
        def timeAdvance(self):

                if self.state['name'] == 'generating':
                        return 1

        def outputFnc(self):

                if self.state['name'] == 'generating':
                        return {self.my_ports['g_out']: [{'type': 'Job', 'repr': {'duration': 0.3}}]}

        def intTransition(self):

                if self.state['name'] == 'generating':
                        new_state = {}
                        new_state['name'] = 'generating'
                        return new_state

        def extTransition(self, my_inputs):

                inputs = {k.getPortName(): v for k, v in my_inputs.iteritems()}
                return AtomicDEVS.extTransition(self, my_inputs)

        def confTransition(self, my_inputs):

                inputs = {k.getPortName(): v for k, v in my_inputs.iteritems()}
                return AtomicDEVS.confTransition(self, my_inputs)

class Processor(AtomicDEVS):
        def __init__(self, name="Processor"):
                AtomicDEVS.__init__(self, name)
                self.state = {'name': "idle"}
                self.my_ports = {"p_in": self.addInPort("p_in"), "p_out": self.addOutPort("p_out")}
        def timeAdvance(self):

                if self.state['name'] == 'processing':
                        return self.state['job']['duration']
                if self.state['name'] == 'idle':
                        return INFINITY

        def outputFnc(self):

                if self.state['name'] == 'processing':
                        return {self.my_ports['p_out']: [{'type': 'Job', 'repr': self.state['job']}]}
                if self.state['name'] == 'idle':
                        return {}

        def intTransition(self):

                if self.state['name'] == 'processing':
                        new_state = {}
                        new_state['name'] = 'idle'
                        return new_state

        def extTransition(self, my_inputs):

                inputs = {k.getPortName(): v for k, v in my_inputs.iteritems()}
                if self.state['name'] == 'idle':
                        new_state = {}
                        new_state['name'] = 'processing'
                        new_state['job'] = inputs['p_in'][0]['repr']
                        return new_state
                else:
                        return AtomicDEVS.extTransition(self, my_inputs)

        def confTransition(self, my_inputs):

                inputs = {k.getPortName(): v for k, v in my_inputs.iteritems()}
                return AtomicDEVS.confTransition(self, my_inputs)

class Collector(AtomicDEVS):
        def __init__(self, name="Collector"):
                AtomicDEVS.__init__(self, name)
                self.state = {'name': "waiting"}
                self.my_ports = {"c_in": self.addInPort("c_in")}
        def timeAdvance(self):

                if self.state['name'] == 'waiting':
                        return INFINITY

        def outputFnc(self):

                if self.state['name'] == 'waiting':
                        return {}

        def intTransition(self):

                return AtomicDEVS.intTransition(self, my_inputs)

        def extTransition(self, my_inputs):

                inputs = {k.getPortName(): v for k, v in my_inputs.iteritems()}
                if self.state['name'] == 'waiting':
                        new_state = {}
                        new_state['name'] = 'waiting'
                        new_state['nr_of_jobs'] = self.state['nr_of_jobs'] + 1
                        return new_state
                else:
                        return AtomicDEVS.extTransition(self, my_inputs)

        def confTransition(self, my_inputs):

                inputs = {k.getPortName(): v for k, v in my_inputs.iteritems()}
                return AtomicDEVS.confTransition(self, my_inputs)

class CoupledProcessor(CoupledDEVS):
        def __init__(self, name="CoupledProcessor"):
                CoupledDEVS.__init__(self, name)
                self.my_ports = {"cp_out": self.addOutPort("cp_out"), "cp_in": self.addInPort("cp_in")}
                self.submodels = {"p1": self.addSubModel(Processor(name="p1")), "p2": self.addSubModel(Processor(name="p2"))}
                self.connectPorts(self.my_ports["cp_in"], self.submodels["p1"].my_ports["p_in"])
                self.connectPorts(self.submodels["p1"].my_ports["p_out"], self.submodels["p2"].my_ports["p_in"])
                self.connectPorts(self.submodels["p2"].my_ports["p_out"], self.my_ports["cp_out"])

class Root(CoupledDEVS):
        def __init__(self, name="Root"):
                CoupledDEVS.__init__(self, name)
                self.my_ports = {}
                self.submodels = {"generator": self.addSubModel(Generator(name="generator")), "coupledprocessor": self.addSubModel(CoupledProcessor(name="coupledprocessor")), "processor": self.addSubModel(Processor(name="processor")), "collector": self.addSubModel(Collector(name="collector"))}
                self.connectPorts(self.submodels["coupledprocessor"].my_ports["cp_out"], self.submodels["processor"].my_ports["p_in"])
                self.connectPorts(self.submodels["generator"].my_ports["g_out"], self.submodels["coupledprocessor"].my_ports["cp_in"])
                self.connectPorts(self.submodels["processor"].my_ports["p_out"], self.submodels["collector"].my_ports["c_in"])