#  Copyright 2020 Modelling, Simulation and Design Lab (MSDL) at
#  McGill University and the University of Antwerp (http://msdl.cs.mcgill.ca/)
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""PythonPDEVS Queue for storing items that are waiting to be processed."""

from pypdevs.DEVS import AtomicDEVS, CoupledDEVS
from pypdevs.infinity import INFINITY
from pypdevsbbl.extra.heap import heap
from pypdevsbbl.extra.queues import FIFO, QueueElement
from pypdevs.util import DEVSException

class SimpleQueue(AtomicDEVS):
	"""Queue with reneging that allows items to wait for storage.

	Items can only be dequeued when a requestdequeue input is obtained. This
	makes it so the queue don't have to handle the complexities concerning
	automatic dequeueing. See the :class:`Queue` for more details on automatic
	dequeueing.

	Reneging will happen automatically if needs be. The requestrenege will
	renege the last item in the queue.

	The enqueue will send the outputs to the overflow output if capacity >= K.

	Note:
		We will be using a maxheap in this queue.

	.. image:: ../img/queue/simple-queue.svg
		:align: center

	Args:
		name (str):			The name of the Queue.
		fc (def):			Queueing Discipline. I.e. the priority assigning
							function. Takes the current time, the value and a
							unique index as arguments. Defaults to
							:meth:`pypdevsbbl.extra.queues.FIFO`.
		K (int):			The maximal capacity of the Queue.
							Defaults to INFINITY (= no maximal capacity).
		dr (numeric):		Standard reneging time delay. Cannot be negative.
							Defaults to INFINITY (= no reneging).
		contents (iter):	A list of contents to be inserted into the queue on
							startup. Defaults to an empty tuple.
		req_am (bool):		When ``True``, the message in a dequeue/renege
							request indicates the amount that should be released.
							When ``False``, no matter which message arrives on
							these ports, only a single item is released. Defaults
							to ``False``.

	State:
		queue (heap):		The internal heap for this priority queue.
		dr (numeric):		The renege time delay.
		T (numeric):		Current absolute time.
		rd (bool):			If a dequeue operation was requested.
		rr (bool):			If a renege operation was requested.
		ovf (any):			The overflowing items.

	Input Ports:
		enqueue	(any):		Obtains values and queues them.
		dr (numeric):		A new minimal time in the system value. Cannot be
							negative.
		requestdequeue:		Indicates that items must dequeue. When :attr:`req_am`
							is ``True``, it expects a natural number. Otherwise,
							any arriving value is accepted and a single item will
							be dequeued.
		requestrenege:		Indicates that the last item in the queue must depart on
							the `renege` port. When :attr:`req_am` is ``True``, it
							expects a natural number. Otherwise, any arriving value is
							accepted and a single item will be reneged.

	Output Ports:
		dequeue (any):		Outputs the next item if it exists.
		renege (any):		Outputs reneged items.
		overflow (any):		Outputs overflowing items.
	"""
	def __init__(self, name:str, fc=FIFO, K=INFINITY, dr=INFINITY, contents=(), req_am=False):
		assert dr >= 0
		super(SimpleQueue, self).__init__(name)
		self.state = {
			"queue": heap(),
			"dr": dr,
			"T": 0,
			"rd": 0,
			"rr": 0,
			"ovf": None
		}
		self.fc = fc
		self.K = K
		self.req_am = req_am

		self._ix = 0
		for val in contents:
			item = QueueElement(self.fc(self.state["T"], val, self._ix), val, self.state["T"])
			self.state["queue"].insert(item)
			self._ix += 1

		self.enqueue = self.addInPort("enqueue")
		self.dr = self.addInPort("dr")
		self.requestdequeue = self.addInPort("requestdequeue")
		self.requestrenege = self.addInPort("requestrenege")

		self.dequeue = self.addOutPort("dequeue")
		self.renege = self.addOutPort("renege")
		self.overflow = self.addOutPort("overflow")

	def minRen(self, init=INFINITY, it=None):
		"""Obtain the item that has a minimal reneging time.

		Also compare it to the given item.

		Args:
			init (numeric):	The initial value to compare it to.
			it (any):		The item that is linked to the ``init``
							value.

		Returns:
			Tuple in the shape of ``(time, item)`` where ``time`` is
			the time delay until next renege and ``item`` the item that
			must be reneged.
		"""
		dr = self.state["dr"]
		if dr == INFINITY:
			return init, it
		val = init
		m = it
		T = self.state["T"]
		items = self.state["queue"].items()
		for item in items:
			ai = item.arrival_time
			# Python can be way too precise
			t = round(dr - T + ai, 10)
			if t < val:
				val = t
				m = item
		return val, m

	def maxRen(self, init=-INFINITY, it=None):
		"""Obtain the item that has a maximal reneging time.

		Also compare it to the given item.

		Args:
			init (numeric):	The initial value to compare it to.
			it (any):		The item that is linked to the ``init``
							value.

		Returns:
			Tuple in the shape of ``(time, item)`` where ``time`` is
			the time delay until the last renege and ``item`` the item
			that must be reneged last.
		"""
		val = init
		m = it
		items = self.state["queue"].items()
		for item in items:
			t = item.arrival_time
			if t > val:
				val = t
				m = item
		return val, m

	def timeAdvance(self):
		if self.state["ovf"] is not None:
			return 0
		if self.state["queue"].empty():
			return INFINITY
		if self.state["rd"] > 0 or self.state["rr"] > 0:
			return 0
		res, _ = self.minRen()
		if res < 0:
			res = 0
		return res

	def intTransition(self):
		self.state["T"] += self.timeAdvance()
		if self.state["ovf"] is not None:
			self.state["ovf"] = None
		elif not self.state["queue"].empty():
			if self.state["rd"] > 0:
				self.state["queue"].pop()
				self.state["rd"] -= 1
			if self.state["rr"] > 0:
				_, i = self.maxRen()
				self.state["queue"].pop(self.state["queue"].indexOf(i))
				self.state["rr"] -= 1
			else:
				v, m = self.minRen()
				if v <= 0:
					self.state["queue"].pop(self.state["queue"].indexOf(m))
		return self.state

	def extTransition(self, inputs):
		self.state["T"] += self.elapsed
		if self.dr in inputs:
			val = inputs[self.dr]
			if val >= 0:
				self.state["dr"] = val
		if self.enqueue in inputs:
			cap = len(self.state["queue"])
			val = inputs[self.enqueue]
			if cap < self.K:
				item = QueueElement(self.fc(self.state["T"], val, self._ix), val, self.state["T"])
				self.state["queue"].insert(item)
				self._ix += 1
			else:
				self.state["ovf"] = val
		if self.requestdequeue in inputs:
			self.state["rd"] = inputs[self.requestdequeue] if self.req_am else 1
		if self.requestrenege in inputs:
			self.state["rr"] = inputs[self.requestrenege] if self.req_am else 1
		return self.state

	def outputFnc(self):
		res = {}
		if self.state["ovf"] is not None:
			res[self.overflow] = self.state["ovf"]
		elif not self.state["queue"].empty():
			if self.state["rd"] > 0:
				res[self.dequeue] = self.state["queue"].peek().item
			if self.state["rr"] > 0:
				_, i = self.maxRen()
				res[self.renege] = self.state["queue"].peek(self.state["queue"].indexOf(i)).item
			else:
				_, m = self.minRen()
				if m is not None:
					res[self.renege] = self.state["queue"].peek(self.state["queue"].indexOf(m)).item
		return res


class QueueTracker(AtomicDEVS):
	"""Helper block for the :class:`SimpleQueue` to be linked in the :class:`Queue`.

	Instead of using a :class:`ConstantGenerator`, this block will be used to
	periodically request dequeues. It also keeps track of the current queue
	size, meaning that, whenever the queue is empty, it can pause until the
	next input.

	.. image:: ../img/queue/queue-tracker.svg
		:align: center

	Args:
		name (str):			The name for the block.
		dt (numeric):		Standard time delay between messages.
							Defaults to INFINITY. Cannot be zero or negative.
		size (int):			The queue size to start from. Must be positive or
							zero. Defaults to 0.
		value (any):		The value to output instead of the queue size.
							When ``None``, the queue size is used. Defaults to
							``None``.

	State:
		dt (numeric):		The time delay between messages.
		el (numeric):		The elapsed time since the last event.
		size (int):			The current queue size.

	Input Ports:
		dt (numeric):		Updates the time delay.
		enter (any):		An item entered the queue.
		leave (any):		An item left the queue.

	Output Ports:
		size (int):		The queue size or a predefined value.
	"""
	def __init__(self, name: str, dt=INFINITY, size=0, value=None):
		assert 0 < dt and 0 <= size
		super(QueueTracker, self).__init__(name)
		self.state = {
			"dt": dt,
			"el": dt,
			"size": size
		}

		self.value = value

		self.dt = self.addInPort("dt")
		self.enter = self.addInPort("enter")
		self.leave = self.addInPort("leave")
		self.size = self.addOutPort("size")

	def extTransition(self, inputs):
		self.state["el"] += self.elapsed
		if self.dt in inputs:
			value = inputs[self.dt]
			if value > 0:
				self.state["dt"] = value
		if self.enter in inputs:
			self.state["size"] += 1
		if self.leave in inputs:
			if self.state["size"] <= 0:
				raise DEVSException("Impossible for an item to leave if the queue is empty.")
			self.state["size"] -= 1
		return self.state

	def intTransition(self):
		self.state["el"] += self.timeAdvance()
		return self.state

	def timeAdvance(self):
		dt = self.state["dt"]
		if self.state["size"] == 0 or (self.state["el"] == dt == INFINITY):
			return INFINITY
		return dt - (self.state["el"] % dt)

	def outputFnc(self):
		return { self.size: self.state["size"] if self.value is None else self.value }


class Queue(CoupledDEVS):
	"""Queue with reneging that allows items to wait for storage.

	This queue will be a wrapper around the :class:`SimpleQueue` and
	allows for automatic dequeueing by continuously sending messages
	to the `requestdequeue` input.

	The names of the internal :class:`SimpleQueue` and :class:`QueueTracker`
	will be made up of this Queue's name, appended with ``-queue`` and
	``-tracker`` respectively.

	Note:
		If you don't want to use automatic dequeueing, it's more
		efficient to make use of the :class:`SimpleQueue` instead.

	Note:
		We will be using a maxheap in this queue.

	.. image:: ../img/queue/queue.svg
		:align: center

	Args:
		name (str):				The name of the Queue.
		fc (def):				Queueing Discipline. I.e. the priority assigning function.
								Takes the current time, the value and a unique index as
								arguments. Defaults to :meth:`pypdevsbbl.extra.queues.FIFO`.
		K (int):				The maximal capacity of the Queue.
								Defaults to INFINITY (= no maximal capacity).
		dd (numeric):			Standard dequeueing time delay. Cannot be zero or negative.
								Defaults to INFINITY (= no automatic dequeueing).
		dr (numeric):			Standard reneging time delay. Cannot be negative.
								Defaults to INFINITY (= no reneging).
		contents (iter):		A list of contents to be inserted into the queue on startup.
								Defaults to an empty tuple.
		req_am (bool):			When ``True``, the message in a dequeue/renege request
								indicates the amount that should be released. When ``False``,
								no matter which message arrives on these ports, only a single
								item is released. Defaults to ``False``.

	Input Ports:
		enqueue	(any):			Obtains values and queues them.
		dd (numeric):			A new minimal time until dequeue. Cannot be zero or negative.
		dr (numeric):			A new minimal time in the system value. Cannot be negative.
		requestdequeue:			Indicates that items must dequeue. When :attr:`req_am` is
								``True``, it expects a natural number. Otherwise, any arriving
								value is accepted and a single item will be dequeued.
		requestrenege:			Indicates that the last item in the queue must depart on the
								`renege` port. When :attr:`req_am` is ``True``, it expects a
								natural number. Otherwise, any arriving value is accepted and
								a single item will be reneged.

	Output Ports:
		dequeue (any):		Outputs the next item if it exists.
		renege (any):		Outputs reneged items.
		overflow (any):		Outputs overflowing items.
		count (int):		Current queue size.
	"""
	def __init__(self, name:str, fc=FIFO, K=INFINITY, dd=INFINITY, dr=INFINITY, contents=(), req_am=False):
		assert dd > 0 and dr >= 0
		super(Queue, self).__init__(name)
		self._queue = self.addSubModel(SimpleQueue(name + "-queue", fc, K, dr, contents, req_am))
		self._tracker = self.addSubModel(QueueTracker(name + "-tracker", dt=dd, size=len(contents),
													  value=1 if req_am else None))

		self.enqueue = self.addInPort("enqueue")
		self.dd = self.addInPort("dd")
		self.dr = self.addInPort("dr")
		self.requestdequeue = self.addInPort("requestdequeue")
		self.requestrenege = self.addInPort("requestrenege")

		self.dequeue = self.addOutPort("dequeue")
		self.renege = self.addOutPort("renege")
		self.overflow = self.addOutPort("overflow")
		self.count = self.addOutPort("count")

		self.connectPorts(self.enqueue, self._queue.enqueue)
		self.connectPorts(self.dr, self._queue.dr)
		self.connectPorts(self.requestdequeue, self._queue.requestdequeue)
		self.connectPorts(self.requestrenege, self._queue.requestrenege)

		self.connectPorts(self._queue.renege, self.renege)
		self.connectPorts(self._queue.dequeue, self.dequeue)
		self.connectPorts(self._queue.overflow, self.overflow)

		self.connectPorts(self.dd, self._tracker.dt)
		self.connectPorts(self.enqueue, self._tracker.enter)
		self.connectPorts(self._queue.dequeue, self._tracker.leave)
		self.connectPorts(self._queue.renege, self._tracker.leave)
		self.connectPorts(self._queue.overflow, self._tracker.leave)
		self.connectPorts(self._tracker.size, self._queue.requestdequeue)
		self.connectPorts(self._tracker.size, self.count)


class Retain(AtomicDEVS):
	"""Retains an item until the next arrival or next time instance.

	There are two modes:
		- *individual*:		Will only output the previous item.
		- *collection*:		Will collect all arrived items in the same time
		  instance and will output a list of all these items.

	.. image:: ../img/queue/retain.svg
		:align: center

	Args:
		name (str):			The name of the block.
		ind (bool):			When true, this block uses the "individual" mode.
							Otherwise, the "collection" mode is used.

	State:
		old (any):			The previous item (or a list thereof).
		cur (any):			The current item.

	Input Ports:
		input (any):		An item that needs to be remembered.

	Output Ports:
		previous (any):		The previous item (or a list thereof).
		current (any):		The current item.
	"""
	def __init__(self, name: str, ind: bool = True):
		super(Retain, self).__init__(name)
		self.state = {
			"old": None if ind else [],
			"cur": None
		}

		self.ind = ind

		self.input = self.addInPort("input")
		self.previous = self.addOutPort("previous")
		self.current = self.addOutPort("current")

	def extTransition(self, inputs):
		if self.input in inputs:
			self.state["cur"] = inputs[self.input]
			if not self.ind and self.elapsed > 0:
				self.state["old"].clear()
		return self.state

	def intTransition(self):
		if self.ind:
			self.state["old"] = self.state["cur"]
		else:
			self.state["old"].append(self.state["cur"])
		self.state["cur"] = None
		return self.state

	def timeAdvance(self):
		if self.state["cur"] is not None:
			return 0
		return INFINITY

	def outputFnc(self):
		if self.ind:
			prev = self.state["old"]
		else:
			prev = self.state["old"][:]

		if self.state["cur"] is not None:
			return {
				self.previous: prev,
				self.current: self.state["cur"]
			}
		return {}


class Advance(AtomicDEVS):
	"""Reneging queue where each item holds their own delay.

	This block is inspired by the ``ADVANCE`` block from GPSS.

	Note:
		The input and the delay must arrive at exactly the same time.
		When in doubt, use the :class:`pypdevsbbl.generic.routing.Sync`
		block to ensure this happens.

	Warning:
		The items traveling through must be unique!

	Warning:
		When using pausing, it is pertinent that the items traveling
		through the block are hashable.

	.. image:: ../img/queue/advance.svg
		:align: center

	Args:
		name (str):			The name of the block.

	State:
		items (dict):		All the scheduled items in the "queue", with
							their expected delay.
		least (tuple):		The item to leave the block first, in the form
							of ``(item, time)``.
		paused (set):		Set of items that are currently paused. Note
							that the items must be hashable.

	Input Ports:
		input (any):		The item to schedule. Must arrive at the same
							time as the `delay` input.
		delay (numeric):	The delay of the item. Must arrive at the same
							time as the `input` message.
		pause (list):		A list of tuples ``(item, bool)`` indicative of
							whether or not the item must pause.

	Output Ports:
		output (any):		The item, at the pre-scheduled delay.
	"""
	def __init__(self, name):
		super(Advance, self).__init__(name)

		self.state = {
			"items": {},
			"least": None,
			"paused": set()
		}

		self.input = self.addInPort("input")
		self.delay = self.addInPort("delay")
		self.pause = self.addInPort("pause")
		self.output = self.addOutPort("output")

	def least(self):
		tc = [x for x in self.state["items"].items() if x[0] not in self.state["paused"]]
		if len(tc) == 0:
			self.state["least"] = None
		else:
			self.state["least"] = min([x for x in self.state["items"].items() if x[0] not in self.state["paused"]],
									  key=lambda x: x[1])

	def timeAdvance(self):
		if self.state["least"] is None:
			return INFINITY
		return self.state["least"][1]

	def outputFnc(self):
		if len(self.state["items"]) == 0:
			return {}
		return { self.output: self.state["least"][0] }

	def intTransition(self):
		ta = self.timeAdvance()
		del self.state["items"][self.state["least"][0]]
		for item in self.state["items"]:
			if item not in self.state["paused"]:
				self.state["items"][item] -= ta
		self.least()
		return self.state

	def extTransition(self, inputs):
		for item in self.state["items"]:
			if item not in self.state["paused"]:
				self.state["items"][item] -= self.elapsed
		if self.pause in inputs:
			for item, s in inputs[self.pause]:
				if s:
					self.state["paused"].add(item)
				elif item in self.state["paused"]:
					self.state["paused"].remove(item)
		elif self.input in inputs and self.delay in inputs:
			self.state["items"][inputs[self.input]] = inputs[self.delay]
		else:
			raise RuntimeError("Invalid input configuration.")
		self.least()
		return self.state
