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

"""
The actual DEVS solvers containing the main DEVS implementation
"""

from collections import defaultdict

from pypdevs.DEVS import *

from pypdevs.util import *
from pypdevs.logger import *

from pypdevs.classicDEVSWrapper import ClassicDEVSWrapper

class Solver(object):
    """
    A unified DEVS solver, containing all necessary functions
    """
    def __init__(self, listeners = {}):
        """
        Constructor
        """
        self.activities = {}
        self.dsdevs_dict = {}
        self.listeners = listeners

    def atomicOutputGenerationEventTracing(self, aDEVS, time):
        """
        Wrapper for the AtomicDEVS output function, which will save event counts

        :param aDEVS: the AtomicDEVS model that generates the output
        :param time: the time at which the output must be generated
        :returns: dict -- the generated output
        """
        retval = Solver.atomicOutputGeneration(self, aDEVS, time)
        for port in retval:
            port.msg_count += len(retval[port])
        return retval

    def atomicOutputGeneration(self, aDEVS, time):
        """
        AtomicDEVS function to generate output, invokes the outputFnc function of the model.

        :param aDEVS: the AtomicDEVS model that generates the output
        :param time: the time at which the output must be generated
        :returns: dict -- the generated output
        """
        aDEVS.my_output = aDEVS.outputFnc()

        # Being here means that this model created output, so it triggered its internal transition
        # save this knowledge in the basesimulator for usage in the actual transition step
        self.transitioning[aDEVS] |= 1

        return aDEVS.my_output

    def massAtomicTransitions(self, trans, clock):
        """
        AtomicDEVS function to perform all necessary transitions, 
        does so on a collection of models for performance.

        :param trans: iterable containing all models and their requested transition
        :param clock: the time at which the transition must happen
        """
        t, age = clock
        partialmod = []
        for aDEVS in trans:
            ttype = trans[aDEVS]

            ###########
            ## Memoization and activity tracking code
            ##   Skipped in local simulation
            if not self.temporary_irreversible:
                # Memo part
                if self.memoization and len(aDEVS.memo) >= 2:
                    found = False
                    prev = aDEVS.memo.pop()
                    memo = aDEVS.memo[-1]
                    if memo.time_last == clock and prev.loadState() == aDEVS.state:
                        if ttype == 1:
                            found = True
                        elif aDEVS.my_input == memo.my_input:
                            # Inputs should be equal too
                            if ttype == 3:
                                found = True
                            elif aDEVS.elapsed == memo.elapsed and ttype == 2:
                                found = True
                    if found:
                        aDEVS.state = memo.loadState()
                        aDEVS.time_last = clock
                        aDEVS.time_next = memo.time_next
                        # Just add the copy
                        aDEVS.old_states.append(memo)
                        if self.do_some_tracing:
                            # Completely skip all these calls if no tracing, saves us a lot of function calls
                            if ttype == 1:
                                self.tracers.tracesInternal(aDEVS)
                            elif ttype == 2:
                                self.tracers.tracesExternal(aDEVS)
                            elif ttype == 3:
                                self.tracers.tracesConfluent(aDEVS)
                        aDEVS.my_input = {}
                        if self.relocation_pending:
                            # Quit ASAP by throwing an exception
                            raise QuickStopException()
                        continue
                    else:
                        aDEVS.memo = []
                activity_tracking_prevalue = aDEVS.preActivityCalculation()
            elif self.activity_tracking:
                activity_tracking_prevalue = aDEVS.preActivityCalculation()
            ###########

            # Make a copy of the message before it is passed to the user
            if self.msg_copy != 2:
                # Prevent a pass statement, which still consumes some time in CPython
                if self.msg_copy == 1:
                    # Using list comprehension inside of dictionary comprehension...
                    aDEVS.my_input = {key: 
                            [i.copy() for i in aDEVS.my_input[key]] 
                            for key in aDEVS.my_input}
                elif self.msg_copy == 0:
                    # Dictionary comprehension
                    aDEVS.my_input = {key: 
                            pickle.loads(pickle.dumps(aDEVS.my_input[key], 
                                                      pickle.HIGHEST_PROTOCOL)) 
                            for key in aDEVS.my_input}

            # NOTE ttype mappings:            (EI)
            #       1 -- Internal transition  (01)
            #       2 -- External transition  (10)
            #       3 -- Confluent transition (11)
            if ttype == 1:
                # Internal only
                aDEVS.elapsed = None
                aDEVS.state = aDEVS.intTransition()
            elif ttype == 2:
                # External only
                aDEVS.elapsed = t - aDEVS.time_last[0]
                aDEVS.state = aDEVS.extTransition(aDEVS.my_input)
            elif ttype == 3:
                # Confluent
                aDEVS.elapsed = 0.
                aDEVS.state = aDEVS.confTransition(aDEVS.my_input)
            else:
                raise DEVSException(
                    "Problem in transitioning dictionary: unknown element %s" 
                    % ttype)

            ta = aDEVS.timeAdvance()
            aDEVS.time_last = clock

            if ta < 0:
                raise DEVSException("Negative time advance in atomic model '" + \
                                    aDEVS.getModelFullName() + "' with value " + \
                                    str(ta) + " at time " + str(t))

            # Update the time, this is just done in the timeNext, as this will propagate to the basesimulator
            aDEVS.time_next = (t + ta, 1 if ta else (age + 1))

            # Save the state
            if not self.temporary_irreversible:
                partialmod.append(aDEVS)
                # But only if there are multiple kernels, since otherwise there would be no other kernel to invoke a revertion
                # This can save us lots of time for local simulation (however, all other code is written with parallellisation in mind...)
                activity = aDEVS.postActivityCalculation(activity_tracking_prevalue)
                aDEVS.old_states.append(self.state_saver(aDEVS.time_last,
                                                         aDEVS.time_next,
                                                         aDEVS.state,
                                                         activity,
                                                         aDEVS.my_input,
                                                         aDEVS.elapsed))
                if self.relocation_pending:
                    # Quit ASAP by throwing an exception
                    for m in partialmod:
                        # Roll back these models to before the transitions
                        m.time_next = m.old_states[-1].time_next
                        m.time_last = m.old_states[-1].time_last
                        m.state = m.old_states[-1].loadState()
                    self.model.scheduler.massReschedule(trans)
                    self.server.flushQueuedMessages()
                    raise QuickStopException()
            elif self.activity_tracking:
                activity = aDEVS.postActivityCalculation(activity_tracking_prevalue)
                self.total_activities[aDEVS.model_id] += activity

            if self.do_some_tracing:
                # Completely skip all these calls if no tracing, saves us a lot of function calls
                if ttype == 1:
                    self.tracers.tracesInternal(aDEVS)
                elif ttype == 2:
                    self.tracers.tracesExternal(aDEVS)
                elif ttype == 3:
                    self.tracers.tracesConfluent(aDEVS)
      
            # Clear the bag
            aDEVS.my_input = {}
        self.server.flushQueuedMessages()

    def atomicInit(self, aDEVS, time):
        """
        AtomicDEVS function to initialise the model

        :param aDEVS: the model to initialise
        """
        aDEVS.time_last = (time[0] - aDEVS.elapsed, 1)
        ta = aDEVS.timeAdvance()

        if ta < 0:
            raise DEVSException("Negative time advance in atomic model '" + \
                                aDEVS.getModelFullName() + "' with value " + \
                                str(ta) + " at initialisation")
        aDEVS.time_next = (aDEVS.time_last[0] + ta, 1)
        # Save the state
        if not self.irreversible:
            aDEVS.old_states.append(self.state_saver(aDEVS.time_last,
                                                     aDEVS.time_next,
                                                     aDEVS.state,
                                                     0.0,
                                                     {},
                                                     0.0))

        # All tracing features
        self.tracers.tracesInit(aDEVS, time)

    def coupledOutputGenerationClassic(self, time):
        """
        CoupledDEVS function to generate the output, calls the atomicDEVS models where necessary. Output is routed too.

        :param time: the time at which output should be generated
        :returns: the models that should be rescheduled
        """
        cDEVS = self.model
        imminent = cDEVS.scheduler.getImminent(time)
        if not imminent:
            # For real time simulation, when a model is interrupted
            return self.transitioning
        reschedule = set(imminent)
        for model in imminent:
            model.time_next = (model.time_next[0], model.time_next[1] + 1)
        # Return value are the models to reschedule
        # self.transitioning are the models that must transition
        if len(imminent) > 1:
            # Perform all selects
            imminent.sort(key=lambda i: i.getModelFullName())
            pending = imminent
            level = 1
            while len(pending) > 1:
                # Take the model each time, as we need to make sure that the selectHierarchy is valid everywhere
                model = pending[0]
                # Make a set first to remove duplicates
                colliding = list(set([m.select_hierarchy[level] for m in pending]))
                chosen = model.select_hierarchy[level-1].select(
                        sorted(colliding, key=lambda i:i.getModelFullName()))
                pending = [m for m in pending 
                             if m.select_hierarchy[level] == chosen]
                level += 1
            child = pending[0]
        else:
            child = imminent[0]
        # Recorrect the timeNext of the model that will transition
        child.time_next = (child.time_next[0], child.time_next[1] - 1)

        outbag = child.my_output = ClassicDEVSWrapper(child).outputFnc()
        self.transitioning[child] = 1

        for outport in outbag:
            for inport, z in outport.routing_outline:
                payload = outbag[outport]
                if z is not None:
                    payload = [z(pickle.loads(pickle.dumps(m))) for m in payload]
                aDEVS = inport.host_DEVS
                aDEVS.my_input[inport] = list(payload)
                self.transitioning[aDEVS] = 2
                reschedule.add(aDEVS)
        # We have now generated the transitioning variable, though we need some small magic to have it work for classic DEVS
        self.transitioning = {ClassicDEVSWrapper(m): self.transitioning[m] 
                              for m in self.transitioning}
        return reschedule
    
    def coupledOutputGeneration(self, time):
        """
        CoupledDEVS function to generate the output, calls the atomicDEVS models where necessary. Output is routed too.

        :param time: the time at which output should be generated
        :returns: the models that should be rescheduled
        """
        cDEVS = self.model
        remotes = {}
        for child in cDEVS.scheduler.getImminent(time):
            outbag = self.atomicOutputGeneration(child, time)
            for outport in outbag:
                payload = outbag[outport]
                if not hasattr(outport, "routing_outline"):
                    raise Exception(outport)
                for inport, z in outport.routing_outline:
                    aDEVS = inport.host_DEVS
                    if z is not None:
                        payload = [z(pickle.loads(pickle.dumps(m))) 
                                   for m in payload]
                    if aDEVS.model_id in self.model.local_model_ids:
                        # This setdefault call is responsible for our non-linear runtime in several situations...
                        aDEVS.my_input.setdefault(inport, []).extend(payload)
                        self.transitioning[aDEVS] |= 2
                    else:
                        remotes.setdefault(aDEVS.model_id, 
                                           {}).setdefault(inport.port_id, 
                                                          []).extend(payload)
        for destination in remotes:
            self.send(destination, time, remotes[destination])
        return self.transitioning

    def coupledInit(self):
        """
        CoupledDEVS function to initialise the model, calls all its _local_ children too.
        """
        cDEVS = self.model
        time_next = (float('inf'), 1)
        # This part isn't fast, but it doesn't matter, since it just inits everything, optimizing here doesn't
        # matter as it is only called once AND every element has to be initted.
        # Only local models should receive this initialisation from us
        for d in self.local:
            self.atomicInit(d, (0.0, 0))
            time_next = min(time_next, d.time_next)
        # NOTE do not immediately assign to the timeNext, as this is used in the GVT algorithm to see whether a node has finished
        cDEVS.time_next = time_next
        self.model.setScheduler(self.model.scheduler_type)
        self.server.flushQueuedMessages()

    def performDSDEVS(self, transitioning):
        """
        Perform Dynamic Structure detection of the model

        :param transitioning: iteratable to be checked for a dynamic structure transiton
        """
        #TODO setting the server is very dirty
        self.dc_altered = set()
        for m in transitioning:
            m.server = self
        iterlist = [aDEVS.parent for aDEVS in transitioning 
                                 if aDEVS.modelTransition(self.dsdevs_dict)]
        # Contains all models that are already checked, to prevent duplicate checking.
        # This was not necessary for atomic models, as they are guaranteed to only be called
        # once, as they have no children to induce a structural change on them
        checked = set()
        while iterlist:
            new_iterlist = []
            for cDEVS in iterlist:
                cDEVS.server = self
                if cDEVS is None:
                    # Problematic
                    #assert warning("Root DEVS returned True in the modelTransition method; ignoring")
                    continue
                if cDEVS in checked:
                    continue
                checked.add(cDEVS)
                if cDEVS.modelTransition(self.dsdevs_dict):
                    new_iterlist.append(cDEVS.parent)
            # Don't update the iterlist while we are iterating over it
            iterlist = new_iterlist
        if self.dc_altered:
            self.model.redoDirectConnection(self.dc_altered)
