# 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 importlib.resources as pkg_resources
import logging
import os
from collections import defaultdict
from typing import List, Sequence

import gltf
import numpy
from direct.showbase.ShowBase import ShowBase
from envision import types as envision_types
from envision.client import Client as EnvisionClient
from panda3d.core import ClockObject, NodePath, Shader, loadPrcFileData

import warnings

with warnings.catch_warnings():
    # XXX: Benign warning, seems no other way to "properly" fix
    warnings.filterwarnings("ignore", "numpy.ufunc size changed")
    from sklearn.metrics.pairwise import euclidean_distances

from . import glsl, models
from .agent_manager import AgentManager
from .bubble_manager import BubbleManager
from .colors import SceneColors
from .controllers import ActionSpaceType, Controllers
from .masks import RenderMasks
from .motion_planner_provider import MotionPlannerProvider
from .provider import ProviderState
from .scenario import Scenario
from .sensors import Collision
from .sumo_road_network import SumoRoadNetwork
from .sumo_traffic_simulation import SumoTrafficSimulation
from .traffic_history_provider import TrafficHistoryProvider
from smarts.core.chassis import AckermannChassis, BoxChassis
from .trap_manager import TrapManager
from .utils import pybullet
from .utils.pybullet import bullet_client as bc
from .utils.visdom_client import VisdomClient
from .vehicle import VehicleState
from .vehicle_index import VehicleIndex
from .waypoints import Waypoints

# disable vsync otherwise we are limited to refresh-rate of screen
loadPrcFileData("", "sync-video false")
loadPrcFileData("", "model-path %s" % os.getcwd())
loadPrcFileData("", "audio-library-name null")
loadPrcFileData("", "gl-version 3 3")
loadPrcFileData("", "notify-level error")
loadPrcFileData("", "print-pipe-types false")

# https://www.panda3d.org/manual/?title=Multithreaded_Render_Pipeline
# loadPrcFileData('', 'threading-model Cull/Draw')

# have makeTextureBuffer create a visible window
# loadPrcFileData('', 'show-buffers true')


class SMARTSNotSetupError(Exception):
    pass


class SMARTS(ShowBase):
    def __init__(
        self,
        agent_interfaces,
        traffic_sim,
        envision: EnvisionClient = None,
        visdom: VisdomClient = None,
        timestep_sec=0.1,
        reset_agents_only=False,
        zoo_addrs=None,
    ):
        try:
            super().__init__(self, windowType="offscreen")
        except Exception as e:
            # Known reasons for this failing:
            raise Exception(
                "SMARTS: Error in initializing framework for opening graphical display and creating scene graph. "
                "A typical reason is display not found. Try running with different configurations of "
                "`export DISPLAY=` using `:0`, `:1`... . If this does not work please consult "
                "the documentation."
            ) from e

        gltf.patch_loader(self.loader)

        self._log = logging.getLogger(self.__class__.__name__)

        self._is_setup = False
        self._scenario: Scenario = None
        self._envision: EnvisionClient = envision
        self._visdom: VisdomClient = visdom
        self._timestep_sec = timestep_sec
        self._traffic_sim = traffic_sim
        self._motion_planner_provider = MotionPlannerProvider()
        self._traffic_history_provider = TrafficHistoryProvider()
        self._providers = [
            self._traffic_sim,
            self._motion_planner_provider,
            self._traffic_history_provider,
        ]

        # We buffer provider state between steps to compensate for TRACI's timestep delay
        self._last_provider_state = None
        self._reset_agents_only = reset_agents_only  # a.k.a "teleportation"
        self._imitation_learning_mode = False

        # Global clock always proceeds by a fixed dt on each tick
        self.taskMgr.clock.setMode(ClockObject.M_non_real_time)
        self.taskMgr.clock.setDt(timestep_sec)
        self._elapsed_sim_time = 0

        self.setBackgroundColor(0, 0, 0, 1)

        # Displayed framerate is misleading since we are not using a realtime clock
        self.setFrameRateMeter(False)

        # For macOS GUI. See our `BulletClient` docstring for details.
        # from .utils.bullet import BulletClient
        # self._bullet_client = BulletClient(pybullet.GUI)
        self._bullet_client = bc.BulletClient(pybullet.DIRECT)
        self._dynamic_action_spaces = {
            ActionSpaceType.Continuous,
            ActionSpaceType.Lane,
            ActionSpaceType.ActuatorDynamic,
            ActionSpaceType.LaneWithContinuousSpeed,
            ActionSpaceType.Trajectory,
            ActionSpaceType.MPC,
        }

        # Set up indices
        self._agent_manager = AgentManager(agent_interfaces, zoo_addrs)
        self._vehicle_index = VehicleIndex()

        # TODO: Should not be stored in SMARTS
        self._vehicle_collisions = defaultdict(list)  # list of `Collision` instances
        self._vehicle_states = []

        self._bubble_manager = None
        self._trap_manager: TrapManager = None

        # SceneGraph-related setup
        self._root_np = None
        self._vehicles_np = None
        self._road_network_np = None
        self._ground_bullet_id = None

    def step(self, agent_actions):
        if not self._is_setup:
            raise SMARTSNotSetupError("Must call reset() or setup() before stepping.")

        try:
            return self._step(agent_actions)
        except (KeyboardInterrupt, SystemExit):
            # ensure we clean-up if the user exits the simulation
            self._log.info("Simulation was interrupted by the user.")
            self.destroy()
            raise  # re-raise the KeyboardInterrupt
        except Exception as e:
            self._log.error(
                "Simulation crashed with exception. Attempting to cleanly shutdown."
            )
            self._log.exception(e)
            self.destroy()
            raise  # re-raise

    def _check_if_acting_on_active_agents(self, agent_actions):
        for agent_id in agent_actions.keys():
            if agent_id not in self._agent_manager.ego_agent_ids:
                self._log.warning(
                    f"Attempted to perform actions on non-existing agent, {agent_id} "
                )

    def _step(self, agent_actions):
        """Steps through the simulation while applying the given agent actions.
        Returns the observations, rewards, and done signals.
        """

        # Due to a limitation of our traffic simulator(SUMO) interface(TRACI), we can
        # only observe traffic state of the previous simulation step.
        #
        # To compensate for this, we:
        #
        # 1. Fetch social agent actions
        # 2. Step all providers and harmonize state
        # 3. Step bubble manager
        # 4. Calculate observation and reward
        # 5. Send observations to social agents
        # 6. Clear done agents
        # 7. Perform visualization
        #
        # In this way, observations and reward are computed with data that is
        # consistently with one step of latencey and the agent will observe consistent
        # data.
        dt = self.taskMgr.clock.get_dt()
        self._elapsed_sim_time = self.taskMgr.clock.get_frame_time()

        # 1. Fetch agent actions
        all_agent_actions = self._agent_manager.fetch_agent_actions(self, agent_actions)

        # 2. Step all providers and harmonize state
        provider_state = self._step_providers(all_agent_actions, dt)
        self._check_if_acting_on_active_agents(agent_actions)

        # 3. Step bubble manager and trap manager
        self._vehicle_index.sync()
        self._bubble_manager.step(self)
        self._trap_manager.step(self)

        # 4. Calculate observation and reward
        # We pre-compute vehicle_states here because we *think* the users will
        # want these during their observation/reward computations.
        # This is a hack to give us some short term perf wins. Longer term we
        # need to expose better support for batched computations
        self._vehicle_states = [v.state for v in self._vehicle_index.vehicles]

        # Agents
        self._agent_manager.step_sensors(self)

        # Panda3D
        # runs through the render pipeline
        # MUST perform this after step_sensors() above, and before observe() below,
        # so that all updates are ready before rendering happens per frame
        self.taskMgr.mgr.poll()

        observations, rewards, scores, dones = self._agent_manager.observe(self)

        response_for_ego = self._agent_manager.filter_response_for_ego(
            (observations, rewards, scores, dones)
        )

        # 5. Send observations to social agents
        self._agent_manager.send_observations_to_social_agents(observations)

        # 6. Clear done agents
        self._teardown_done_agents_and_vehicles(dones)

        # 7. Perform visualization
        self._try_emit_envision_state(provider_state, observations, scores)
        self._try_emit_visdom_obs(observations)

        observations, rewards, scores, dones = response_for_ego
        extras = dict(scores=scores)
        return observations, rewards, dones, extras

    def _teardown_done_agents_and_vehicles(self, dones):
        def done_vehicle_ids(dones):
            vehicle_ids = set()
            for agent_id, done in dones.items():
                if self._agent_manager.is_boid_agent(agent_id):
                    vehicle_ids.update(id_ for id_ in done if done[id_])
                elif done:
                    ids = self._vehicle_index.vehicle_ids_by_actor_id(agent_id)
                    # 0 if shadowing, 1 if active
                    assert len(ids) <= 1, f"{len(ids)} <= 1"
                    vehicle_ids.update(ids)

            return vehicle_ids

        def done_agent_ids(dones):
            agent_ids = set()
            for agent_id, done in dones.items():
                if self._agent_manager.is_boid_agent(agent_id):
                    if not self.agent_manager.is_boid_keep_alive_agent(
                        agent_id
                    ) and all(dones[agent_id].values()):
                        agent_ids.add(agent_id)
                elif done:
                    agent_ids.add(agent_id)

            return agent_ids

        # XXX: These can not be put inline because we do queries that must proceed
        #      the actual teardown.
        vehicles_to_teardown = done_vehicle_ids(dones)
        agents_to_teardown = done_agent_ids(dones)

        self._agent_manager.teardown_ego_agents(agents_to_teardown)
        self._agent_manager.teardown_social_agents(agents_to_teardown)
        self._teardown_vehicles(vehicles_to_teardown)

    def reset(self, scenario: Scenario):
        if scenario == self._scenario and self._reset_agents_only:
            vehicle_ids_to_teardown = []
            agent_ids = self._agent_manager.teardown_ego_agents()
            for agent_id in agent_ids:
                ids = self._vehicle_index.vehicle_ids_by_actor_id(agent_id)
                vehicle_ids_to_teardown.extend(ids)
            self._teardown_vehicles(set(vehicle_ids_to_teardown))
            self._trap_manager.init_traps(
                scenario.road_network, scenario.waypoints, scenario.missions
            )
            self._agent_manager.init_ego_agents(self)
            self._sync_panda3d()
        else:
            self.teardown()
            self.setup(scenario)

        # Tell history provide to ignore vehicles if we have assigned mission to them
        self._traffic_history_provider.set_replaced_ids(scenario.missions.keys())

        self.taskMgr.clock.reset()
        self._elapsed_sim_time = 0

        self._vehicle_states = [v.state for v in self._vehicle_index.vehicles]
        observations, _, _, _ = self._agent_manager.observe(self)
        observations_for_ego = self._agent_manager.reset_agents(observations)

        # Visualization
        self._try_emit_visdom_obs(observations)
        if len(self._agent_manager.ego_agent_ids):
            while len(observations_for_ego) < 1:
                observations_for_ego, _, _, _ = self.step({})

        self._reset_providers()

        return observations_for_ego

    def setup(self, scenario: Scenario):
        self._scenario = scenario

        self._root_np = NodePath("sim")
        self._root_np.reparentTo(self.render)

        with pkg_resources.path(
            glsl, "unlit_shader.vert"
        ) as vshader_path, pkg_resources.path(
            glsl, "unlit_shader.frag"
        ) as fshader_path:
            unlit_shader = Shader.load(
                Shader.SL_GLSL,
                vertex=str(vshader_path.absolute()),
                fragment=str(fshader_path.absolute()),
            )
        self._root_np.setShader(unlit_shader)
        self._setup_road_network()
        self._vehicles_np = self._root_np.attachNewNode("vehicles")

        self._bubble_manager = BubbleManager(scenario.bubbles, scenario.road_network)
        self._trap_manager = TrapManager(scenario)

        self._setup_bullet_client(self._bullet_client)
        provider_state = self._setup_providers(self._scenario)
        self._agent_manager.setup_agents(self)

        self._harmonize_providers(provider_state)
        self._last_provider_state = provider_state

        self._is_setup = True

    def add_provider(self, provider):
        self._providers.append(provider)

    def switch_ego_agent(self, agent_interface):
        self._agent_manager.switch_initial_agent(agent_interface)
        self._is_setup = False

    def _setup_road_network(self):
        glb_path = self.scenario.map_glb_filepath
        if self._road_network_np:
            self._log.debug(
                "road_network={} already exists. Removing and adding a new "
                "one from glb_path={}".format(self._road_network_np, glb_path)
            )
        model_np = self.loader.loadModel(glb_path, noCache=True)

        np = self._root_np.attachNewNode("road_network")
        model_np.reparent_to(np)
        np.hide(RenderMasks.OCCUPANCY_HIDE)
        np.setColor(SceneColors.Road.value)
        self._road_network_np = np

    def _setup_bullet_client(self, client: bc.BulletClient):
        client.resetSimulation()
        client.configureDebugVisualizer(pybullet.COV_ENABLE_GUI, 0)

        # PyBullet defaults the timestep to 240Hz. Several parameters are tuned with
        # this value in mind. For example the number of solver iterations and the error
        # reduction parameters (erp) for contact, friction and non-contact joints.
        # Attempting to get around this we set the number of substeps so that
        # timestep * substeps = 240Hz. Bullet (C++) does something to this effect as
        # well (https://git.io/Jvf0M), but PyBullet does not expose it.
        client.setPhysicsEngineParameter(
            fixedTimeStep=self._timestep_sec,
            numSubSteps=int(self._timestep_sec * 240),
            numSolverIterations=10,
            solverResidualThreshold=0.001,
            # warmStartingFactor=0.99
        )

        client.setGravity(0, 0, -9.8)

        plane_path = self._scenario.plane_filepath

        # 1e6 is the default value for plane length and width.
        plane_scale = (
            max(self._scenario.map_bounding_box[0], self._scenario.map_bounding_box[1])
            / 1e6
        )
        if not os.path.exists(plane_path):
            with pkg_resources.path(models, "plane.urdf") as path:
                plane_path = str(path.absolute())

        self._ground_bullet_id = client.loadURDF(
            plane_path,
            useFixedBase=True,
            basePosition=self._scenario.map_bounding_box[2],
            globalScaling=1.1 * plane_scale,
        )

    def teardown(self):
        self._agent_manager.teardown()
        self._vehicle_index.teardown()

        self._bullet_client.resetSimulation()
        self._traffic_sim.teardown()
        self._teardown_providers()

        if self._root_np is not None:
            self._root_np.clearLight()
            self._root_np.removeNode()

        if self._bubble_manager is not None:
            self._bubble_manager.teardown()
            self._bubble_manager = None

        if self._trap_manager is not None:
            self._trap_manager.teardown()
            self._trap_manager = None

        self._vehicles_np = None
        self._ground_bullet_id = None
        self._road_network_np = None
        self._is_setup = False

    def destroy(self):
        self.teardown()

        if self._envision:
            self._envision.teardown()

        if self._visdom:
            self._visdom.teardown()

        self._agent_manager.destroy()
        self._traffic_sim.destroy()
        self._bullet_client.disconnect()

        super().destroy()

    def _teardown_vehicles(self, vehicle_ids):
        self._vehicle_index.teardown_vehicles_by_vehicle_ids(vehicle_ids)
        self._clear_collisions(vehicle_ids)

    def attach_sensors_to_vehicles(self, agent_spec, vehicle_ids):
        self._agent_manager.attach_sensors_to_vehicles(
            self, agent_spec.interface, vehicle_ids
        )

    def observe_from(self, vehicle_ids):
        return self._agent_manager.observe_from(self, vehicle_ids)

    @property
    def road_stiffness(self):
        return self._bullet_client.getDynamicsInfo(self._ground_bullet_id, -1)[9]

    @property
    def dynamic_action_spaces(self):
        return self._dynamic_action_spaces

    @property
    def traffic_sim(self) -> SumoTrafficSimulation:
        return self._traffic_sim

    @property
    def waypoints(self) -> Waypoints:
        return self.scenario.waypoints

    @property
    def road_network(self) -> SumoRoadNetwork:
        return self.scenario.road_network

    @property
    def np(self):
        return self._root_np

    @property
    def vehicles_np(self):
        return self._vehicles_np

    @property
    def bc(self):
        return self._bullet_client

    @property
    def envision(self):
        return self._envision

    @property
    def elapsed_sim_time(self):
        return self._elapsed_sim_time

    def teardown_agents_without_vehicles(self, agent_ids: Sequence):
        """
        Teardown agents in the given list that have no vehicles registered as
        controlled-by or shadowed-by

        Params:
            agent_ids: Sequence of agent ids
        """
        agents_to_teardown = {
            agent_id
            for agent_id in agent_ids
            # Only clean-up when there are no controlled agents left (e.g. boids)
            if len(
                self._vehicle_index.vehicles_by_actor_id(
                    agent_id, include_shadowers=True
                )
            )
            == 0
        }

        agents_to_teardown = {
            id_
            for id_ in agents_to_teardown
            if not self.agent_manager.is_boid_keep_alive_agent(id_)
        }
        self.agent_manager.teardown_social_agents(filter_ids=agents_to_teardown)

    def _teardown_vehicles_and_agents(self, vehicle_ids):
        shadow_and_controlling_agents = set()
        for vehicle_id in vehicle_ids:
            agent_id = self._vehicle_index.actor_id_from_vehicle_id(vehicle_id)
            if agent_id:
                shadow_and_controlling_agents.add(agent_id)

            shadow_agent_id = self._vehicle_index.shadow_actor_id_from_vehicle_id(
                vehicle_id
            )
            if shadow_agent_id:
                shadow_and_controlling_agents.add(shadow_agent_id)

        self._vehicle_index.teardown_vehicles_by_vehicle_ids(vehicle_ids)
        self.teardown_agents_without_vehicles(shadow_and_controlling_agents)

    def _pybullet_provider_sync(self, provider_state: ProviderState):
        current_vehicle_ids = {v.vehicle_id for v in provider_state.vehicles}
        previous_sv_ids = self._vehicle_index.social_vehicle_ids()
        exited_vehicles = previous_sv_ids - current_vehicle_ids
        self._teardown_vehicles_and_agents(exited_vehicles)

        # Update our pybullet world given this provider state
        for vehicle in provider_state.vehicles:
            vehicle_id = vehicle.vehicle_id
            # either this is a pybullet agent vehicle, or it is a social vehicle
            if vehicle_id in self._vehicle_index.agent_vehicle_ids():
                # this is an agent vehicle
                agent_id = self._vehicle_index.actor_id_from_vehicle_id(vehicle_id)
                agent_interface = self._agent_manager.agent_interface_for_agent_id(
                    agent_id
                )
                agent_action_space = agent_interface.action_space
                if agent_action_space not in self._dynamic_action_spaces:
                    # This is not a pybullet agent, but it has an avatar in this world
                    # to make it's observations. Update the avatar to match the new
                    # state of this vehicle
                    pybullet_vehicle = self._vehicle_index.vehicle_by_id(vehicle_id)
                    assert isinstance(pybullet_vehicle.chassis, BoxChassis)
                    pybullet_vehicle.control(pose=vehicle.pose, speed=vehicle.speed)
            else:
                # This vehicle is a social vehicle
                if vehicle_id in self._vehicle_index.social_vehicle_ids():
                    social_vehicle = self._vehicle_index.vehicle_by_id(vehicle_id)
                else:
                    # It is a new social vehicle we have not seen yet.
                    # Create it's avatar.
                    social_vehicle = self._vehicle_index.build_social_vehicle(
                        sim=self,
                        vehicle_state=vehicle,
                        actor_id=vehicle_id,
                        vehicle_id=vehicle_id,
                        vehicle_type=vehicle.vehicle_type,
                    )
                # Update the social vehicle avatar to match the vehicle state
                social_vehicle.control(pose=vehicle.pose, speed=vehicle.speed)

    def _pybullet_provider_step(self, agent_actions) -> ProviderState:
        self._perform_agent_actions(agent_actions)

        self._bullet_client.stepSimulation()

        self._process_collisions()

        provider_state = ProviderState()
        pybullet_agent_ids = {
            agent_id
            for agent_id, interface in self._agent_manager.agent_interfaces.items()
            if interface.action_space in self._dynamic_action_spaces
        }

        for vehicle_id in self._vehicle_index.agent_vehicle_ids():
            agent_id = self._vehicle_index.actor_id_from_vehicle_id(vehicle_id)
            if agent_id not in pybullet_agent_ids:
                continue

            vehicle = self._vehicle_index.vehicle_by_id(vehicle_id)
            vehicle.step(self._elapsed_sim_time)
            provider_state.vehicles.append(
                VehicleState(
                    vehicle_id=vehicle.id,
                    vehicle_type="passenger",
                    pose=vehicle.pose,
                    dimensions=vehicle.chassis.dimensions,
                    speed=vehicle.speed,
                    source="PYBULLET",
                )
            )

        return provider_state

    @property
    def vehicle_index(self):
        return self._vehicle_index

    @property
    def agent_manager(self):
        return self._agent_manager

    @property
    def providers(self):
        # TODO: Add check to ensure that action spaces are disjoint between providers
        # TODO: It's inconsistent that pybullet is not here
        return self._providers

    def _setup_providers(self, scenario) -> ProviderState:
        provider_state = ProviderState()
        for provider in self.providers:
            provider_state.merge(provider.setup(scenario))

        return provider_state

    def _teardown_providers(self):
        for provider in self.providers:
            provider.teardown()

        self._last_provider_state = None

    def _harmonize_providers(self, provider_state: ProviderState):
        for provider in self.providers:
            provider.sync(provider_state)

        self._pybullet_provider_sync(provider_state)
        self._sync_panda3d()

    def _reset_providers(self):
        for provider in self.providers:
            provider.reset()

    def _step_providers(self, actions, dt) -> List[VehicleState]:
        accumulated_provider_state = ProviderState()

        def agent_controls_vehicles(agent_id):
            vehicles = self._vehicle_index.vehicles_by_actor_id(agent_id)
            return len(vehicles) > 0

        def matches_provider_action_spaces(agent_id, action_spaces):
            interface = self._agent_manager.agent_interface_for_agent_id(agent_id)
            return interface.action_space in action_spaces

        # PyBullet
        pybullet_actions = {
            agent_id: action
            for agent_id, action in actions.items()
            if agent_controls_vehicles(agent_id)
            and matches_provider_action_spaces(agent_id, self._dynamic_action_spaces)
        }
        accumulated_provider_state.merge(self._pybullet_provider_step(pybullet_actions))

        for provider in self.providers:
            provider_state = self._step_provider(
                provider, actions, dt, self._elapsed_sim_time
            )
            if provider == self._traffic_sim:
                # Remove agent vehicles from provider vehicles
                provider_state.filter(self._vehicle_index.agent_vehicle_ids())

            accumulated_provider_state.merge(provider_state)

        self._harmonize_providers(accumulated_provider_state)
        return accumulated_provider_state

    def _step_provider(self, provider, actions, dt, elapsed_sim_time):
        def agent_controls_vehicles(agent_id):
            vehicles = self._vehicle_index.vehicles_by_actor_id(agent_id)
            return len(vehicles) > 0

        provider_actions = {}
        for agent_id, action in actions.items():
            agent_interface = self._agent_manager.agent_interface_for_agent_id(agent_id)
            if (
                agent_interface
                and agent_controls_vehicles(agent_id)
                and agent_interface.action_space in provider.action_spaces
            ):
                vehicle_ids = [
                    v.id
                    for v in self._vehicle_index.vehicles_by_actor_id(
                        agent_id, include_shadowers=True
                    )
                ]

                if self._agent_manager.is_boid_agent(agent_id):
                    for vehicle_id, vehicle_action in action.items():
                        assert vehicle_id in vehicle_ids
                        provider_actions[vehicle_id] = vehicle_action
                else:
                    assert len(vehicle_ids) == 1
                    provider_actions[vehicle_ids[0]] = action

        provider_state = provider.step(provider_actions, dt, elapsed_sim_time)
        return provider_state

    @property
    def scenario(self):
        return self._scenario

    @property
    def traffic_sim(self):
        return self._traffic_sim

    @property
    def np(self):
        return self._root_np

    @property
    def timestep_sec(self):
        return self._timestep_sec

    @property
    def road_stiffness(self):
        return self._bullet_client.getDynamicsInfo(self._ground_bullet_id, -1)[9]

    def neighborhood_vehicles_around_vehicle(self, vehicle, radius=None):
        other_states = [v for v in self._vehicle_states if v.vehicle_id != vehicle.id]
        if radius is None:
            return other_states

        other_positions = [state.pose.position for state in other_states]
        if not other_positions:
            return []

        distances = euclidean_distances(other_positions, [vehicle.position]).reshape(
            -1,
        )
        indices = numpy.argwhere(distances <= radius).flatten()
        return [other_states[i] for i in indices]

    def vehicle_did_collide(self, vehicle_id):
        return (
            len(
                [
                    c
                    for c in self._vehicle_collisions[vehicle_id]
                    if c.collidee_id != self._ground_bullet_id
                ]
            )
            > 0
        )

    def vehicle_collisions(self, vehicle_id):
        return [
            c
            for c in self._vehicle_collisions[vehicle_id]
            if c.collidee_id != self._ground_bullet_id
        ]

    def _clear_collisions(self, vehicle_ids):
        for vehicle_id in vehicle_ids:
            self._vehicle_collisions.pop(vehicle_id, None)

    def _perform_agent_actions(self, agent_actions):
        for agent_id, action in agent_actions.items():
            agent_vehicles = self._vehicle_index.vehicles_by_actor_id(agent_id)
            if len(agent_vehicles) == 0:
                self._log.warning(
                    f"{agent_id} doesn't have a vehicle, is the agent done? (dropping action)"
                )
            else:
                agent_interface = self._agent_manager.agent_interface_for_agent_id(
                    agent_id
                )
                is_boid_agent = self._agent_manager.is_boid_agent(agent_id)

                for vehicle in agent_vehicles:
                    vehicle_action = action[vehicle.id] if is_boid_agent else action

                    controller_state = self._vehicle_index.controller_state_for_vehicle_id(
                        vehicle.id
                    )
                    sensor_state = self._vehicle_index.sensor_state_for_vehicle_id(
                        vehicle.id
                    )
                    # TODO: Support performing batched actions
                    Controllers.perform_action(
                        self,
                        agent_id,
                        vehicle,
                        vehicle_action,
                        controller_state,
                        sensor_state,
                        agent_interface.action_space,
                        agent_interface.vehicle_type,
                    )

    def _sync_panda3d(self):
        for vehicle in self._vehicle_index.vehicles:
            vehicle.sync_to_panda3d()

    def _process_collisions(self):
        self._vehicle_collisions = defaultdict(list)  # list of `Collision` instances

        for vehicle_id in self._vehicle_index.agent_vehicle_ids():
            vehicle = self._vehicle_index.vehicle_by_id(vehicle_id)
            # We are only concerned with vehicle-vehicle collisions
            collidee_bullet_ids = set(
                [p.bullet_id for p in vehicle.chassis.contact_points]
            )
            collidee_bullet_ids.discard(self._ground_bullet_id)

            if not collidee_bullet_ids:
                continue

            for bullet_id in collidee_bullet_ids:
                collidee = self._bullet_id_to_vehicle(bullet_id)
                collision = self._node_to_collision(collidee.np.node())
                self._vehicle_collisions[vehicle_id].append(collision)

    def _bullet_id_to_vehicle(self, bullet_id):
        for vehicle in self._vehicle_index.vehicles:
            if bullet_id == vehicle.chassis.bullet_id:
                return vehicle

        assert False, "Only collisions with agent or social vehicles is supported"

    def _node_to_collision(self, node):
        for vehicle in self._vehicle_index.vehicles:
            if node == vehicle.np.node():
                actor_id = self._vehicle_index.actor_id_from_vehicle_id(vehicle.id)
                # TODO: Should we specify the collidee as the vehicle ID instead of
                #       the agent/social ID?
                return Collision(collidee_id=actor_id)

        assert False, "Only collisions with agent or social vehicles is supported"

    def _try_emit_envision_state(self, provider_state, obs, scores):
        if not self._envision:
            return

        traffic = {}
        for v in provider_state.vehicles:
            if v.vehicle_id in self._vehicle_index.agent_vehicle_ids():
                # this is an agent controlled vehicle
                agent_id = self._vehicle_index.actor_id_from_vehicle_id(v.vehicle_id)
                agent_obs = obs[agent_id]
                is_boid_agent = self._agent_manager.is_boid_agent(agent_id)
                vehicle_obs = agent_obs[v.vehicle_id] if is_boid_agent else agent_obs

                if self._agent_manager.is_ego(agent_id):
                    actor_type = envision_types.TrafficActorType.Agent
                    mission_route_geometry = self._vehicle_index.sensor_state_for_vehicle_id(
                        v.vehicle_id
                    ).mission_planner.route.geometry
                else:
                    actor_type = envision_types.TrafficActorType.SocialAgent
                    mission_route_geometry = None

                point_cloud = vehicle_obs.lidar_point_cloud or ([], [], [])
                point_cloud = point_cloud[0]  # (points, hits, rays), just want points

                # TODO: driven path should be read from vehicle_obs
                driven_path = self._vehicle_index.vehicle_by_id(
                    v.vehicle_id
                ).driven_path_sensor()

                road_waypoints = []
                if vehicle_obs.road_waypoints:
                    road_waypoints = [
                        path
                        for paths in vehicle_obs.road_waypoints.lanes.values()
                        for path in paths
                    ]
                traffic[v.vehicle_id] = envision_types.TrafficActorState(
                    name=self._agent_manager.agent_name(agent_id),
                    actor_type=actor_type,
                    vehicle_type=envision_types.VehicleType.Car,
                    position=v.pose.position,
                    heading=v.pose.heading,
                    speed=v.speed,
                    actor_id=envision_types.format_actor_id(
                        agent_id, v.vehicle_id, is_multi=is_boid_agent,
                    ),
                    events=vehicle_obs.events,
                    waypoint_paths=(vehicle_obs.waypoint_paths or []) + road_waypoints,
                    point_cloud=point_cloud,
                    driven_path=driven_path,
                    mission_route_geometry=mission_route_geometry,
                )
            elif v.vehicle_id in self._vehicle_index.social_vehicle_ids():
                # this is a social vehicle
                traffic[v.vehicle_id] = envision_types.TrafficActorState(
                    actor_type=envision_types.TrafficActorType.SocialVehicle,
                    vehicle_type=v.vehicle_type,
                    position=list(v.pose.position),
                    heading=v.pose.heading,
                    speed=v.speed,
                )

        bubble_geometry = [
            list(bubble.geometry.exterior.coords)
            for bubble in self._bubble_manager.bubbles
        ]

        state = envision_types.State(
            traffic=traffic,
            scenario_id=self.scenario.scenario_hash,
            bubbles=bubble_geometry,
            scene_colors=SceneColors.EnvisionColors.value,
            scores=scores,
        )
        self._envision.send(state)

    def _try_emit_visdom_obs(self, obs):
        if not self._visdom:
            return

        self._visdom.send(obs)
