# Source code for pypdevsbbl.generic.collectors

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

    """Set of PythonPDEVS collectors for keeping track of all data.

    The classes in this file will allow for obtaining some well
    required statistical data from your models.
    """

    from pypdevs.DEVS import AtomicDEVS
    from pypdevs.infinity import INFINITY
    from pypdevsbbl.extra.mathutils import histostep, draw_boxplot

    class BaseCollector(AtomicDEVS):
    	"""Base class for all the collectors.

    	This building block itself is not useful, but mainly acts as a root
    	from which all collectors can inherit.

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

    	Input Ports:
    		input (any):	Obtains values.
    		clear (any):	Clears the contents of the block.
    	"""
    	def __init__(self, name: str):
    		super(BaseCollector, self).__init__(name)
    		self.state = {}
    		self.input = self.addInPort("input")
    		self.clear = self.addInPort("clear")

    	def clearContents(self):
    		"""Clears the contents of the block."""
    		raise NotImplementedError()

    	def plot(self, axis, arg):
    		"""Plots the data of the collector.

    		Args:
    			axis:	The :mod:`matplotlib` axis to plot on.
    			arg:	Additional argument.
    		"""
    		raise NotImplementedError()

    class Collector(BaseCollector):
    	"""Collects base information about items that arrive on `input`.

    	The Collector block keeps track of the basic information about all items that
    	arrive. This includes minima, maxima, means, sums, sums of squares and the
    	amount of items that it has seen.

    	Using this data, the min, max, mean, variance, standard derivation and count
    	of the passed items can be obtained. See their respective functions below.

    	Note:
    	 	In order to do this, the items must be numeric and therefore have the
    		``__lt__``, ``__gt__``, ``__add__``, ``__pow__`` and ``__div__`` operators
    		available.

    	.. image:: ../paper/images/blocks/collectors/collector.svg
    		:align: center

    	Args:
    		name (str):			The name for the Collector.
    		positive (bool):	When ``True``, only data from strict positive values will be
    							captured. Defaults to ``False``.

    	State:
    		min (numeric):		The minimal value that arrived. When no items passed
    							through, this will be INFINITY.
    		max (numeric):		The maximal value that arrived. When no items passed
    							through, this will be -INFINITY.
    		n (int):			The amount of items that have arrived.
    		sum (numeric):		The sum of all values that have arrived.
    		squares (numeric):	The sum of the squares of all values that have arrived.

    	Input Ports:
    		input	(numeric):	Obtains values.
    							Note that INFINITY inputs will shape the data in an undefined way.
    		clear (any):		Clears the contents of the block.
    	"""
    	def __init__(self, name: str, positive: bool = False):
    		super(Collector, self).__init__(name)
    		self.positive = positive
    		self.clearContents()

    	def plot(self, axis, cols=1):
    		"""Plots the data of the collector.

    		Args:
    			axis:	The :mod:`matplotlib` axis to plot on.
    			cols:	How many columns to split the info into.
    					Defaults to 1.
    		"""
    		axis.clear()
    		axis.set_xlim(0, 1)
    		axis.set_ylim(0, 8)
    		axis.get_xaxis().set_visible(False)
    		axis.get_yaxis().set_visible(False)
    		axis.set_title(self.name + " Results")
    		annots = [
    			"$n = %i$" % self.count(),
    			"$\\min = %.5f$" % self.min(),
    			"$\\max = %.5f$" % self.max(),
    			"$\\sum_{i=1}^{%i} i = %.5f$" % (self.count(), self.sum()),
    			"$\\sum_{i=1}^{%i} i^2 = %.5f$" % (self.count(), self.sumOfSquares()),
    			"$\\mu = %.5f$" % self.mean(),
    			"$\\sigma = %.5f$" % self.std(),
    			"$\\sigma^2 = %.5f$" % self.variance()
    		]
    		annots.reverse()

    		for ix in range(len(annots)):
    			ann = annots[ix]
    			pos = (0.05 + (ix % cols), 0.5 + (ix // cols))
    			axis.annotate(ann, pos)

    	def extTransition(self, inputs):
    		if self.clear in inputs:
    			self.clearContents()
    		if self.input in inputs:
    			value = inputs[self.input]
    			if not self.positive or value > 0:
    				if value < self.state["min"]:
    					self.state["min"] = value
    				if value > self.state["max"]:
    					self.state["max"] = value
    				self.state["n"] += 1
    				self.state["sum"] += value
    				self.state["squares"] += value ** 2
    		return self.state

    	def clearContents(self):
    		self.state["min"] = INFINITY
    		self.state["max"] = -INFINITY
    		self.state["n"] = 0
    		self.state["sum"] = 0
    		self.state["squares"] = 0

    	def min(self):
    		"""Conveniance method to obtain the minimal value.

    		Returns:
    			The minimal value that arrived.
    		"""
    		return self.state["min"]

    	def max(self):
    		"""Conveniance method to obtain the maximal value.

    		Returns:
    			The maximal value that arrived.
    		"""
    		return self.state["max"]

    	def count(self):
    		"""Conveniance method to obtain the amount of items.

    		Returns:
    			The amount of items that have arrived.
    		"""
    		return self.state["n"]

    	def sum(self):
    		"""Conveniance method to obtain the sum.

    		Returns:
    			The sum of the items that have arrived.
    		"""
    		return self.state["sum"]

    	def sumOfSquares(self):
    		"""Conveniance method to obtain the sum of squares.

    		Returns:
    			The sum of squares of the items that have arrived.
    		"""
    		return self.state["squares"]

    	def mean(self):
    		"""Get the mean of all values that have arrived.

    		It is computed as :math:`\\mu = \\dfrac{\\sum_i i}{n}`, with
    		:math:`i` an item that arrived.

    		Returns:
    			The mean or 0 if no items have arrived.
    		"""
    		n = self.count()
    		if n  == 0:
    			return 0
    		return self.sum() / n

    	def variance(self):
    		"""Get the variance of all values that have arrived.

    		It is computed as
    		:math:`\\sigma^2 = \\dfrac{\\sum_i i^2}{n} + \\mu^2 - \\dfrac{2\\mu\\sum_i i}{n}`,
    		with :math:`i` an item that arrived.

    		Returns:
    			The variance or 0 if no items have arrived.
    		"""
    		mu = self.mean()
    		n = self.count()
    		s = self.sum()
    		sq = self.sumOfSquares()
    		if n == 0:
    			return 0
    		return sq / n + mu ** 2 - (2 * mu * s) / n

    	def std(self):
    		"""Get the standard derivation of all values that have arrived.

    		It is computed as :math:`\\sigma = \\sqrt{v}`, with
    		:math:`v` the result from the :meth:`variance` function.

    		Returns:
    			The standard derivation or 0 if no items have arrived.
    		"""
    		return self.variance() ** 0.5

    class EstimateCollector(BaseCollector):
    	"""Collects information about items that arrive.

    	The EstimateCollector block is a block that keeps track of everything that
    	passes through, allowing it to yield important information.

    	Whereas the :class:`Collector` can collect general information data, this
    	collector will use the `P-squared Algorithm for Dynamic Calculation of
    	Quantiles and Histograms without Scoring Observations`.

    	It has a :class:`CollectorState` class, which is used to handle the information
    	of passed-through items. Besides the ``PythonPDEVS``-functionality, it also
    	has a function that can obtain all boxplot- and histogram-data, called
    	:meth:`CollectorState.histoplot`.

    	See Also:
    		- :meth:`pypdevsbbl.extra.mathutils.histogram`
    		- :meth:`pypdevsbbl.extra.mathutils.histostep`

    	.. image:: ../paper/images/blocks/collectors/estimate-collector.svg
    		:align: center

    	Args:
    		name (str):			The name of the EstimateCollector.
    		b (int):			The amount of cells you want for the histogram.
    							4 is often used, because it yields [min, Q1, median, Q3, max].
    							When 1 or less, the histogram will only keep track of the min
    							and the max. Defaults to 4.

    	:State:		See :class:`EstimateCollector.CollectorState`.

    	Input Ports:
    		input	(numeric):	Obtains values.
    							Note that ``INFINITY`` inputs will shape the data in an undefined way.
    		clear (any):		Clears the contents of the block.
    	"""

    	class CollectorState:
    		"""Helperclass for the EstimateCollector.

    		Note:
    			This class allows for the creation of boxplots and histogram, thanks to the
    			:mod:`pypdevsbbl.extra.mathutils` module. This makes use of the ``_hq``, ``_hq0``,
    			``_hn`` and ``_hn0`` attributes. Changing these will make it so the histogram,
    			min, max, median, Q1 and Q3 also changes, so please... **DON'T**.

    		For the sake of the P^2 algorithm, b (number of cells) is chosen to be 4.

    		Attributes:
    			N (int):		The total number of items that have arrived.
    			N0 (int):		The total number of non-zero items that have arrived.

    			_b (int):		The amount of cells in the histogram.
    			_hq (list):		Datavalues of the bars in the histogram of all items.
    							**DON'T CHANGE THIS LIST!**
    			_hq0 (list):	Datavalues of the bars in the histogram of all non-zero items.
    							**DON'T CHANGE THIS LIST!**
    			_hn (list):		Positions of the bars in the histogram of all items.
    							**DON'T CHANGE THIS LIST!**
    			_hn0 (list):	Positions of the bars in the histogram of all non-zero items.
    							**DON'T CHANGE THIS LIST!**

    		Args:
    			b (int):		The amount of cells you want for the histogram.
    							When 1 or less, the histogram will only keep track of the min and
    							the max. Defaults to 4.
    		"""

    		def __init__(self, b=4):
    			self.N = 0
    			self.N0 = 0

    			# Lists for the P^2 Algorithm for Dynamic Calculation of Quantiles and Histograms without Scoring Observations
    			#	See mathutils.histogram and mathutils.histostep
    			self._b = 1 if b < 2 else b
    			self._hq = []
    			self._hq0 = []
    			self._hn = list(range(self._b + 1))
    			self._hn0 = list(range(self._b + 1))

    		def __str__(self):
    			return "CollectorState { N = %i; N0 = %i}" % (self.N, self.N0)

    		def __getitem__(self, item):
    			if item in ["N", "N0"]:
    				return getattr(self, item)
    			return getattr(self, "_" + item)

    		def __setitem__(self, key, value):
    			raise KeyError("Canot set items externally.")

    		def _hstep(self, val):
    			"""Do the histostep function of the mathutils on this state.

    			Additionally, this will do some size checks to make sure histostep can be executed correctly.
    			The general histogram will change.

    			See Also:
    				:meth:`pypdevsbbl.extra.mathutils.histogram`

    			Warning:
    				Only use this function if you know what you're doing.

    			Args:
    				val (numeric):	A value that passes through and must update the histogram.
    			"""
    			if self.N <= len(self._hn):
    				self._hq.append(val)
    				self._hq.sort()
    			else:
    				histostep(val, self._b, self._hq, self._hn, self.N)

    		def _hstep0(self, val):
    			"""Do the histostep function of the mathutils on this state.

    			Additionally, this will do some size checks to make sure histostep can be executed correctly.
    			The non-zero histogram will change.

    			See Also:
    				:meth:`pypdevsbbl.extra.mathutils.histogram`

    			Warning:
    				Only use this function if you know what you're doing.

    			Args:
    				val (numeric):	A non-zero value that passes through and must update the histogram.
    			"""
    			if self.N <= len(self._hn):
    				self._hq0.append(val)
    				self._hq0.sort()
    			else:
    				histostep(val, self._b, self._hq0, self._hn0, self.N0)

    		def min(self):
    			"""Get the smallest value that has arrived.

    			Returns:
    				The smallest value, or INFINITY if no values have arrived yet.
    			"""
    			if self.N > 0:
    				return self._hq[0]
    			return INFINITY

    		def min0(self):
    			"""Get the smallest non-zero value that has arrived.

    			Returns:
    				The smallest non-zero value, or INFINITY if no values have arrived yet.
    			"""
    			if self.N > 0:
    				return self._hq0[0]
    			return INFINITY

    		def max(self):
    			"""Get the largest value that has arrived.

    			Returns:
    				The largest value, or -INFINITY if no values have arrived yet.
    			"""
    			if self.N > 0:
    				return self._hq[-1]
    			return -INFINITY

    		def max0(self):
    			"""Get the largest non-zero value that has arrived.

    			Returns:
    				The largest non-zero value, or -INFINITY if no values have arrived yet.
    			"""
    			if self.N > 0:
    				return self._hq0[-1]
    			return -INFINITY

    		def histoplot(self):
    			"""Obtain all boxplot and histogram data from the state.

    			This function is meant to be used as an unpacking of all data and only
    			reads from the state. In fact, its result is a *copy* of the actual data,
    			making sure it remains immutable.

    			Returns:
    				A 4-tuple (q, h, q0, h0), where
    					* q (list):		Datavalues of the bars in the histogram of all items.
    					* h (list):		Heights of the bars in the histogram, w.r.t. the item count.
    					* q0 (list):	Datavalues of the bars in the histogram of all non-zero items.
    					* h0 (list):	Heights of the bars in the histogram, w.r.t. the non-zero item count.

    			Examples:
    				*	To plot the histogram of all items, use q as x-values and h as its y-values.
    				*	To plot the histogram of all non-zero items, use q0 as x-values and h0 as its y-values.
    				*	To plot the boxplot of all items, unpack q and use it as input in the
    					mathutils.draw_boxplot function (and plot it thereafter).
    				*	To plot the boxplot of all non-zero items, unpack q0 and use it as input in the
    					mathutils.draw_boxplot function (and plot it thereafter).
    			"""
    			q = self._hq
    			q0 = self._hq0
    			n = self._hn
    			n0 = self._hn0
    			return q, [] if self.N == 0 else [x / self.N for x in n],\
    				   q0, [] if self.N0 == 0 else [x / self.N0 for x in n0]

    	def __init__(self, name: str, b: int = 4):
    		super(EstimateCollector, self).__init__(name)
    		self.state = EstimateCollector.CollectorState(b)

    	def plot(self, axis, type='hist'):
    		"""Plots the data of the collector.

    		Args:
    			axis:	The :mod:`matplotlib` axis to plot on.
    			type:	A string indicating which plot to create.
    					Defaults to `'hist'`.

    					For the histogram, use any of
    					``['hist', 'histogram', 'bars', 'H', 'h', 0]``.

    					For the boxplot, use any of
    					``['bxpt', 'boxplot', 'box', 'B', 'b', 1]``
    		"""
    		if type in ['hist', 'histogram', 'bars', 'H', 'h', 0]:
    			self.plot_histogram(axis)
    		elif type in ['bxpt', 'boxplot', 'box', 'B', 'b', 1]:
    			self.plot_boxplot(axis)
    		else:
    			raise ValueError("Unknown type '%s'." % str(type))

    	def plot_boxplot(self, axis):
    		"""Plots the data of the collector as a boxplot.

    		Args:
    			axis:	The :mod:`matplotlib` axis to plot on.
    		"""
    		q, h, _, _ = self.state.histoplot()
    		axis.clear()
    		axis.set_title(self.name + " Boxplot")
    		if len(q) == 5:
    			bbx = draw_boxplot(q[0], q[1], q[2], q[3], q[4])
    			axis.bxp(bbx)
    		else:
    			axis.boxplot(q)

    	def plot_histogram(self, axis):
    		"""Plots the data of the collector as a histogram.

    		Args:
    			axis:	The :mod:`matplotlib` axis to plot on.
    		"""
    		axis.clear()
    		axis.set_title(self.name + " Histogram of Percentiles")
    		q, h, _, _ = self.state.histoplot()
    		axis.bar(["%.5f" % x for x in q], h[:len(q)], ec='black')

    	def extTransition(self, inputs):
    		if self.clear in inputs:
    			self.clearContents()
    		if self.input in inputs:
    			val = inputs[self.input]
    			self.state.N += 1
    			if val > 0:
    				self.state.N0 += 1
    				self.state._hstep0(val)
    			self.state._hstep(val)
    		return self.state

    	def clearContents(self):
    		self.state = EstimateCollector.CollectorState(self.state._b)

    class TableCollector(BaseCollector):
    	"""Builds a table of all collected data.

    	Contrary to the other collectors, the TableCollector will create
    	a table from ALL data that arrived. This implies that the usage
    	of this component will allocate a lot of memory and, if you're using
    	mere numerical data, will be overkill. Nevertheless, if your data is
    	not numerical, or you require to know when certain data arrived, this
    	collector is the one you need.

    	For every item that arrives, it will add another record in a table,
    	which is accessible as a dictionary (see the :meth:`table` method).

    	.. image:: ../paper/images/blocks/collectors/table-collector.svg
    		:align: center

    	Args:
    		name (str):		The name for the TableizeCollector.

    	State:
    		table (list):	The table, as a pair/tuple of (`elapsed time`, `item`)

    	Input Ports:
    		input (any):	Obtains values.
    		clear (any):	Clears the contents of the block.
    	"""
    	def __init__(self, name: str):
    		super(TableCollector, self).__init__(name)
    		self.clearContents()

    	def clearContents(self):
    		self.state["table"] = []
    		self.state["elapsed"] = 0.0

    	def extTransition(self, inputs):
    		self.state["elapsed"] += self.elapsed
    		if self.clear in inputs:
    			self.clearContents()
    		if self.input in inputs:
    			item = inputs[self.input]
    			self.state["table"].append((self.state["elapsed"], item))
    		return self.state

    	def plot(self, axis, kwargs=None):
    		"""Plots the data of the collector.

    		Args:
    			axis:	The :mod:`matplotlib` axis to plot on.
    			kwargs:	Dictionary of keyword arguments to pass to the
    					:meth:`matplotlib.axes.Axes.plot` function.
    					Defaults to ``None`` (= no arguments).
    		"""
    		data = self.table()
    		axis.clear()
    		axis.set_title(self.name + " Results")
    		axis.set_xlabel("time")
    		axis.set_ylabel("value")
    		kw = kwargs if kwargs is not None else {}
    		axis.plot([x[0] for x in data], [x[1] for x in data], **kw)

    	def table(self):
    		"""Fetches the table that was built upto this point in time.

    		This table will be returned as a dict, in the form of
    			`elapsed time` -> `item`

    		Returns:
    			The table.
    		"""
    		return self.state["table"]

    class Counter(BaseCollector):
    	"""Counts the amount of arrived messages.

    	.. image:: ../paper/images/blocks/collectors/counter.svg
    		:align: center

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

    	State:
    		item (any):		The current item.
    		count (int):	The current count.

    	Input Ports:
    		input (any):	The items that need to be counted.
    		clear (any):	Clears the contents of the block.

    	Output Ports:
    		output (any):	The item that was counted.
    		count (int):	How many items have been counted, including the one
    						outputted on the `output` port.
    	"""
    	def __init__(self, name: str):
    		super(Counter, self).__init__(name)
    		self.count = self.addOutPort("count")
    		self.output = self.addOutPort("output")
    		self.clearContents()

    	def plot(self, axis, pos=(0.5, 0.5)):
    		"""Plots the counted info

    		Args:
    			axis:	The :mod:`matplotlib` axis to plot on.
    			pos:	The position to plot the data. Must be
    					in [0, 1] for both x and y to be visible.
    					Defaults to ``(0.5, 0.5)``.
    		"""
    		axis.clear()
    		axis.set_xlim(0, 1)
    		axis.set_ylim(0, 1)
    		axis.get_xaxis().set_visible(False)
    		axis.get_yaxis().set_visible(False)
    		axis.set_title(self.name + " Counted")
    		axis.annotate("$n = %i$" % self.state["count"], pos)

    	def clearContents(self):
    		self.state["item"] = None
    		self.state["count"] = 0

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

    	def extTransition(self, inputs):
    		if self.clear in inputs:
    			self.clearContents()
    		if self.input in inputs:
    			self.state["item"] = inputs[self.input]
    			self.state["count"] += 1
    		return self.state

    	def outputFnc(self):
    		if self.state["item"] is not None:
    			return {
    				self.output: self.state["item"],
    				self.count: self.state["count"]
    			}
    		return {}

    	def intTransition(self):
    		self.state["item"] = None
    		return self.state