#  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:: ../img/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:: ../img/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:: ../img/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:: ../img/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

