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

"""This file contains the tracer that can provide plots of the :mod:`pypdevsbbl.generic.collectors`."""

from pypdevsbbl.generic.collectors import BaseCollector
from pypdevsbbl.tracers import BaseTracer

import matplotlib.pyplot as plt
import inspect

# TODO: make use of matplotlib's animate module?
class PlotTracer(BaseTracer):
    """Plots data from the provided collectors.

    Note:
        Because of how :mod:`matplotlib` and :mod:`pypdevs` work, this
        tracer only seems to be working within a `Jupyter` notebook.

        See :file:`src/notebook/PlotTracer.ipynb` for an example on how
        to do this.

    Args:
        uid:                A unique identifier for distributed simulation.
                            Will be set automatically.
        server:             The server that runs the distributed simulation.
                            Will be set automatically.
        *collectors:        The collectors to plot. This is a pair in the
                            form of ``(BaseCollector, Axis)`` or a triple in
                            the form of ``(BaseDEVS, Axis, func)``, where
                            ``BaseCollector`` or ``BaseDEVS`` are instances
                            of those classes, ``Axis`` a :class:`matplotlib.Axes`
                            and ``func`` a callable that takes the first two
                            members and plots the data on the axis.

                            Additionally, the tuples/lists may be extended
                            with a thrird/fourth member, which is an additional
                            argument to pass to the plot function. What this is
                            for each collector is described in their docs.
    """
    def __init__(self, uid, server, *collectors):
        super(PlotTracer, self).__init__(uid, server, None)

        self.blocks = {}

        for ix in range(len(collectors)):
            cls = collectors[ix]
            if 2 <= len(cls) <= 4:
                col = cls[0]
                ax = cls[1]
                func = PlotTracer.defPlot
                arg = inspect.Parameter.empty
                if len(cls) == 3:
                    if callable(cls[2]):
                        func = cls[2]
                    else:
                        arg = cls[2]
                if len(cls) == 4:
                    func = cls[2]
                    if not callable(func):
                        raise ValueError("Third argument must be callable.")
                    arg = cls[3]

                self.blocks.setdefault(col, []).append((ax, func, arg))
            else:
                raise TypeError("Cannot unpack collector info. Invalid number of arguments.")

    @staticmethod
    def defPlot(collector, axis, extra=inspect.Parameter.empty):
        """Default plotting function of the collectors."""
        if not issubclass(type(collector), BaseCollector):
            raise ValueError("No collector to plot found. Please specify a plotting function.")
        sig = inspect.signature(collector.plot)
        dft = None
        for p in sig.parameters:
            if p != 'axis':
                dft = sig.parameters[p].default
                break
        if extra is not inspect.Parameter.empty:
            dft = extra
        collector.plot(axis, dft)


    def traceExternal(self, aDEVS):
        dvs = aDEVS.model
        if dvs in self.blocks:
            ll = self.blocks[dvs]
            for ax, func, arg in ll:
                if arg is inspect.Parameter.empty:
                    func(dvs, ax)
                else:
                    func(dvs, ax, arg)
            plt.show()
