#  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.

"""A set of building blocks that can be used for generating data.

When using the classes in this file, you can start creating simple
(and complex) flow systems or, generally, anything that requires
input.
"""

# TODO: give all blocks a maximum amount of items it can generate

from pypdevs.DEVS import AtomicDEVS
from pypdevs.infinity import INFINITY


class ConstantGenerator(AtomicDEVS):
	"""AtomicDEVS that generates a constant value every dt time.

	.. image:: ../img/generators/constant-generator.svg
		:align: center

	Args:
		name (str):			The name of the generator.
		value (any):		The item to be generated. Defaults to 0.
		dt (numeric):		Time delay between the generation of new events.
							Defaults to 1.

	State:
		value (any):		The item to generate.
		dt (numeric):		Time delay between the generation of new events.
							Cannot be negative or zero.
		elapsed (numeric):	The elapsed time since last internal transition.

	Input Ports:
		value (any):		A new value to output.
		dt (numeric):		A new time delay. Cannot be negative or zero.

	Output Ports:
		output	(any):		Outputs the generated value.
	"""
	def __init__(self, name: str, value=0, dt=1):
		assert dt > 0
		super(ConstantGenerator, self).__init__(name)

		self.state = {
			"dt": dt,
			"value": value,
			"elapsed": dt
		}

		self.value = self.addInPort("value")
		self.dt = self.addInPort("dt")
		self.output = self.addOutPort("output")

	def intTransition(self):
		self.state["elapsed"] = 0
		return self.state

	def extTransition(self, inputs):
		self.state["elapsed"] += self.elapsed
		if self.value in inputs:
			self.state["value"] = inputs[self.value]
		if self.dt in inputs:
			dt = inputs[self.dt]
			if 0 < dt:
				self.state["dt"] = dt
		return self.state

	def timeAdvance(self):
		return max(self.state["dt"] - self.state["elapsed"], 0.0)

	def outputFnc(self):
		return { self.output: self.state["value"] }


class FunctionGenerator(AtomicDEVS):
	"""AtomicDEVS that generates a new value based on the previous one and the elapsed time.

	Note:
		By default, he first output ``ic`` (initial condition) will be outputted at time 0.

	.. image:: ../img/generators/function-generator.svg
		:align: center

	Args:
		name (str):			The name of the generator.
		function (def):		A function that takes the previous value and the elapsed time as
							arguments and returns a new item from those.
							Defaults to an incrementor of the ic.
		ic (any):			The initial condition of the generator.
							Defaults to 0.
		dt (numeric):		Time delay between the generation of new events.
							Defaults to 1.

	State:
		next (any):			The next item that will be generated.
		time (numeric):		The amount of time that has elapsed.
		dt (numeric):		Time delay between the generation of new events.
							Cannot be negative or zero. Defaults to 1.
		elapsed (numeric):	The elapsed time since last internal transition.

	Input Ports:
		dt (numeric):		A new time delay. Cannot be negative or zero.

	Output Ports:
		output	(any):		Outputs the generated value.
	"""
	def __init__(self, name: str, function=lambda p, t: p+1, ic=0, dt=1):
		assert dt > 0
		super(FunctionGenerator, self).__init__(name)

		self.state = {
			"dt": dt,
			"elapsed": dt,
			"time": -dt,
			"prev": None
		}
		self.ic = ic
		self.function = function

		self.dt = self.addInPort("dt")
		self.output = self.addOutPort("output")

	def timeAdvance(self):
		return max(self.state["dt"] - self.state["elapsed"], 0.0)

	def extTransition(self, inputs):
		self.state["elapsed"] += self.elapsed
		if self.dt in inputs:
			dt = inputs[self.dt]
			if 0 < dt:
				self.state["dt"] = dt
		return self.state

	def intTransition(self):
		self.state["elapsed"] = 0
		self.state["time"] += self.state["dt"]
		self.state["prev"] = self.ic if self.state["prev"] is None else self.function(self.state["prev"], self.state["time"])
		return self.state

	def outputFnc(self):
		return {
			self.output: self.ic if self.state["prev"] is None else self.function(self.state["prev"], self.state["time"] + self.state["dt"])
		}


class TableGenerator(AtomicDEVS):
	"""AtomicDEVS that generates a new value based on the table.

	.. image:: ../img/generators/table-generator.svg
		:align: center

	Args:
		name (str):			The name of the generator.
		table (iter):		An iterable containing (time, value) tuples/pairs.
							Defaults to empty list.
		repeat (bool):		When ``True``, the table will be repeated when it
							is finished. Defaults to ``False``

	State:
		time (numeric):		The amount of time until the next transition.
		index (int):		The index in the table at this point.

	Output Ports:
		output	(any):		Outputs the generated value.
	"""
	def __init__(self, name: str, table, repeat: bool = False):
		super(TableGenerator, self).__init__(name)

		self.repeat = repeat
		self.table = table
		self.state = {
			"time": self.table[0][0] if len(table) >= 1 else INFINITY,
			"index": 0
		}

		self.output = self.addOutPort("output")

	def timeAdvance(self):
		return self.state["time"]

	def intTransition(self):
		i1 = self.state["index"]
		i2 = i1 + 1
		time = INFINITY
		if self.repeat and len(self.table) <= i2:
			i2 = 0
			time = self.table[0][0]
		elif len(self.table) > i2 and self.table[i2][0] < INFINITY:
			time = self.table[i2][0] - self.table[i1][0]
		self.state["index"] = i2
		self.state["time"] = time
		return self.state

	def outputFnc(self):
		if len(self.table) > self.state["index"]:
			return { self.output: self.table[self.state["index"]][1] }
		return {}


class SingleFire(AtomicDEVS):
	"""Output an item once.

	Useful to prevent the overhead of a :class:`TableGenerator` with a
	single record.

	.. image:: ../img/generators/single-fire.svg
		:align: center

	Args:
		name (str):		The name of the block.
		item (any):		The item that needs to be outputted.
						Cannot be ``None``.
		t (numeric):	Absolute simulation time when to fire the event.

	State:
		item (any):		The item to fire.
		t (numeric):	The absolute time until the event or ``INFINITY``
						if the event has already passed.

	Output Ports:
		output (any):	The item.
	"""
	def __init__(self, name: str, item, t = 0.0):
		assert item is not None
		super(SingleFire, self).__init__(name)
		self.state = {
			"item": item,
			"t": t
		}

		self.output = self.addOutPort("output")

	def timeAdvance(self):
		return self.state["t"]

	def intTransition(self):
		self.state["t"] = INFINITY
		return self.state

	def outputFnc(self):
		return { self.output: self.state["item"] }


class BulkGenerator(AtomicDEVS):
	"""Creates :math:`n` messages every time it receives :math:`n`.

	.. image:: ../img/generators/bulk-generator.svg
		:align: center

	Args:
		name (str):			The name of the block.
		dt (numeric):		The time delay between the messages. Cannot be
							negative. Defaults to 0.

	State:
		index (int):		The index of the last item that was outputted.
		items (list):		The items to generate.
		dt (numeric):		Time delay between the generation of new events.
							Cannot be negative or zero.
		elapsed (numeric):	The elapsed time since last internal transition.

	Input Ports:
		input (int):		A amount of items to output in bulk. When there
							are still items remaining, the provided values
							are appended to the list.
		dt (numeric):		A new time delay. Cannot be negative.

	Output Ports:
		output (numeric):	Outputs the generated values.
	"""

	def __init__(self, name: str, dt = 0.0):
		assert dt >= 0
		super(BulkGenerator, self).__init__(name)

		self.state = {
			"index": -1,
			"items": [],
			"dt": dt,
			"elapsed": dt
		}

		self.input = self.addInPort("input")
		self.dt = self.addInPort("dt")
		self.output = self.addOutPort("output")

	def intTransition(self):
		self.state["elapsed"] = 0.0
		self.state["items"].pop(0)
		self.state["index"] += 1
		return self.state

	def extTransition(self, inputs):
		self.state["elapsed"] += self.elapsed
		if self.dt in inputs:
			dt = inputs[self.dt]
			if 0 <= dt:
				self.state["dt"] = dt
		if self.input in inputs:
			n = int(inputs[self.input])
			ix = self.state["index"]
			self.state["items"] += list(range(ix + 1, ix + n + 1))
		return self.state

	def timeAdvance(self):
		if len(self.state["items"]) > 0:
			return max(self.state["dt"] - self.state["elapsed"], 0.0)
		return INFINITY

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


from pypdevsbbl.extra.rngstreams import createStream, U01
class RandomNumberGenerator(AtomicDEVS):
	""":class:`AtomicDEVS` that generates a random number every dt time.

	It makes use of the RngStreams random number generator i.i.d. (0, 1).
	Using the ``dist`` argument, one can simply change the distribution over which
	numbers need to be generated. The :mod:`pypdevsbbl.extra.rngstreams.distributions`
	module provides an extensive implementation of a lot of distributions.

	See Also:
		- :class:`pypdevsbbl.generic.math.Random`
		- :class:`RandomDelayGenerator`
		- :mod:`pypdevsbbl.extra.rngstreams`
		- :mod:`pypdevsbbl.extra.rngstreams.distributions`

	.. image:: ../img/generators/rng.svg
		:align: center

	Args:
		name (str):			The name of the generator (both the block and the engine).
		seed (numeric):		The seed for the generator. When ``None``, it will not set a
							seed. Defaults to ``None``.
		dist (def):			A stateless function that implements an inversive cumulative
							distribution function. It takes at least a single argument,
							which falls in (0, 1).
							Defaults to a uniform distribution over (0, 1).
		args (tuple):		Optional additional arguments to add to the distribution.
		dt (numeric):		Time delay between the generation of new events.
							Defaults to 1. Cannot be negative or zero.

	State:
		dt (numeric):		Time delay between the generation of new events.
		num (numeric):		The random number that's being/needs to be outputted.
		elapsed (numeric):	The elapsed time since last internal transition.

	Input Ports:
		dt (numeric):		A new time delay. Cannot be negative or zero.

	Output Ports:
		output	(numeric):	Outputs the generated value.
	"""
	def __init__(self, name: str, seed=None, dist=lambda x: x, args=(), dt=1):
		assert dt > 0
		super(RandomNumberGenerator, self).__init__(name)

		self.dist = dist
		self.args = args
		self.gen = createStream(name, seed)

		self.state = {
			"dt": dt,
			"num": self.next(),
			"elapsed": dt
		}
		self.dt = self.addInPort("dt")
		self.output = self.addOutPort("output")

	def timeAdvance(self):
		return max(self.state["dt"] - self.state["elapsed"], 0)

	def extTransition(self, inputs):
		self.state["elapsed"] += self.elapsed
		if self.dt in inputs:
			dt = inputs[self.dt]
			if 0 < dt:
				self.state["dt"] = dt
		return self.state

	def intTransition(self):
		self.state["elapsed"] = 0.0
		self.state["num"] = self.next()
		return self.state

	def outputFnc(self):
		return { self.output: self.state["num"] }

	def next(self):
		"""Obtains a random variable w.r.t. this generator's config.

		Do note that this updates the gen object, meaning it must only be called
		in functions that may update the state.
		"""
		return self.dist(U01(self.gen), *self.args)


class RandomDelayGenerator(AtomicDEVS):
	""":class:`AtomicDEVS` that generates a new item every random time.

	The items that are generated get a sequential numbering (which can be altered
	using the :attr:`func` argument), but their inter-arrival times follow the
	provided distribution.

	It makes use of the RngStreams random number generator i.i.d. (0, 1).
	Using the ``dist`` argument, one can simply change the distribution over which
	numbers need to be generated. The :mod:`pypdevsbbl.extra.rngstreams.distributions`
	module provides an extensive implementation of a lot of distributions.

	See Also:
		- :class:`RandomNumberGenerator`
		- :class:`pypdevsbbl.generic.math.Random`
		- :mod:`pypdevsbbl.extra.rngstreams`
		- :mod:`pypdevsbbl.extra.rngstreams.distributions`

	.. image:: ../img/generators/random-delay-generator.svg
		:align: center

	Args:
		name (str):			The name of the generator (both the block and the engine).
		seed (numeric):		The seed for the generator. When ``None``, it will not set a
							seed. Defaults to ``None``.
		dist (def):			A stateless function that implements an inversive cumulative
							distribution function. It takes at least a single argument,
							which falls in (0, 1).
							Defaults to a uniform distribution over (0, 1).
		args (tuple):		Optional additional arguments to add to the distribution.
		func (def):			Transformation function for the newly generated item.
		dt (numeric):		Time of first item to generate. When ``None``, no output
							will be created at time 0. Defaults to 0.
		max (numeric):		Maximal amount of items to generate.
		base (numeric):		The base number to start from when assigning the sequential
							numbering. When ``None``, the block's hash is used. Defaults
							to 0.

	State:
		dt (numeric):		Time delay between the generation of new events.
		num (numeric):		The random number that's being/needs to be outputted.
		elapsed (numeric):	The elapsed time since last internal transition.

	Input Ports:
		halt (bool):		Halts the generator until further notice.

	Output Ports:
		output (any):		Outputs the value.
	"""
	# TODO: test
	def __init__(self, name: str, seed=None, dist=lambda x: x, args=(), func=lambda x, t: x, dt=0.0,
				 max=INFINITY, base=0):
		assert dt is None or dt >= 0
		super(RandomDelayGenerator, self).__init__(name)

		self.dist = dist
		self.args = args
		self.func = func
		self.gen = createStream(name, seed)

		if dt is None:
			dt = 0.0
			while dt == 0.0:
				dt = self.next()

		self.state = {
			"time": 0.0,
			"dt": dt,
			"num": base if base is not None else id(self),
			"halted": False,
			"max": max
		}
		self.halt = self.addInPort("halt")
		self.output = self.addOutPort("output")

	def timeAdvance(self):
		if self.state["halted"] or self.state["max"] == 0:
			return INFINITY
		return self.state["dt"]

	def extTransition(self, inputs):
		self.state["time"] += self.elapsed
		if self.halt in inputs:
			self.state["halted"] = inputs[self.halt]
			self.state["dt"] = self.next()
		return self.state

	def intTransition(self):
		self.state["time"] += self.state["dt"]
		self.state["num"] += 1
		self.state["max"] -= 1
		if self.state["max"] > 0:
			self.state["dt"] = self.next()
		return self.state

	def outputFnc(self):
		return { self.output: self.func(self.state["num"], self.state["time"] + self.timeAdvance()) }

	def next(self):
		"""Obtains a random variable w.r.t. this generator's config.

		Do note that this updates the gen object, meaning it must only be called
		in functions that may update the state.
		"""
		return self.dist(U01(self.gen), *self.args)
