# Source code for pydsol.model.source

    """ """
    """
    Created on: 22-7-2021 10:19

    @author: IvS
    """
    import numpy as np
    import itertools

    from pydsol.model.entities import Entity
    from pydsol.model.basic_logger import get_module_logger

    logger = get_module_logger(__name__)

    __all__ = ["Source"]

    class Source(object):
        """This class defines a basic source for a discrete event simulation model. A Source is the start-station
        of each entity, meaning the entities enter the system via the source. Entities are created at a source,
        given a specific interarrival time."""
        id_iter = itertools.count(1)

        def __init__(self, simulator, interarrival_time="default", num_entities=1, entity_type=Entity, **kwargs):
            """

            Parameters
            ----------
            simulator: Simulator object
                simulator of the model.
            interarrival_time:
                time between the creation of entities. Default is np.random.exponential(0.25).
            num_entities: int
                number of entities to create at once. Default is 1.
            kwargs:
                kwargs are the keyword arguments that are used to expand the source class.
                *name: str
                    user-specified source for the link
            """
            self.simulator = simulator
            self.num_entities = num_entities
            self.entity_type = entity_type

            self.next = None  # this should be the next process

            self.id = next(self.id_iter)
            self.name = "{0} {1}".format(self.__class__.__name__, str(self.id))
            if "name" in kwargs:
                self.name = kwargs["name"]

            self.interarrival_time = interarrival_time
            if "distribution" in kwargs:
                self.distribution = kwargs["distribution"]

            self.kwargs = kwargs

            self.simulator.schedule_event_now(self, "create_entities")

        def create_entities(self, **kwargs):
            """
            Create entities via SimEvent, given the interarrival time and the number of entities.

            Parameters
            ----------
            entity_type: class
                class where to make instances of, for example class Entity
            kwargs:
                kwargs are the keyword arguments that are used to invoke the method or expand the function.

            """

            # To make it work with the set seed of the simulator
            if self.interarrival_time == "default":
                self.interarrival_time = np.random.exponential(0.25)

            # Create new entity
            for _ in range(self.num_entities):
                entity = self.entity_type(self.simulator, self.simulator.simulator_time, **kwargs)
                logger.info("Time {0:.2f}: {1} is created at {2}".format(self.simulator.simulator_time, entity.name,
                                                                          self.name))
                self.exit_source(entity)

            if "distribution" in self.__dict__:
                try:
                    interarrival_time = self.distribution(*self.interarrival_time)
                except TypeError:
                    try:
                        interarrival_time = self.distribution(self.interarrival_time)
                    except TypeError:
                        raise "Wrong types for transfer in time ({1}) and/or distribution ({0}) for scheduling an event".format(
                            self.interarrival_time, self.distribution)
            else:
                interarrival_time = self.interarrival_time

            relative_delay = interarrival_time

            # Schedule event to create next entity according to the interarrival time
            self.simulator.schedule_event_rel(relative_delay, self, "create_entities")

        def exit_source(self, entity, **kwargs):
            """Schedules the event to exit the source and enter the output node.

            Parameters
            ----------
            entity: object
                the target on which a state change is scheduled.
            kwargs:
                kwargs are the keyword arguments that are used to expand the function.

            """
            self.simulator.schedule_event_now(self, "enter_output_node", entity=entity)

        def enter_output_node(self, entity, **kwargs):
            """Combine the entity with a Vehicle if an vehicle type is given. Combined or not, it schedules an event
            for exiting the output node.

            Parameters
            ----------
            entity: object
                the target on which a state change is scheduled.
            kwargs:
                kwargs are the keyword arguments that are used to expand the function.
                *vehicle_type: Vehicle class
                    subclass of Vehicle on which the entity should travel to the next destination.
                *vehicle_speed: int, optional
                    speed of the vehicle.

            """
            if "vehicle_type" in self.kwargs:
                if "vehicle_speed" in self.kwargs:
                    vehicle = self.kwargs["vehicle_type"](self.simulator, self.kwargs["vehicle_speed"])
                else:
                    vehicle = self.kwargs["vehicle_type"](self.simulator)
                vehicle.entities_on_vehicle.append(entity)
                logger.info("Time {0:.2f}: {1} loaded on {2}".format(self.simulator.simulator_time, entity.name,
                                                                     vehicle.name))
                entity = vehicle

            self.simulator.schedule_event_now(self, "exit_output_node", entity=entity)

        def exit_output_node(self, entity, **kwargs):
            """Exit the resource by selecting a link on which the entity should travel to the next destination by
            weighted choice.

            Parameters
            ----------
            entity: object
                the target on which a state change is scheduled.
            kwargs:
                kwargs are the keyword arguments that are used to expand the function.

            """
            try:
                # Selection based on weights in links
                next_list = self.next if isinstance(self.next, list) else [self.next]
                weights = [link.selection_weight for link in next_list]
                link_by_weight = np.random.choice(np.array(next_list), p=weights / np.sum(weights))
                link_by_weight.enter_input_node(entity)
            except AttributeError:
                try:
                    if len(next_list) > 1:
                        next_process = np.random.choice(np.array(next_list))
                        next_process.enter_input_node(entity)
                    if len(next_list) == 1:
                        self.next.enter_input_node(entity)
                except AttributeError:
                    raise AttributeError("{0} has no next process assigned".format(self.name))