# Source code for pypdevs.server

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

    """
    Server for DEVS simulation
    """
    from pypdevs.basesimulator import BaseSimulator
    from pypdevs.controller import Controller
    import pypdevs.middleware as middleware
    from pypdevs.threadpool import ThreadPool
    import threading

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

    class Server(object):
        """
        A server to host MPI, will delegate all of its calls to the active simulation kernel.
        """
        # Don't forward some of the internally provided functions, but simply raise an AttributeError
        noforward = frozenset(["__str__",
                               "__getstate__",
                               "__setstate__",
                               "__repr__"])

        def __init__(self, name, total_size):
            """
            Constructor

            :param name: the name of the server, used for addressing (in MPI terms, this is the rank)
            :param total_size: the total size of the network in which the model lives
            """
            self.name = name
            self.kernel = None
            self.size = total_size
            self.proxies = [MPIRedirect(i) for i in range(total_size)]
            from pypdevs.MPIRedirect import LocalRedirect
            self.proxies[name] = LocalRedirect(self)
            self.queued_messages = []
            self.queued_time = None
            if total_size > 1:
                self.threadpool = ThreadPool(2)
                self.bootMPI()

        def getProxy(self, rank):
            """
            Get a proxy to a specified rank.

            This rank is allowed to be the local server, in which case a local shortcut is created.

            :param rank: the rank to return a proxy to, should be an int
            :returns: proxy to the server, either of type MPIRedirect or LocalRedirect
            """
            return self.proxies[rank]

        def checkLoadCheckpoint(self, name, gvt):
            """
            Reconstruct the server from a checkpoint.

            :param name: name of the checkpoint
            :param gvt: the GVT to restore to
            :returns: bool -- whether or not the checkpoint was successfully loaded
            """
            rank = self.name
            #assert debug("Accessing file " + str("%s_%s_%s.pdc" % (name, gvt, rank)))
            try:
                infile = open("%s_%s_%s.pdc" % (name, gvt, rank), 'r')
                pickle.load(infile)
                return True
            except KeyboardInterrupt:
                # If the user interrupts, still reraise
                raise
            except Exception as e:
                # Something went wrong
                print("Error found: " + str(e))
                return False

        def loadCheckpoint(self, name, gvt):
            """
            Reconstruct the server from a checkpoint.

            :param name: name of the checkpoint
            :param gvt: the GVT to restore to
            """
            rank = self.name
            #assert debug("Accessing file " + str("%s_%s_%s.pdc" % (name, gvt, rank)))
            infile = open("%s_%s_%s.pdc" % (name, gvt, rank), 'r')
            self.kernel = pickle.load(infile)
            self.kernel.server = self
            from pypdevs.MPIRedirect import LocalRedirect
            self.proxies[self.name] = LocalRedirect(self)
            infile.close()
            #assert debug("Closing file")
            self.kernel.loadCheckpoint()

        def setPickledData(self, pickled_data):
            """
            Set the pickled representation of the model.

            For use on the controller itself, as this doesn't need to unpickle the model.

            :param pickled_data: the pickled model
            """
            self.kernel.pickled_model = pickled_data

        def prepare(self, scheduler):
            """
            Prepare the server to receive the complete model over MPI

            :param scheduler: the scheduler to use
            """
            data = middleware.COMM_WORLD.bcast(None, root=0)
            if data is not None:
                self.saveAndProcessModel(data, scheduler)
                middleware.COMM_WORLD.barrier()

        def saveAndProcessModel(self, pickled_model, scheduler):
            """
            Receive the model and set it on the server, but also saves it for further reinitialisation.

            :param pickled_model: pickled representation of the model
            :param scheduler: the scheduler to use
            """
            self.sendModel(pickle.loads(pickled_model), scheduler)
            self.kernel.pickled_model = pickled_model

        def getName(self):
            """
            Returns the name of the server

            Is practically useless, since the server is previously addressed using its name. This does have a use as a ping function though.
            """
            # Actually more of a ping function...
            return self.name

        # All calls to this server are likely to be forwarded to the currently
        #  active simulation kernel, so provide an easy forwarder
        def __getattr__(self, name):
            """
            Remote calls happen on the server object, though it is different from the simulation kernel itself. Therefore, forward the actual function call to the correct kernel.

            :param name: the name of the method to call
            :returns: requested attribute
            """
            # For accesses that are actually meant for the currently running kernel
            if name in Server.noforward:
                raise AttributeError()
            return getattr(self.kernel, name)

        def processMPI(self, data, comm, remote):
            """
            Process an incomming MPI message and reply to it if necessary

            :param data: the data that was received
            :param comm: the MPI COMM object
            :param remote: the location from where the message was received
            """
            # Receiving a new request
            resend_tag = data[0]
            function = data[1]
            args = data[2]
            kwargs = data[3]
            result = getattr(self, function)(*args, **kwargs)
            if resend_tag is not None:
                if result is None:
                    result = 0
                comm.send(result, dest=remote, tag=resend_tag)

        def listenMPI(self):
            """
            Listen for incomming MPI messages and process them as soon as they are received
            """
            comm = middleware.COMM_WORLD
            status = middleware.MPI.Status()
            while 1:
                #assert debug("[" + str(comm.Get_rank()) + "]Listening to remote " + str(middleware.MPI.ANY_SOURCE) + " -- " + str(middleware.MPI.ANY_TAG))
                # First check if a message is present, otherwise we would have to do busy polling
                data = comm.recv(source=middleware.MPI.ANY_SOURCE,
                                 tag=middleware.MPI.ANY_TAG, status=status)
                tag = status.Get_tag()
                #assert debug("Got data from " + str(status.Get_source()) + " (" + str(status.Get_tag()) + "): " + str(data))
                if tag == 0:
                    # Flush all waiters, as we will never receive an answer when we close the receiver...
                    self.finishWaitingPool()
                    break
                elif tag == 1:
                    # NOTE Go back to listening ASAP, so do the processing on another thread
                    if data[1] == "receive" or data[1] == "receiveAntiMessages":
                        self.threadpool.addTask(Server.processMPI,
                                                 self,
                                                 list(data),
                                                 comm,
                                                 status.Get_source())
                    else:
                        # Normal 'control' commands are immediately executed, as they would otherwise have the potential to deadlock the node
                        threading.Thread(target=Server.processMPI,
                                         args=[self,
                                               list(data),
                                               comm,
                                               status.Get_source()]
                                        ).start()
                else:
                    # Receiving an answer to a previous request
                    try:
                        event = MPIRedirect.waiting[tag]
                        MPIRedirect.waiting[tag] = data
                        event.set()
                    except KeyError:
                        # Probably processed elsewhere already, just skip
                        pass
                    except AttributeError:
                        # Key was already set elsewhere
                        pass

        def finishWaitingPool(self):
            """
            Stop the complete MPI request queue from blocking, used when stopping simulation is necessary while requests are still outstanding.
            """
            for i in MPIRedirect.waiting:
                try:
                    i.set()
                except AttributeError:
                    # It was not a lock...
                    pass
                except KeyError:
                    # It was deleted in the meantime
                    pass

        def bootMPI(self):
            """
            Boot the MPI receivers when necessary, on an other thread to prevent blocking
            """
            if self.size > 1:
                listener = threading.Thread(target=Server.listenMPI, args=[self])
                # Make sure that this is a daemon on the controller, as otherwise this thread will prevent the atexit from stopping
                # Though on every other node this should NOT be a daemon, as this is the only part still running

                if middleware.COMM_WORLD.Get_rank() == 0:
                    listener.daemon = True
                listener.start()

        def sendModel(self, data, scheduler):
            """
            Receive a complete model and set it.

            :param data: a tuple containing the model, the model_ids dictionary, scheduler name, and a flag for whether or not the model was flattened to allow pickling
            :param scheduler: the scheduler to use
            """
            model, model_ids, flattened = data
            if self.name == 0:
                self.kernel = Controller(self.name, model, self)
            else:
                self.kernel = BaseSimulator(self.name, model, self)
            self.kernel.sendModel(model, model_ids, scheduler, flattened)

        def finish(self):
            """
            Stop the currently running simulation
            """
            sim = self.kernel
            with sim.simlock:
                # Shut down all threads on the topmost simulator
                sim.finished = True
                sim.should_run.set()
                self.finishWaitingPool()

                # Wait until they are done
                sim.sim_finish.wait()

        def queueMessage(self, time, model_id, action):
            """
            Queue a delayed action from being sent, to make it possible to batch them.

            Will raise an exception if previous messages form a different time were not yet flushed!
            This flushing is not done automatically, as otherwise the data would be received at a further timestep
            which causes problems with the GVT algorithm.

            :param time: the time at which the action happens
            :param model_id: the model_id that executed the action
            :param action: the action to execute (as a string)
            """
            if self.queued_time is None:
                self.queued_time = time
            elif time != self.queued_time:
                raise DEVSException("Queued message at wrong time! Probably forgot a flush")
            self.queued_messages.append([model_id, action])

        def flushQueuedMessages(self):
            """
            Flush all queued messages to the controller. This will block until all of them are queued.
            It is required to flush all messages right after all of them happened and this should happen within the critical section!
            """
            if self.queued_time is not None:
                self.getProxy(0).massDelayedActions(self.queued_time,
                                                    self.queued_messages)
                self.queued_messages = []
                self.queued_time = None