import sys
sys.path.append("../../pypdevs/src/")

from DEVS import *
from infinity import INFINITY

class GeneratorState:
	def __init__(self, name=""):
		self.name = name

class ProcessorState:
	def __init__(self, name="", job=None):
		self.name = name
		self.job = job

class CollectorState:
	def __init__(self, name="", nr_of_jobs=None):
		self.name = name
		self.nr_of_jobs = nr_of_jobs

class Job:
	def __init__(self, jobSize=None):
		self.jobSize = jobSize

	def __str__(self):
		return "Job(" + str(self.jobSize) + ")"

class Generator(AtomicDEVS):
	def __init__(self, name="Generator"):
		AtomicDEVS.__init__(self, name)

		self.state = GeneratorState(name="generating")
		self.my_ports = {"p_out": self.addOutPort("p_out")}

	def timeAdvance(self):
		if self.state.name == "generating":
			return 1

	def outputFnc(self):
		if self.state.name == "generating":
			return {self.my_ports[k]: v for k, v in {'p_out': [Job(0.3)]}.iteritems()}

	def intTransition(self):
		def cond_int_generating_to_generating():
			return True
		def action_int_generating_to_generating():
			return {}
		if self.state.name == "generating" and cond_int_generating_to_generating():
			return GeneratorState(name="generating", **action_int_generating_to_generating())

		else:
			return AtomicDEVS.intTransition(self)

	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 = ProcessorState(name="idle")
		self.my_ports = {"p_in": self.addInPort("p_in"), "p_out": self.addOutPort("p_out")}

	def timeAdvance(self):
		if self.state.name == "idle":
			return INFINITY
		if self.state.name == "processing":
			return self.state.job.jobSize

	def outputFnc(self):
		if self.state.name == "idle":
			return {}
		if self.state.name == "processing":
			return {self.my_ports[k]: v for k, v in {'p_out': [self.state.job]}.iteritems()}

	def intTransition(self):
		def cond_int_processing_to_idle():
			return True
		def action_int_processing_to_idle():
			return {}
		if self.state.name == "processing" and cond_int_processing_to_idle():
			return ProcessorState(name="idle", **action_int_processing_to_idle())

		else:
			return AtomicDEVS.intTransition(self)

	def extTransition(self, my_inputs):
		inputs = {k.getPortName(): v for k, v in my_inputs.iteritems()}
		def cond_ext_idle_to_processing():
			return True
		def action_ext_idle_to_processing():
			return {"job": inputs['p_in'][0]}
		if self.state.name == "idle" and cond_ext_idle_to_processing():
			return ProcessorState(name="processing", **action_ext_idle_to_processing())

		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 = CollectorState(name="waiting", nr_of_jobs=0)
		self.my_ports = {"p_in": self.addInPort("p_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)

	def extTransition(self, my_inputs):
		inputs = {k.getPortName(): v for k, v in my_inputs.iteritems()}
		def cond_ext_waiting_to_waiting():
			return True
		def action_ext_waiting_to_waiting():
			return {"nr_of_jobs": self.state.nr_of_jobs + 1}
		if self.state.name == "waiting" and cond_ext_waiting_to_waiting():
			return CollectorState(name="waiting", **action_ext_waiting_to_waiting())

		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 = {"p_in": self.addInPort("p_in"), "p_out": self.addOutPort("p_out")}
		self.submodels = {}
		self.submodels["p1"] = self.addSubModel(Processor(name="p1"))
		self.submodels["p2"] = self.addSubModel(Processor(name="p2"))

		self.connectPorts(self.my_ports["p_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["p_out"])

class Root(CoupledDEVS):
	def __init__(self, name="Root"):
		CoupledDEVS.__init__(self, name)
		self.my_ports = {}
		self.submodels = {}
		self.submodels["g"] = self.addSubModel(Generator(name="g"))
		self.submodels["cp"] = self.addSubModel(CoupledProcessor(name="cp"))
		self.submodels["p"] = self.addSubModel(Processor(name="p"))
		self.submodels["c"] = self.addSubModel(Collector(name="c"))

		self.connectPorts(self.submodels["g"].my_ports["p_out"], self.submodels["cp"].my_ports["p_in"])
		self.connectPorts(self.submodels["cp"].my_ports["p_out"], self.submodels["p"].my_ports["p_in"])
		self.connectPorts(self.submodels["p"].my_ports["p_out"], self.submodels["c"].my_ports["p_in"])