<reponame>liuzuxin/metadrive<gh_stars>0
from metadrive.component.static_object.traffic_object import TrafficCone, TrafficWarning
from metadrive.component.vehicle.vehicle_type import LVehicle
from metadrive.constants import TerminationState, DEFAULT_AGENT
from metadrive.envs.safe_metadrive_env import SafeMetaDriveEnv
from metadrive.manager.object_manager import TrafficObjectManager


class ComplexObjectManager(TrafficObjectManager):
    def reset(self):
        ret = super(ComplexObjectManager, self).reset()
        lane = self.engine.current_map.road_network.graph[">>>"]["1C0_0_"][0]
        breakdown_vehicle = self.engine.object_manager.spawn_object(
            self.engine.traffic_manager.random_vehicle_type(),
            vehicle_config={
                "spawn_lane_index": lane.index,
                "spawn_longitude": 30
            }
        )
        self.engine.object_manager.accident_lanes.append(lane)
        lane_ = self.engine.current_map.road_network.graph[">>>"]["1C0_0_"][1]
        breakdown_vehicle = self.engine.object_manager.spawn_object(
            LVehicle, vehicle_config={
                "spawn_lane_index": lane_.index,
                "spawn_longitude": 30
            }
        )
        self.engine.object_manager.accident_lanes.append(lane_)

        alert = self.engine.object_manager.spawn_object(
            TrafficWarning, lane=lane, longitude=22, lateral=0, pbr_model=False
        )

        # part 1
        lane = self.engine.current_map.road_network.graph["1C0_1_"]["2S0_0_"][2]
        pos = [
            (-20, lane.width / 3), (-15.6, lane.width / 4), (-12.1, 0), (-8.7, -lane.width / 4),
            (-4.2, -lane.width / 2), (-0.7, -lane.width * 3 / 4), (4.1, -lane.width), (7.3, -lane.width),
            (11.5, -lane.width), (15.5, -lane.width), (20.0, -lane.width), (23.2, -lane.width), (29.1, -lane.width),
            (32.9, -lane.width / 2), (37.0, 0), (40.0, lane.width / 2)
        ]

        for p in pos:
            cone = self.engine.object_manager.spawn_object(
                TrafficCone, lane=lane, longitude=p[0], lateral=p[1] + lane.width / 2
            )
        self.engine.object_manager.accident_lanes.append(lane)
        from metadrive.component.vehicle.vehicle_type import SVehicle, XLVehicle
        v_pos = [8, 14]
        v_type = [SVehicle, XLVehicle]
        for v_long, v_t in zip(v_pos, v_type):
            v = self.engine.object_manager.spawn_object(
                v_t, vehicle_config={
                    "spawn_lane_index": lane.index,
                    "spawn_longitude": v_long
                }
            )

        # part 2
        lane = self.engine.current_map.road_network.graph["3R0_0_"]["3R0_1_"][0]
        self.engine.object_manager.accident_lanes.append(lane)
        pos = [
            (-20, lane.width / 3), (-15.6, lane.width / 4), (-12.1, 0), (-8.7, -lane.width / 4),
            (-4.2, -lane.width / 2), (-0.7, -lane.width * 3 / 4), (4.1, -lane.width), (7.3, -lane.width),
            (11.5, -lane.width), (15.5, -lane.width), (20.0, -lane.width), (23.2, -lane.width), (29.1, -lane.width),
            (32.9, -lane.width / 2), (37.0, 0), (40.0, lane.width / 2)
        ]

        for p in pos:
            p_ = (p[0] + 5, -p[1])
            cone = self.engine.object_manager.spawn_object(
                TrafficCone, lane=lane, longitude=p_[0], lateral=-p[1] - lane.width / 2
            )

        v_pos = [14, 19]
        for v_long in v_pos:
            v = self.engine.object_manager.spawn_object(
                self.engine.traffic_manager.random_vehicle_type(),
                vehicle_config={
                    "spawn_lane_index": lane.index,
                    "spawn_longitude": v_long
                }
            )

        alert = self.engine.object_manager.spawn_object(TrafficWarning, lane=lane, longitude=-35, lateral=0)

        alert = self.engine.object_manager.spawn_object(TrafficWarning, lane=lane, longitude=-60, lateral=0)

        # part 3
        lane = self.engine.current_map.road_network.graph["4C0_0_"]["4C0_1_"][2]
        self.engine.object_manager.accident_lanes.append(lane)
        pos = [
            (-12.1, 0), (-8.7, -lane.width / 4), (-4.2, -lane.width / 2), (-0.7, -lane.width * 3 / 4),
            (4.1, -lane.width), (7.3, -lane.width), (11.5, -lane.width), (15.5, -lane.width), (20.0, -lane.width),
            (23.2, -lane.width), (29.1, -lane.width), (32.9, -lane.width / 2), (37.0, 0), (40.0, lane.width / 2)
        ]

        for p in pos:
            p_ = (p[0] + 5, p[1] * 3.5 / 3)
            cone = self.engine.object_manager.spawn_object(
                TrafficCone, lane=lane, longitude=p_[0], lateral=p[1] + lane.width / 2
            )

        v_pos = [14, 19]
        for v_long in v_pos:
            v = self.engine.object_manager.spawn_object(
                self.engine.traffic_manager.random_vehicle_type(),
                vehicle_config={
                    "spawn_lane_index": lane.index,
                    "spawn_longitude": v_long
                }
            )

        # part 4
        lane = self.engine.current_map.road_network.graph["4C0_1_"]["5R0_0_"][0]
        self.engine.object_manager.accident_lanes.append(lane)
        # pos = [(-12, lane.width / 4), (-8.1, 0), (-4, -lane.width / 4), (-0.1, -lane.width / 2), (4, -lane.width)]
        #
        # for p in pos:
        #     p_ = (p[0] + 60, -p[1] * 3.5 / 3)
        #     cone = self.engine.object_manager.spawn_object(TrafficCone, lane=lane, longitude=p_[0], lateral=p_[1])

        return ret


class ComplexEnv(SafeMetaDriveEnv):
    """
    now for test use and demo use only
    """
    def default_config(self):
        config = super(ComplexEnv, self).default_config()
        config.update(
            {
                "environment_num": 1,
                "traffic_density": 0.05,
                "start_seed": 5,
                "accident_prob": 0.0,
                # "traffic_mode":"respawn",
                "debug_physics_world": False,
                "debug": False,
                "map": "CSRCR"
            }
        )
        return config

    def setup_engine(self):
        super(ComplexEnv, self).setup_engine()
        self.engine.update_manager("object_manager", ComplexObjectManager())


def test_object_collision_detection(render=False):
    env = ComplexEnv(
        {
            "manual_control": True,
            "traffic_density": 0.0,
            "use_render": render,
            "debug": False,
            "vehicle_config": {
                "show_lidar": True
            }
        }
    )
    try:
        o = env.reset()
        lane_index = (">>", ">>>", 0)
        alert = env.engine.object_manager.spawn_object(
            TrafficWarning, lane=env.current_map.road_network.get_lane(lane_index), longitude=22, lateral=0
        )
        lane_index = (">>", ">>>", 2)
        alert = env.engine.object_manager.spawn_object(
            TrafficCone, lane=env.current_map.road_network.get_lane(lane_index), longitude=22, lateral=0
        )
        crash_obj = False
        detect_obj = False
        for i in range(1, 100000 if render else 2000):
            o, r, d, info = env.step([0, 1])
            for obj in env.observations[DEFAULT_AGENT].detected_objects:
                if isinstance(obj, TrafficCone):
                    detect_obj = True
            if render:
                env.render()
            if info[TerminationState.CRASH_OBJECT]:
                crash_obj = True
                break
        assert crash_obj and detect_obj, "Can not crash with object!"
    finally:
        env.close()


if __name__ == "__main__":
    test_object_collision_detection(render=True)
