#  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 for statistical data."""

from pypdevsbbl.tracers import BaseTracer
import numpy as np

from pypdevs.DEVS import CoupledDEVS
from pypdevs.minimal import CoupledDEVS as minCoupledDEVS


class StatisticsTracer(BaseTracer):
    """Gathers information about the messages in a (sub)model.

    It can compute:
        - The average time of messages in the system.
        - The average amount of messages in the system for a certain
          time period.
        - The (amount of) messages that entered the system.
        - The (amount of) messages that left the system.
        - The (amount of) messages that were newly generated.
        - The (amount of) messages that were discarded or left behind.

    The internal depiction of a message is a tiplet:
    ``(message, arrival time, departure time)``.

    Note:
        The tracer makes use of the :meth:`id` function to determine
        which message entered and left. When your messages may change
        order (when using a different queueing discipline or so), make
        sure the messages can be uniquely identified afterwards.

        It would also be non-sensical to have a
        :class:`pypdevsbbl.generic.generators.ConstantGenerator` generate a
        single value and change the order in which they arrive. All
        messages are the same, so how would you determine they have
        a different order of arrival?

    Args:
        uid:                A unique identifier for distributed simulation.
                            Will be set automatically.
        server:             The server that runs the distributed simulation.
                            Will be set automatically.
        filename (str):     The name of the file to write the results to.
                            When ``None``, all results will be written to the
                            standard output.
        inports (list):     A list of the input ports to listen on for events
                            that enter the system. These may be ports of a
                            coupled DEVS.
        outports (list):    A list of the output ports to listen on for events
                            that leave the system. These may be ports of a
                            coupled DEVS.
    """
    def __init__(self, uid, server, filename, inports, outports):
        super(StatisticsTracer, self).__init__(uid, server, filename)
        for p in inports:
            if not p.is_input:
                raise TypeError("Inport must be an input port in %s.%s" % (p.host_DEVS.name, p.name))
        for p in outports:
            if p.is_input:
                raise TypeError("Outport must be an output port in %s.%s" % (p.host_DEVS.name, p.name))
        self.table = {}

        # Internally solve coupled DEVS
        # This is needed because of the flattening the DEVS simulator applies.
        self.inports = []
        self.inDEVS = []
        for port in inports:
            if port.type() != 'INPORT':
                raise TypeError("Can only listen on inports for messages that enter the submodel.")
            dvs = port.host_DEVS
            if issubclass(type(dvs), (CoupledDEVS, minCoupledDEVS)):
                if len(port.outline) > 0:  # Useless to listen on a port that's not connected
                    # The message will be inputted to all connected ports, so only listening on one port
                    # should be sufficient.
                    # All ports in the port.outline list are INPORTs by definition.
                    p = port.outline[0]
                    self.inports.append(p)
                    self.inDEVS.append(p.host_DEVS)
            else: # Default fuctionality
                self.inports.append(port)
                self.inDEVS.append(dvs)
        self.outports = []
        self.outDEVS = []
        for port in outports:
            if port.type() != 'OUTPORT':
                raise TypeError("Can only listen on outports for messages that leave the submodel.")
            dvs = port.host_DEVS
            if issubclass(type(dvs), (CoupledDEVS, minCoupledDEVS)):
                # All messages arriving on this port will be outputted, so we have a set of outputs
                # All ports in the port.inline list are OUTPORTs by definition.
                for p in port.inline:
                    self.outports.append(p)
                    if p.host_DEVS not in self.outDEVS:
                        self.outDEVS.append(p.host_DEVS)
            else: # Default functionality
                self.outports.append(port)
                self.inDEVS.append(dvs)

    def startTracer(self, recover):
        super(StatisticsTracer, self).startTracer(recover)

        self.file.write("STATISTICS:\n")
        self.file.write("===========\n\tRUNNING...")

    def stopTracer(self):
        self.file.write("\r\tSUMMARY:\t\n" + self.summary('\t\t') + "\n")

        super(StatisticsTracer, self).stopTracer()

    def traceExternal(self, aDEVS):
        """Trace messages that enter the system."""
        dvs = aDEVS.model
        if dvs in self.inDEVS:
            inps = aDEVS.my_input
            for port in self.inports:
                if port in inps:
                    for p in inps[port]:
                        self.table.setdefault(id(p), [p, None, None])[1] = aDEVS.time_last[0]

    def traceInternal(self, aDEVS):
        """Trace messages that leave the system."""
        dvs = aDEVS.model
        if dvs in self.outDEVS:
            outs = aDEVS.my_output
            for port in self.outports:
                if port in outs:
                    for p in outs[port]:
                        self.table.setdefault(id(p), [p, None, None])[2] = aDEVS.time_last[0]

    def summary(self, prefix='\t'):
        """Gets the current stats as a string."""
        summ = [
            "Average Time in System:                   %.5f" % self.average_time_in_system(),
            "Average Num of Messages per Time Unit:    %.5f" % self.average_num(),
            "Total Throughput:                         %i" % self.n_passed(),
            "Total Amount of Arrivals:                 %i" % self.n_entries(),
            "Total Amount of Departures:               %i" % self.n_exits(),
            "Amount of Newly Generated Messages:       %i" % self.n_new(),
            "Amount of Bequeathed Messages:            %i" % self.n_bequeathed()
        ]

        return prefix + ("\n" + prefix).join(summ)

    def average_time_in_system(self):
        """Reports the average time a message remained in the system.

        Note:
            This ignores all newly created and bequeathed messages.
        """
        msgs = self.passed()
        if len(msgs) == 0:
            return 0
        avg = 0
        for _, a, d in msgs:
            avg += d - a
        return float(avg) / len(msgs)

    def average_num(self, dx=1.0, start=0, end=None):
        """Reports the average amount of messages for a certain time period.

        Note:
            This ignores all newly created and bequeathed messages.

        Args:
            dx (numeric):       The time period to look over. Defaults to 1.
            start (numeric):    The point in time to start measuring from.
                                When ``None``, it will use the first seen
                                arrival as a starting point. Defaults to 0
                                (= usually the start of the simulation).
            end (numeric):      The point in time to end measuring at. When
                                ``None``, it will use the last seen departure
                                as an end point. This point is inclusive in
                                the range. Defaults to ``None``.
        """
        # Create dictionaries for all arrival and departures
        msgs = self.passed()
        if len(msgs) == 0:
            return 0
        arr = {}
        dep = {}
        for v, a, d in msgs:
            arr.setdefault(a, []).append((v, a, d))
            dep.setdefault(d, []).append((v, a, d))

        # Dicts are not sorted
        sarr = sorted(arr.keys())
        sdep = sorted(dep.keys())

        times = []
        prev = []
        last = min(arr.keys()) if start is None else start
        for t in np.arange(last + dx, (max(dep.keys()) if end is None else end) + 2 * dx, dx):
            curr = prev

            # Add all arriving messages in the time period
            for sa in (x for x in sarr if last <= x < t):
                curr += arr[sa]

            # Remove all departing messages in the time period
            for sd in (x for x in sdep if last <= x < t):
                for m in dep[sd]:
                    curr.remove(m)

            prev = curr[:]
            times.append(curr)
            last = t

        return sum(map(len, times)) / len(times)


    def entries(self):
        """Reports all messages that entered the system."""
        return [self.table[x] for x in self.table if self.table[x][1] is not None]

    def n_entries(self):
        """Reports the amount of messages that entered the system."""
        return len(self.entries())

    def exits(self):
        """Reports all messages that left the system."""
        return [self.table[x] for x in self.table if self.table[x][2] is not None]

    def n_exits(self):
        """Reports the amount of messages that left the system."""
        return len(self.exits())

    def passed(self):
        """Reports all messages that passed through the system."""
        return [self.table[x] for x in self.table if self.table[x][1] is not None and self.table[x][2] is not None]

    def n_passed(self):
        """Reports the amount of messages that passed through the system."""
        return len(self.passed())

    def new(self):
        """Reports all messages that were generated in the submodel."""
        return [self.table[x] for x in self.table if self.table[x][1] is None and self.table[x][2] is not None]

    def n_new(self):
        """Reports the amount of messages that were generated in the submodel."""
        return len(self.new())

    def bequeathed(self):
        """Reports all messages that were dropped or left behind in the submodel."""
        return [self.table[x] for x in self.table if self.table[x][1] is not None and self.table[x][2] is None]

    def n_bequeathed(self):
        """Reports the amount of messages that were dropped or left behind in the submodel."""
        return len(self.bequeathed())

    forgotten = bequeathed
    """Alias for :meth:`bequeathed`, because, although less accurate, it is easier to remember."""

    n_forgotten = n_bequeathed
    """Alias for :meth:`n_bequeathed`, because, although less accurate, it is easier to remember."""



