text: # Copyright (C) 2020. Huawei Technologies Co., Ltd. All rights reserved.
  #
  # Permission is hereby granted, free of charge, to any person obtaining a copy
  # of this software and associated documentation files (the "Software"), to deal
  # in the Software without restriction, including without limitation the rights
  # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  # copies of the Software, and to permit persons to whom the Software is
  # furnished to do so, subject to the following conditions:
  #
  # The above copyright notice and this permission notice shall be included in
  # all copies or substantial portions of the Software.
  #
  # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
  # THE SOFTWARE.
  import collections.abc as collections_abc
  import logging
  import random
  from dataclasses import dataclass, field
  from sys import maxsize
  from typing import Any, Dict, Optional, Sequence, Tuple, Union
  
  from shapely.geometry import GeometryCollection, MultiPolygon, Point, Polygon
  from shapely.ops import unary_union
  
  from smarts.core import gen_id
  from smarts.core.sumo_road_network import SumoRoadNetwork
  from smarts.core.utils.id import SocialAgentId
  
  
  class _SumoParams(collections_abc.Mapping):
      """For some Sumo params (e.x. LaneChangingModel) the arguments are in title case
      with a given prefix. Subclassing this class allows for an automatic way to map
      between PEP8-compatible naming and Sumo's.
      """
  
      def __init__(self, prefix, whitelist=[], **kwargs):
          def snake_to_title(word):
              return "".join(x.capitalize() or "_" for x in word.split("_"))
  
          # XXX: On rare occasions sumo doesn't respect their own conventions
          #      (e.x. junction model's impatience).
          self._params = {key: kwargs.pop(key) for key in whitelist if key in kwargs}
  
          for key, value in kwargs.items():
              self._params[f"{prefix}{snake_to_title(key)}"] = value
  
      def __iter__(self):
          return iter(self._params)
  
      def __getitem__(self, key):
          return self._params[key]
  
      def __len__(self):
          return len(self._params)
  
      def __hash__(self):
          return hash(frozenset(self._params.items()))
  
      def __eq__(self, other):
          return self.__class__ == other.__class__ and hash(self) == hash(other)
  
  
  class LaneChangingModel(_SumoParams):
      """Models how the actor acts with respect to lane changes."""
  
      def __init__(self, **kwargs):
          super().__init__("lc", **kwargs)
  
  
  class JunctionModel(_SumoParams):
      """Models how the actor acts with respect to waiting at junctions."""
  
      def __init__(self, **kwargs):
          super().__init__("jm", whitelist=["impatience"], **kwargs)
  
  
  @dataclass(frozen=True)
  class Distribution:
      """A gaussian distribution used for randomized parameters."""
  
      mean: float
      """The mean value of the gaussian distribution."""
      sigma: float
      """The sigma value of the gaussian distribution."""
  
      def sample(self):
          """The next sample from the distribution."""
          return random.gauss(self.mean, self.sigma)
  
  
  @dataclass
  class UniformDistribution:
      """A uniform distribution, return a random number N
      such that a <= N <= b for a <= b and b <= N <= a for b < a.
      """
  
      a: float
      b: float
  
      def __post_init__(self):
          if self.b < self.a:
              self.a, self.b = self.b, self.a
  
      def sample(self):
          return random.uniform(self.a, self.b)
  
  
  @dataclass
  class TruncatedDistribution:
      """A truncated normal distribution, by default, location=0, scale=1"""
  
      a: float
      b: float
      loc: float = 0
      scale: float = 1
  
      def __post_init__(self):
          assert self.a != self.b
          if self.b < self.a:
              self.a, self.b = self.b, self.a
  
      def sample(self):
          from scipy.stats import truncnorm
  
          return truncnorm.rvs(self.a, self.b, loc=self.loc, scale=self.scale)
  
  
  @dataclass(frozen=True)
  class Actor:
      """This is the base description/spec type for traffic actors."""
  
      pass
  
  
  @dataclass(frozen=True, unsafe_hash=True)
  class TrafficActor(Actor):
      """Used as a description/spec for traffic actors (e.x. Vehicles, Pedestrians,
      etc). The defaults provided are for a car, but the name is not set to make it
      explicit that you actually want a car.
      """
  
      name: str
      """The name of the traffic actor. It must be unique."""
      accel: float = 2.6
      """The acceleration value of the actor."""
      decel: float = 4.5
      """The deceleration value of the actor."""
      tau: float = 1.0
      """The minimum time headway"""
      sigma: float = 0.5
      """The driver imperfection"""
      depart_speed: Union[float, str] = "max"
      """The starting speed of the actor"""
      emergency_decel: float = 4.5
      """maximum deceleration ability of vehicle in case of emergency"""
      speed: Distribution = Distribution(mean=1.0, sigma=0.1)
      """The speed distribution of this actor in m/s."""
      imperfection: Distribution = Distribution(mean=0.5, sigma=0)
      """Imperfection within range [0..1]"""
      min_gap: Distribution = Distribution(mean=2.5, sigma=0)
      """Minimum gap in meters."""
      max_speed: float = 55.5
      """The vehicle's maximum velocity (in m/s), defaults 200 km/h for vehicles"""
      vehicle_type: str = "passenger"
      """The type of vehicle this actor uses. ("passenger", "bus", "coach", "truck", "trailer")"""
      lane_changing_model: LaneChangingModel = field(
          default_factory=LaneChangingModel, hash=False
      )
      junction_model: JunctionModel = field(default_factory=JunctionModel, hash=False)
  
      @property
      def id(self) -> str:
          """The identifier tag of the traffic actor."""
          return "actor-{}-{}".format(self.name, hash(self))
  
  
  @dataclass(frozen=True)
  class SocialAgentActor(Actor):
      """Used as a description/spec for zoo traffic actors. These actors use a
      pre-trained model to understand how to act in the environment.
      """
  
      name: str
      """The name of the social actor. Must be unique."""
  
      # A pre-registered zoo identifying tag you provide to help SMARTS identify the
      # prefab of a social agent.
      agent_locator: str
      """The locator reference to the zoo registration call. Expects a string in the format
      of ‘path.to.file:locator-name’ where the path to the registration call is in the form
      {PYTHONPATH}[n]/path/to/file.py
      """
      policy_kwargs: Dict[str, Any] = field(default_factory=dict)
      """Additional keyword arguments to be passed to the constructed class overriding the
      existing registered arguments.
      """
      initial_speed: float = None
      """Set the initial speed, defaults to 0."""
  
      # pasted straight from TrafficActor class
      vehicle_type: str = "passenger"
  
  
  @dataclass(frozen=True)
  class BoidAgentActor(SocialAgentActor):
      """Used as a description/spec for boid traffic actors. Boid actors control multiple
      vehicles.
      """
  
      id: str = field(default_factory=lambda: f"boid-{gen_id()}")
  
      # The max number of vehicles that this agent will control at a time. This value is
      # honored when using a bubble for boid dynamic assignment.
      capacity: "BubbleLimits" = None
      """The capacity of the boid agent to take over vehicles."""
  
  
  @dataclass(frozen=True)
  class Route:
      """A route is represented by begin and end edge IDs, with an optional list of
      itermediary edge IDs. When an intermediary is not specified the router will
      decide what it should be.
      """
  
      ## edge, lane index, offset
      begin: Tuple[str, int, Any]
      """The (edge, lane_index, offset) details of the start location for the route.
  
      edge:
          The starting edge by name.
      lane_index:
          The lane index from the rightmost lane.
      offset:
          The offset in metres into the lane. Also acceptable\\: "max", "random"
      """
      ## edge, lane index, offset
      end: Tuple[str, int, Any]
      """The (edge, lane_index, offset) details of the end location for the route.
  
      edge:
          The starting edge by name.
      lane_index:
          The lane index from the rightmost lane.
      offset:
          The offset in metres into the lane. Also acceptable\\: "max", "random"
      """
  
      # Edges we want to make sure this route includes
      via: Tuple[str, ...] = field(default_factory=tuple)
      """The ids of edges that must be included in the route between `begin` and `end`."""
  
      @property
      def id(self) -> str:
          return "route-{}-{}-{}-".format(
              "_".join(map(str, self.begin)),
              "_".join(map(str, self.end)),
              hash(self),
          )
  
      @property
      def edges(self):
          return (self.begin[0],) + self.via + (self.end[0],)
  
  
  @dataclass(frozen=True)
  class RandomRoute:
      """An alternative to types.Route which specifies to sstudio to generate a random
      route.
      """
  
      id: str = field(default_factory=lambda: f"random-route-{gen_id()}")
  
  
  @dataclass(frozen=True)
  class Flow:
      """A route with an actor type emitted at a given rate."""
  
      route: Route
      """The route for the actor to attempt to follow."""
      rate: float
      """Vehicles per hour."""
      begin: float = 0
      """Start time in seconds."""
      # XXX: Defaults to 1 hour of traffic. We may want to change this to be "continual
      #      traffic", effectively an infinite end.
      end: float = 1 * 60 * 60
      """End time in seconds."""
      actors: Dict[TrafficActor, float] = field(default_factory=dict)
      """An actor to weight mapping associated as\\: { actor: weight }
      actor:
          The traffic actors that are provided.
      weight:
          The chance of this actor appearing as a ratio over total weight.
      """
  
      @property
      def id(self) -> str:
          return "flow-{}-{}-".format(
              self.route.id, str(hash(frozenset(self.actors.items())))
          )
  
      def __hash__(self):
          # Custom hash since self.actors is not hashable, here we first convert to a
          # frozenset.
          return hash((self.route, self.rate, frozenset(self.actors.items())))
  
      def __eq__(self, other):
          return self.__class__ == other.__class__ and hash(self) == hash(other)
  
  
  @dataclass(frozen=True)
  class JunctionEdgeIDResolver:
      """A utility for resolving a junction connection edge"""
  
      start_edge_id: str
      start_lane_index: int
      end_edge_id: str
      end_lane_index: int
  
      def to_edge(self, sumo_road_network: SumoRoadNetwork):
          return sumo_road_network.get_edge_in_junction(
              self.start_edge_id,
              self.start_lane_index,
              self.end_edge_id,
              self.end_lane_index,
          )
  
  
  @dataclass
  class Via:
      """A point on an edge that an actor must pass through"""
  
      edge_id: Union[str, JunctionEdgeIDResolver]
      """The edge this via is on"""
      lane_index: int
      """The lane this via sits on"""
      lane_offset: int
      """The offset along the lane where this via sits"""
      required_speed: float
      """The speed that a vehicle should travel through this via"""
      hit_distance: float = -1
      """The distance at which this waypoint can be hit. Negative means half the lane radius."""
  
  
  @dataclass(frozen=True)
  class Traffic:
      """The descriptor for traffic."""
  
      flows: Sequence[Flow]
      """Flows are used to define a steady supply of vehicles."""
  
  
  @dataclass(frozen=True)
  class EntryTactic:
      pass
  
  
  @dataclass(frozen=True)
  class TrapEntryTactic(EntryTactic):
      """An entry tactic that hijacks a vehicle to start the mission."""
  
      wait_to_hijack_limit_s: float
      """The amount of seconds a hijack will wait to get a vehicle before just emitting"""
      zone: "MapZone" = None
      """The zone of the hijack area"""
      exclusion_prefixes: Tuple[str, ...] = tuple()
      """The prefixes of vehicles to avoid hijacking"""
      default_entry_speed: float = None
      """The speed that the vehicle starts at when defaulting to emitting"""
  
  
  @dataclass(frozen=True)
  class UTurn:
      trigger_radius: int = 100
      """This task will be triggered if any vehicles within this radius"""
      initial_speed: float = None
      """This is the initial speed for the vehicle before it starts u-turn maneuver"""
      target_lane_index: int = 0
  
      @property
      def name(self):
          return "uturn"
  
  
  @dataclass(frozen=True)
  class CutIn:
      trigger_radius: int = 30
      """This task will be triggered if any vehicles within this radius"""
  
      complete_on_edge_id: Union[str, JunctionEdgeIDResolver] = None
      """The edge this task will be completed on"""
  
      @property
      def name(self):
          return "cut_in"
  
  
  @dataclass(frozen=True)
  class Mission:
      """The descriptor for an actor's mission."""
  
      route: Route
      """The route for the actor to attempt to follow."""
  
      via: Tuple[Via, ...] = ()
      """Points on an edge that an actor must pass through"""
  
      start_time: float = 0.1
      """The earliest simulation time that this mission starts but may start later in couple with
      `entry_tactic`.
      """
  
      entry_tactic: EntryTactic = None
      """A specific tactic the mission should employ to start the mission."""
  
      task: Tuple[CutIn, UTurn] = None
      """A task for the actor to accomplish."""
  
  
  @dataclass(frozen=True)
  class EndlessMission:
      """The descriptor for an actor's mission that has no end."""
  
      begin: Tuple[str, int, float]
      """The (edge, lane_index, offset) details of the start location for the route.
  
      edge:
          The starting edge by name.
      lane_index:
          The lane index from the rightmost lane.
      offset:
          The offset in metres into the lane. Also acceptable\\: 'max', 'random'
      """
      via: Tuple[Via, ...] = ()
      """Points on an edge that an actor must pass through"""
      start_time: float = 0.1
      """The earliest simulation time that this mission starts"""
      entry_tactic: EntryTactic = None
      """A specific tactic the mission should employ to start the mission"""
  
  
  @dataclass(frozen=True)
  class LapMission:
      """The descriptor for an actor's mission that defines mission that repeats
      in a closed loop.
      """
  
      route: Route
      """The route for the actor to attempt to follow"""
      num_laps: int
      """The amount of times to repeat the mission"""
      via: Tuple[Via, ...] = ()
      """Points on an edge that an actor must pass through"""
      start_time: float = 0.1
      """The earliest simulation time that this mission starts"""
      entry_tactic: EntryTactic = None
      """A specific tactic the mission should employ to start the mission"""
  
  
  @dataclass(frozen=True)
  class GroupedLapMission:
      """The descriptor for a group of actor missions that repeat in a closed loop."""
  
      route: Route
      """The route for the actors to attempt to follow"""
      offset: int
      """The offset of the "starting line" for the group"""
      lanes: int
      """The number of lanes the group occupies"""
      actor_count: int
      """The number of actors to be part of the group"""
      num_laps: int
      """The amount of times to repeat the mission"""
      via: Tuple[Via, ...] = ()
      """Points on an edge that an actor must pass through"""
  
  
  @dataclass(frozen=True)
  class Zone:
      """The base for a descriptor that defines a capture area."""
  
      def to_geometry(self, road_network: SumoRoadNetwork) -> Polygon:
          """Generates the geometry from this zone."""
          raise NotImplementedError
  
  
  @dataclass(frozen=True)
  class MapZone(Zone):
      """A descriptor that defines a capture area."""
  
      start: Tuple[str, int, float]
      """The (edge, lane_index, offset) details of the starting location.
  
      edge:
          The starting edge by name.
      lane_index:
          The lane index from the rightmost lane.
      offset:
          The offset in metres into the lane. Also acceptable\\: 'max', 'random'
      """
      length: float
      """The length of the geometry along the center of the lane. Also acceptable\\: 'max'"""
      n_lanes: 2
      """The number of lanes from right to left that this zone covers."""
  
      def to_geometry(self, road_network: SumoRoadNetwork) -> Polygon:
          def resolve_offset(offset, geometry_length, lane_length):
              if offset == "base":
                  return 0
              # push off of end of lane
              elif offset == "max":
                  return lane_length - geometry_length
              elif offset == "random":
                  return random.uniform(0, lane_length - geometry_length)
              else:
                  return float(offset)
  
          def pick_remaining_shape_after_split(geometry_collection, expected_point, lane):
  
              # smarts_paavi changes => update to use geoms attr to access each geom seperately
  
              lane_shape = geometry_collection
              if not isinstance(lane_shape, GeometryCollection):
                  return lane_shape
  
              # For simplicty, we only deal w/ the == 1 or 2 case
              if len(lane_shape.geoms) not in {1, 2}:
                  return None
  
              if len(lane_shape.geoms) == 1:
                  return lane_shape.geoms[0]
  
              # We assume that there are only two splited shapes to choose from
              keep_index = 0
              if lane_shape.geoms[1].minimum_rotated_rectangle.contains(expected_point):
                  # 0 is the discard piece, keep the other
                  keep_index = 1
  
              lane_shape = lane_shape.geoms[keep_index]
  
              return lane_shape
  
          lane_shapes = []
          edge_id, lane_idx, offset = self.start
          edge = road_network.edge_by_id(edge_id)
          buffer_from_ends = 1e-6
          for lane_idx in range(lane_idx, lane_idx + self.n_lanes):
              lane = edge.getLanes()[lane_idx]
              lane_length = lane.getLength()
              geom_length = self.length
  
              if geom_length > lane_length:
                  logging.debug(
                      f"Geometry is too long={geom_length} with offset={offset} for "
                      f"lane={lane.getID()}, using length={lane_length} instead"
                  )
                  geom_length = lane_length
  
              assert geom_length > 0  # Geom length is negative
  
              lane_shape = SumoRoadNetwork._buffered_lane_or_edge(
                  lane, width=lane.getWidth() + 0.3
              )
  
              lane_offset = resolve_offset(offset, geom_length, lane_length)
              lane_offset += buffer_from_ends
              geom_length = max(geom_length - buffer_from_ends, buffer_from_ends)
              lane_length = max(lane_length - buffer_from_ends, buffer_from_ends)
  
              min_cut = min(lane_offset, lane_length)
              # Second cut takes into account shortening of geometry by `min_cut`.
              max_cut = min(min_cut + geom_length, lane_length)
  
              midpoint = Point(
                  road_network.world_coord_from_offset(
                      lane, lane_offset + geom_length * 0.5
                  )
              )
  
              lane_shape = road_network.split_lane_shape_at_offset(
                  lane_shape, lane, min_cut
              )
              lane_shape = pick_remaining_shape_after_split(lane_shape, midpoint, lane)
              if lane_shape is None:
                  continue
  
              lane_shape = road_network.split_lane_shape_at_offset(
                  lane_shape,
                  lane,
                  max_cut,
              )
              lane_shape = pick_remaining_shape_after_split(lane_shape, midpoint, lane)
              if lane_shape is None:
                  continue
  
              lane_shapes.append(lane_shape)
  
          geom = unary_union(MultiPolygon(lane_shapes))
          return geom
  
  
  @dataclass(frozen=True)
  class PositionalZone(Zone):
      """A descriptor that defines a capture area at a specific XY location."""
  
      # center point
      pos: Tuple[float, float]
      """A (x,y) position of the zone in the scenario."""
      size: Tuple[float, float]
      """The (length, width) dimensions of the zone."""
  
      def to_geometry(self, road_network: SumoRoadNetwork = None) -> Polygon:
          w, h = self.size
          p0 = (self.pos[0] - w / 2, self.pos[1] - h / 2)  # min
          p1 = (self.pos[0] + w / 2, self.pos[1] + h / 2)  # max
          return Polygon([p0, (p0[0], p1[1]), p1, (p1[0], p0[1])])
  
  
  @dataclass(frozen=True)
  class BubbleLimits:
      hijack_limit: int = maxsize
      """The maximum number of vehicles the bubble can hijack"""
      shadow_limit: int = maxsize
      """The maximum number of vehicles the bubble can shadow"""
  
      def __post_init__(self):
          if self.shadow_limit is None:
              raise ValueError("Shadow limit must be a non-negative real number")
          if self.hijack_limit is None or self.shadow_limit < self.hijack_limit:
              raise ValueError("Shadow limit must be >= hijack limit")
  
  
  @dataclass(frozen=True)
  class Bubble:
      """A descriptor that defines a capture bubble for social agents."""
  
      zone: Zone
      """The zone which to capture vehicles."""
      actor: SocialAgentActor
      """The actor specification that this bubble works for."""
      margin: float = 2  # Used for "airlocking"; must be > 0
      """The exterior buffer area for airlocking. Must be > 0."""
      # If limit != None it will only allow that specified number of vehicles to be
      # hijacked. N.B. when actor = BoidAgentActor the lesser of the actor capacity
      # and bubble limit will be used.
      limit: BubbleLimits = None
      """The maximum number of actors that could be captured."""
      exclusion_prefixes: Tuple[str, ...] = field(default_factory=tuple)
      """Used to exclude social actors from capture."""
      id: str = field(default_factory=lambda: f"bubble-{gen_id()}")
      follow_actor_id: str = None
      """Actor ID of agent we want to pin to. Doing so makes this a "travelling bubble"
      which means it moves to follow the `follow_actor_id`'s vehicle. Offset is from the
      vehicle's center position to the bubble's center position.
      """
      follow_offset: Tuple[float, float] = None
      """Maintained offset to place the travelling bubble relative to the follow
      vehicle if it were facing north.
      """
      keep_alive: bool = False
      """If enabled, the social agent actor will be spawned upon first vehicle airlock
      and be reused for every subsequent vehicle entering the bubble until the episode
      is over.
      """
  
      def __post_init__(self):
          if self.margin <= 0:
              raise ValueError("Airlocking margin must be greater than 0")
  
          if self.follow_actor_id is not None and self.follow_offset is None:
              raise ValueError(
                  "A follow offset must be set if this is a travelling bubble"
              )
  
          if self.keep_alive and not self.is_boid:
              # TODO: We may want to remove this restriction in the future
              raise ValueError(
                  "Only boids can have keep_alive enabled (for persistent boids)"
              )
  
      @staticmethod
      def to_actor_id(actor, mission_group):
          return SocialAgentId.new(actor.name, group=mission_group)
  
      @property
      def is_boid(self):
          return isinstance(self.actor, BoidAgentActor)
  
  
  @dataclass(frozen=True)
  class RoadSurfacePatch:
      """A descriptor that defines a patch of road surface with a different friction coefficent."""
  
      zone: Zone
      """The zone which to capture vehicles."""
      begin_time: int
      """The start time in seconds of when this surface is active."""
      end_time: int
      """The end time in seconds of when this surface is active."""
      friction_coefficient: float
      """The surface friction coefficient."""
  
  
  @dataclass(frozen=True)
  class _ActorAndMission:
      actor: Actor
      mission: Union[Mission, EndlessMission, LapMission]
  
  
  @dataclass(frozen=True)
  class Scenario:
      traffic: Optional[Dict[str, Traffic]] = None
      ego_missions: Optional[Sequence[Mission]] = None
      # e.g. { "turning_agents": ([actors], [missions]), ... }
      social_agent_missions: Optional[
          Dict[str, Tuple[Sequence[SocialAgentActor], Sequence[Mission]]]
      ] = None
      """Every dictionary item {group: (actors, missions)} gets run simultaneously. If
      actors > 1 and missions = 0 or actors = 1 and missions > 0 we cycle through them
      every episode. Otherwise actors must be the same length as missions.
      """
      bubbles: Optional[Sequence[Bubble]] = None
      friction_maps: Optional[Sequence[RoadSurfacePatch]] = None
      traffic_histories: Optional[Sequence[str]] = None
  