#  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 a tracer that can can gather the execution cost of a simulation.

The execution cost is either measured at the input port(s) or the output port(s) of a
block. For readability, choose one and keep it consistent whenever possible.
"""

from pypdevsbbl.tracers import BaseTracer, superTrace


class FootprintTracer(BaseTracer):
    """Gathers cost information of the simulation.

    The provided cost can be whatever is required in the current simulation.
    It can be the carbon footprint, the pricetag or even the amount of
    electricity. For this reason, multiple :class:`FootprintTracer` may be
    registered to the simulation. Give each :class:`FootprintTracer` a name
    to correctly associate the traced cost with the real-life consequences.

    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.
        name (str):         The name of the :class:`FootprintTracer`. Used to
                            differentiate the results of different instances.
        mapping (dict):     A mapping from a port of a building block to a
                            number or a function that returns a number. This
                            number represents the cost for each item that is
                            inputted or outputted on that port. The function
                            takes the item as an argument. E.g. an object with
                            a weight of 60 kilograms might be more expensive
                            when compared to an object that weighs 10 kilograms.
        full (bool):        When ``True``, output the full trace. Otherwise,
                            only output the results.
        cpi (bool):         Whether or not to print the cost per item at the
                            end. These costs will be sorted with the largest
                            cost first. Defaults to ``False``.
    """
    def __init__(self, uid, server, filename, name, mapping, full=True, cpi=False):
        super(FootprintTracer, self).__init__(uid, server, filename)

        self.name = name
        self.mapping = mapping
        self.footprint = {}
        self.items = {}
        self.full = full
        self.cpi = cpi

        # Make sure all mappings are formatted in the same way
        for port in self.mapping:
            if not callable(self.mapping[port]):
                val = self.mapping[port]
                self.mapping[port] = lambda x, v=val: v

        self.inports = [x for x in self.mapping.keys() if x.is_input]
        self.inDEVS = [x.host_DEVS for x in self.inports]
        self.outports = [x for x in self.mapping.keys() if not x.is_input]
        self.outDEVS = [x.host_DEVS for x in self.outports]


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

        self.file.write("FOOTPRINT OF '%s':\n" % self.name)
        self.file.write("================" + ("=" * len(self.name)) + "\n")


    def stopTracer(self):
        if self.full:
            self.file.write("\n")
        self.file.write("\tTOTAL FOOTPRINT COST: %s\n" % str(self.total()))
        if self.cpi:
            self.file.write("\tCOST PER ITEM:\n")
            for it, val in sorted(self.items.items(), key=lambda x: x[1], reverse=True):
                self.file.write("\t\t%s: %f\n" % (str(it), self.items[it]))
        self.file.write("\n")

        super(FootprintTracer, self).stopTracer()


    def traceExternal(self, aDEVS):
        """Traces all items arriving on input ports."""
        dvs = aDEVS.model
        if dvs in self.inDEVS:
            for port in aDEVS.my_input:
                if port in self.inports:
                    for x in aDEVS.my_input[port]:
                        lst = self.footprint.setdefault(port, 0)
                        v = self.mapping[port](x)
                        self.footprint[port] += v
                        self.items.setdefault(x, 0)
                        self.items[x] += v
                        if self.full:
                            superTrace(self.server, self.uid, aDEVS,
                                       ["\t[INP]", port.getPortFullName(), "(", "%.5f" % x, ") :",
                                        lst, "to", "%.5f" % self.footprint[port], "\n"])


    def traceInternal(self, aDEVS):
        """Traces all items arriving on output ports."""
        dvs = aDEVS.model
        if dvs in self.outDEVS:
            for port in aDEVS.my_output:
                if port in self.outports:
                    for x in aDEVS.my_output[port]:
                        lst = self.footprint.setdefault(port, 0)
                        v = self.mapping[port](x)
                        self.footprint[port] += v
                        self.items.setdefault(x, 0)
                        self.items[x] += v
                        if self.full:
                            superTrace(self.server, self.uid, aDEVS,
                                       ["\t[OUT]", port.getPortFullName(), "(", "%.5f" % x, ") :",
                                        lst, "to", "%.5f" % self.footprint[port], "\n"])


    def total(self):
        """Obtain the total footprint of the simulation."""
        return sum(self.footprint.values())
