""" Class Scenario

Creation date: 2018 11 05
Author(s): Erwin de Gelder

To do:
Move falls_into method to ScenarioCategory and rename it to "comprises".

Modifications:
2018 11 22: Make is possible to instantiate a Scenario from JSON code.
2018 12 06: Add functionality to return the derived Tags.
2018 12 06: Make it possible to return full JSON code (incl. attributes' JSON code).
2018 12 06: to_openscenario function added.
2018 12 07: fall_into method for checking if Scenario falls into ScenarioCategory.
2019 05 22: Make use of type_checking.py to shorten the initialization.
2019 10 13: Update of terminology.
2019 11 04: Add options to automatically assign unique ids to actor/activities.
2020 08 23: Remove the update_uid options because uid is automatically generated by ScenarioElement.
2020 08 24: Make Scenario a subclass of TimeInterval.
2020 08 24: Enable instantiation of Scenario from json without needing full json code.
2020 08 25: Add functionality to evaluate the state variables (+derivatives) at given times.
2020 10 05: Change way of creating object from JSON code.
2020 10 12: Remove Dynamic/StaticPhysicalThing and use PhysicalElement instead.
2020 10 15: Add function get_actor_by_name.
"""

from typing import Callable, List, Tuple, Union
import numpy as np
from .activity import Activity, _activity_from_json
from .actor import Actor, _actor_from_json
from .physical_element import PhysicalElement, _physical_element_from_json
from .scenario_category import derive_actor_tags, _check_acts, _print_tags, _get_acts
from .scenario_element import DMObjects, _attributes_from_json, _object_from_json
from .state_variable import StateVariable
from .time_interval import TimeInterval, _time_interval_props_from_json
from .type_checking import check_for_list, check_for_tuple


class Scenario(TimeInterval):
    """ Scenario - either a real-world scenario or a test case.

    A scenario is a quantitative description of the ego vehicle, its activities
    and/or goals, its dynamic environment (consisting of traffic environment and
    conditions) and its static environment. From the perspective of the ego
    vehicle, a scenario contains all relevant events.

    When instantiating the Scenario object, the name, start event, end event,
    unique id (uid), and tags are passed. To set the activities, actors, dynamic
    physical things, acts, and static physical things, use the corresponding
    methods, i.e., set_activities(), set_actors(),
    set_dynamic_physical_things(), set_acts(), and set_static_physical_things(),
    respectively.

    Attributes:
        physical_elements (List[PhysicalElement]): All the things that make up
            the static environment and part of the dynamic environment.
        actors (List[Actor]): Actors that are participating in this scenario.
            This list should always include the ego vehicle.
        activities (List[Activity]): Activities that are relevant for this
            scenario.
        acts (List[tuple[Actor, Activity]]): The acts describe which actors
            perform which activities.
        name (str): A name that serves as a short description of the scenario.
        uid (int): A unique ID.
        tags (List[Tag]): A list of tags that formally defines the scenario
            category. These tags determine whether a scenario category comprises
            this scenario or not.
        start (Event): The starting event.
        end (Event): The end event.
    """

    def __init__(self, **kwargs):
        # Assign the attributes
        TimeInterval.__init__(self, **kwargs)
        self.physical_elements = []        # Type: List[PhysicalElement]
        self.actors = []                   # Type: List[Actor]
        self.activities = []               # Type: List[Activity]
        self.acts = []                     # Type: List[Tuple(Actor, Activity)]

        # Set attributes if provided by kwargs.
        if "physical_elements" in kwargs:
            self.set_physical_elements(kwargs["physical_elements"])
        if "actors" in kwargs:
            self.set_actors(kwargs["actors"])
        if "activities" in kwargs:
            self.set_activities(kwargs["activities"])
        if "acts" in kwargs:
            self.set_acts(kwargs["acts"])

    def set_physical_elements(self, physical_elements: List[PhysicalElement]) -> None:
        """ Set the physical things.

        Check whether the physical elements are correctly defined.

        :param physical_elements: List of physical elements.
        """
        # Check whether the physical elements are correctly defined.
        check_for_list("physical_elements", physical_elements, PhysicalElement, can_be_none=False)

        # Assign physical elements to an attribute.
        self.physical_elements = physical_elements

    def set_activities(self, activities: List[Activity]) -> None:
        """ Set the activities.

        Check whether the activities are correctly defined. Activities should be
        a list with instantiations of Activity.

        :param activities: List of activities that are used for this Scenario.
        """
        # Check whether the activities are correctly defined.
        check_for_list("activities", activities, Activity, can_be_none=False)

        # Assign actitivies to an attribute.
        self.activities = activities  # Type: List[Activity]

    def set_actors(self, actors: List[Actor]) -> None:
        """ Set the actors.

        Check whether the actors are correctly defined. Actors should be a list
        with instantiations of Actor.

        :param actors: List of actors that participate in the Scenario.
        """
        # Check whether the actors are correctly defined.
        check_for_list("actors", actors, Actor, can_be_none=False)

        # Assign actors to an attribute.
        self.actors = actors  # Type: List[Actor]

    def set_acts(self, acts_scenario: List[Tuple[Actor, Activity]], verbose: bool = True) -> None:
        """ Set the acts

        Check whether the acts are correctly defined. Each act should be a tuple
        with an actor and an activity, i.e., (Actor, Activity). Acts is a list
        containing multiple tuples (Actor, Activity).

        :param acts_scenario: The acts describe which actors perform which
            activities and a certain time. The actors and activities that are
            used in acts should also be passed with the actors and activities
            arguments. If not, a warning will be shown and the corresponding
            actor/activity will be added to the list of actors/activities.
        :param verbose: Set to False if warning should be surpressed.
        """
        check_for_list("acts", acts_scenario, tuple)
        for act in acts_scenario:
            check_for_tuple("act", act, (Actor, Activity))

        # Set the acts.
        self.acts = acts_scenario

        # Check whether the actors/activities defined with the acts are already listed. If not,
        # the corresponding actor/activity will be added and a warning will be shown.
        _check_acts(self.acts, self.actors, self.activities, verbose=verbose)

    def derived_tags(self) -> dict:
        """ Return all tags, including the tags of the attributes.

        The Scenario has tags, but also its attributes can have tags. More
        specifically, each PhysicalElement, each Actor, and each Activity might
        have tags. A dictionary will be returned. Each item of the dictionary
        contains a list of tags corresponding to either the own object (i.e.,
        Scenario), an Actor, or an PhysicalElement.

        The tags that might be associated with the Activity are returned with
        the Actor if the corresponding Actor is performing that Activity
        according to the defined acts.

        :return: List of tags.
        """
        # Instantiate the dictionary.
        tags = {}

        # Provide the tags of the very own object (Scenario).
        if self.tags:
            tags["{:s}::Scenario".format(self.name)] = self.tags

        # Provide the tags for each Actor.
        tags = derive_actor_tags(self.actors, self.acts, tags=tags)

        # Provide the tags for each PhysicalElement.
        for physical_element in self.physical_elements:
            if physical_element.get_tags():
                tags["{:s}::PhysicalElement".format(physical_element.name)] = \
                    physical_element.get_tags()

        # Return the tags.
        return tags

    def print_tags(self) -> None:
        """ Print the derived tags. """
        print(_print_tags(self.derived_tags()))

    def get_state(self, actor: Actor, state: StateVariable, time: Union[float, List, np.ndarray]) \
            -> Union[None, float, np.ndarray]:
        """ Obtain the values of the state variable at the given time instants.

        :param actor: The actor of which the state variable is to be retrieved.
        :param state: The state variable that is to be retrieved.
        :param time: The time instance(s).
        :return: The value of the state variable at the given time instants.
        """
        return self._get_state(actor, state, time)

    def get_state_dot(self, actor: Actor, state: StateVariable,
                      time: Union[float, List, np.ndarray]) -> Union[None, float, np.ndarray]:
        """ Obtain the derivative of the values of the state variable at the given time instants.

        :param actor: The actor of which the state variable derivative is to be
            retrieved.
        :param state: The state variable that is to be retrieved.
        :param time: The time instance(s).
        :return: The value of the state variable at the given time instants.
        """
        return self._get_state(actor, state, time, derivative=True)

    def _get_state(self, actor: Actor, state: StateVariable, time: Union[float, List, np.ndarray],
                   derivative=False) -> Union[None, float, np.ndarray]:
        vec_time = self._time2vec(time)
        is_valid = False

        # Loop through the acts.
        for my_actor, my_activity in self.acts:
            # Only continue with right actor and activity.
            if my_actor == actor and my_activity.category.state == state:
                tstart = my_activity.get_tstart()
                tend = my_activity.get_tend()
                if tstart is None or tend is None:
                    continue
                # Check if the time span contains time instances that we want to evaluate.
                mask = np.logical_and(vec_time >= tstart, vec_time <= tend)
                if np.any(mask):
                    if not derivative:
                        tmp_values = my_activity.get_state(time=vec_time[mask])
                    else:
                        tmp_values = my_activity.get_state_dot(time=vec_time[mask])
                    if not is_valid:
                        if len(tmp_values.shape) == 1:
                            values = np.ones(len(vec_time)) * np.nan
                        else:
                            values = np.ones((len(vec_time), tmp_values.shape[0])) * np.nan
                        is_valid = True
                    values[mask] = tmp_values.T

        if not is_valid:
            return None
        if isinstance(time, (float, int)):
            return values[0]
        return values

    @staticmethod
    def _time2vec(time: Union[float, List, np.ndarray]) -> np.ndarray:
        if isinstance(time, (float, int)):
            vec_time = np.array([time])
        elif isinstance(time, List):
            vec_time = np.array(time)
        elif isinstance(time, np.ndarray):
            vec_time = time
        else:
            raise TypeError("<time> needs to be of type <float>, <List>, or <np.ndarray>.")
        return vec_time

    def get_actor_by_name(self, name: str) -> Union[Actor, None]:
        """ Get the actor with the provided name.

        If there is no actor with the given name, None is returned. Note that
        as soon as an actor is found with the given name, this actor is retuned.
        Therefore, this function might fail if there are multiple actors with
        the same name.

        :param name: The name of the actor that is to be returned.
        :return: The actor with the given name.
        """
        for actor in self.actors:
            if actor.name == name:
                return actor
        return None

    def to_json(self) -> dict:
        scenario = TimeInterval.to_json(self)
        scenario["physical_elements"] = [dict(name=thing.name, uid=thing.uid)
                                         for thing in self.physical_elements]
        scenario["actors"] = [dict(name=actor.name, uid=actor.uid) for actor in self.actors]
        scenario["activities"] = [dict(name=activity.name, uid=activity.uid)
                                  for activity in self.activities]
        scenario["acts"] = []
        for actor, activity in self.acts:
            scenario["acts"].append({"actor": actor.uid, "activity": activity.uid})
        scenario["derived_tags"] = self.derived_tags()
        for key, tags in scenario["derived_tags"].items():
            scenario["derived_tags"][key] = [tag.to_json() for tag in tags]
        return scenario

    def to_json_full(self) -> dict:
        scenario = self.to_json()
        scenario.update(TimeInterval.to_json_full(self))
        scenario["physical_elements"] = [element.to_json_full()
                                         for element in self.physical_elements]
        scenario["actors"] = [actor.to_json_full() for actor in self.actors]
        scenario["activities"] = [activity.to_json_full() for activity in self.activities]
        return scenario


def _scenario_props_from_json(json: dict, attribute_objects: DMObjects, **kwargs) -> dict:
    props = _time_interval_props_from_json(json, attribute_objects,
                                           start=None if "start" not in kwargs else kwargs["start"],
                                           end=None if "end" not in kwargs else kwargs["end"])
    props.update(_attributes_from_json(
        json, attribute_objects,
        dict(physical_elements=(_physical_element_from_json, "physical_element"),
             actors=(_actor_from_json, "actor"),
             activities=(_activity_from_json, "activity")),
        **kwargs))
    props["acts"] = _get_acts(json, props["actors"], props["activities"])
    return props


def _scenario_from_json(json: dict, attribute_objects: DMObjects, **kwargs) -> Scenario:
    return Scenario(**_scenario_props_from_json(json, attribute_objects, **kwargs))


def scenario_from_json(json: dict, attribute_objects: DMObjects = None, **kwargs) -> Scenario:
    """ Get Scenario object from JSON code

    It is assumed that all the attributes are fully defined. Alternatively,
    the attributes can be passed via optional arguments. The following optional
    arguments are allowed:
    - start: The start event.
    - end: The end event.
    - physical_elements: The physical elements that define the static
        environment.
    - actors: The physical elements that are dynamic.
    - activities: The activities that describe the evolution of the dynamic
        environment.

    :param json: JSON code of Scenario.
    :param attribute_objects: A structure for storing all objects (optional).
    :return: Scenario object.
    """
    return _object_from_json(json, _scenario_from_json, "scenario", attribute_objects, **kwargs)


def _create_scenario_attributes(json: dict, class_name: str, func_from_json: Callable) -> List:
    """ Create list of objects of the class `class_name`.

    :param json: The dictionary with the JSON code.
    :param class_name: The name of the class for which the objects are to be
        instantiated.
    :param func_from_json: Function to instantiate an object.
    :return: List of objects that are member of the class `class_name`.
    """
    # Create the actor categories (ActorCategory) and actors (Actor).
    categories = []
    categories_ids = []
    objects = []
    for json_object in json[class_name]:
        if json_object["category"]["id"] in categories_ids:
            # Category object is already defined.
            category = categories[categories_ids.index(json_object["category"]["id"])]
            objects.append(func_from_json(json_object, category=category))
        else:
            objects.append(func_from_json(json_object))
            categories.append(objects[-1].category)  # This category has just been created
            categories_ids.append(categories[-1].uid)
    return objects
