#!/usr/bin/env python

# Copyright (c) 2021 Computer Vision Center (CVC) at the Universitat Autonoma de
# Barcelona (UAB).
#
# This work is licensed under the terms of the MIT license.
# For a copy, see <https://opensource.org/licenses/MIT>.

"""Example script to generate traffic in the simulation"""

"""Adopted as a class by Tharindu"""

import glob
import os
import sys
import time

try:
    sys.path.append(glob.glob('../carla/dist/carla-*%d.%d-%s.egg' % (
        sys.version_info.major,
        sys.version_info.minor,
        'win-amd64' if os.name == 'nt' else 'linux-x86_64'))[0])
except IndexError:
    pass

import carla

import argparse
import logging
from numpy import random


class GenerateTraffic(object):

    def __init__(self, world=None, client=None, num_vehicle=10, num_pedestrians=10, safe=True, filterv="vehicle.*",
                 generationv="All", filterw="walker.pedestrian.*", generationw="All", tm_port=8000, asynch=False,
                 hybrid=False, seed=None, seedw=0, car_lights_on=True, hero=False, respawn=False, no_rendering=False):

        self.world = world
        self.client = client
        self.number_of_vehicles = num_vehicle
        self.number_of_walkers = num_pedestrians
        self.safe = safe  # Avoid spawning vehicles prone to accidents
        self.filterv = filterv  # Filter vehicle model (default: "vehicle.*")
        self.generationv = generationv  # restrict to certain vehicle generation (values: "1","2","All" - default:
        # "All")
        self.filterw = filterw  # Filter pedestrian type (default: "walker.pedestrian.*")
        self.generationw = generationw  # restrict to certain pedestrian generation (values: "1","2","All" - default:
        # "2")
        self.tm_port = tm_port  # Port to communicate with TM (default: 8000)
        self.asynch = asynch  # Activate asynchronous mode execution
        self.hybrid = hybrid  # Activate hybrid mode for Traffic Manager
        if seed is not None:  # Set random device seed and deterministic mode for Traffic Manager
            self.seed = seed
        else:
            self.seed = int(time.time())
        self.seedw = seedw  # Set the seed for pedestrians module
        self.car_lights_on = car_lights_on  # Enable automatic car light management
        self.hero = hero  # Set one of the vehicles as hero')
        self.respawn = respawn  # Automatically respawn dormant vehicles (only in large maps)')
        self.no_rendering = no_rendering  # Activate no rendering mode

        self.traffic_manager = None
        self.vehicles_list = []
        self.walkers_list = []
        self.all_id = []
        self.synchronous_master = False

        self.is_exit = False

    def get_actor_blueprints(self, filter, generation):
        bps = self.world.get_blueprint_library().filter(filter)

        if generation.lower() == "all":
            return bps

        # If the filter returns only one bp, we assume that this one needed
        # and therefore, we ignore the generation
        if len(bps) == 1:
            return bps

        try:
            int_generation = int(generation)
            # Check if generation is in available generations
            if int_generation in [1, 2]:
                bps = [x for x in bps if int(x.get_attribute('generation')) == int_generation]
                return bps
            else:
                print("   Warning! Actor Generation is not valid. No actor will be spawned.")
                return []
        except:
            print("   Warning! Actor Generation is not valid. No actor will be spawned.")
            return []

    def tick_world(self):
        self.world.tick()

    def spawn_vehicles_and_pedestrians(self):

        logging.basicConfig(format='%(levelname)s: %(message)s', level=logging.INFO)

        try:
            world = self.client.get_world()

            self.traffic_manager = self.client.get_trafficmanager(self.tm_port)
            self.traffic_manager.set_global_distance_to_leading_vehicle(2.5)
            if self.respawn:
                self.traffic_manager.set_respawn_dormant_vehicles(True)
            if self.hybrid:
                self.traffic_manager.set_hybrid_physics_mode(True)
                self.traffic_manager.set_hybrid_physics_radius(70.0)
            if self.seed is not None:
                self.traffic_manager.set_random_device_seed(self.seed)

            settings = world.get_settings()
            if not self.asynch:
                self.traffic_manager.set_synchronous_mode(True)
                if not settings.synchronous_mode:
                    self.synchronous_master = True
                    settings.synchronous_mode = True
                    settings.fixed_delta_seconds = 0.05
                else:
                    self.synchronous_master = False
            else:
                print("You are currently in asynchronous mode. If this is a traffic simulation, \
                you could experience some issues. If it's not working correctly, switch to synchronous \
                mode by using traffic_manager.set_synchronous_mode(True)")

            if self.no_rendering:
                settings.no_rendering_mode = True
            world.apply_settings(settings)

            blueprints = self.get_actor_blueprints(self.filterv, self.generationv)
            blueprintsWalkers = self.get_actor_blueprints(self.filterw, self.generationw)

            if self.safe:
                blueprints = [x for x in blueprints if int(x.get_attribute('number_of_wheels')) == 4]
                blueprints = [x for x in blueprints if not x.id.endswith('microlino')]
                blueprints = [x for x in blueprints if not x.id.endswith('carlacola')]
                blueprints = [x for x in blueprints if not x.id.endswith('cybertruck')]
                blueprints = [x for x in blueprints if not x.id.endswith('t2')]
                blueprints = [x for x in blueprints if not x.id.endswith('sprinter')]
                blueprints = [x for x in blueprints if not x.id.endswith('firetruck')]
                blueprints = [x for x in blueprints if not x.id.endswith('ambulance')]

            blueprints = sorted(blueprints, key=lambda bp: bp.id)

            spawn_points = world.get_map().get_spawn_points()
            number_of_spawn_points = len(spawn_points)

            if self.number_of_vehicles < number_of_spawn_points:
                random.shuffle(spawn_points)
            elif self.number_of_vehicles > number_of_spawn_points:
                msg = 'requested %d vehicles, but could only find %d spawn points'
                logging.warning(msg, self.number_of_vehicles, number_of_spawn_points)
                self.number_of_vehicles = number_of_spawn_points

            # @todo cannot import these directly.
            SpawnActor = carla.command.SpawnActor
            SetAutopilot = carla.command.SetAutopilot
            FutureActor = carla.command.FutureActor

            # --------------
            # Spawn vehicles
            # --------------
            batch = []
            hero = self.hero
            for n, transform in enumerate(spawn_points):
                if n >= self.number_of_vehicles:
                    break
                blueprint = random.choice(blueprints)
                if blueprint.has_attribute('color'):
                    color = random.choice(blueprint.get_attribute('color').recommended_values)
                    blueprint.set_attribute('color', color)
                if blueprint.has_attribute('driver_id'):
                    driver_id = random.choice(blueprint.get_attribute('driver_id').recommended_values)
                    blueprint.set_attribute('driver_id', driver_id)
                if hero:
                    blueprint.set_attribute('role_name', 'hero')
                    hero = False
                else:
                    blueprint.set_attribute('role_name', 'autopilot')

                # spawn the cars and set their autopilot and light state all together
                batch.append(SpawnActor(blueprint, transform)
                             .then(SetAutopilot(FutureActor, True, self.traffic_manager.get_port())))

            for response in self.client.apply_batch_sync(batch, self.synchronous_master):
                if response.error:
                    logging.error(response.error)
                else:
                    self.vehicles_list.append(response.actor_id)

            # Set automatic vehicle lights update if specified
            if self.car_lights_on:
                all_vehicle_actors = world.get_actors(self.vehicles_list)
                for actor in all_vehicle_actors:
                    self.traffic_manager.update_vehicle_lights(actor, True)

            # -------------
            # Spawn Walkers
            # -------------
            # some settings
            percentagePedestriansRunning = 0.0  # how many pedestrians will run
            percentagePedestriansCrossing = 0.0  # how many pedestrians will walk through the road
            if self.seedw:
                world.set_pedestrians_seed(self.seedw)
                random.seed(self.seedw)
            # 1. take all the random locations to spawn
            spawn_points = []
            for i in range(self.number_of_walkers):
                spawn_point = carla.Transform()
                loc = world.get_random_location_from_navigation()
                if loc != None:
                    spawn_point.location = loc
                    spawn_points.append(spawn_point)
            # 2. we spawn the walker object
            batch = []
            walker_speed = []
            for spawn_point in spawn_points:
                walker_bp = random.choice(blueprintsWalkers)
                # set as not invincible
                if walker_bp.has_attribute('is_invincible'):
                    walker_bp.set_attribute('is_invincible', 'false')
                # set the max speed
                if walker_bp.has_attribute('speed'):
                    if random.random() > percentagePedestriansRunning:
                        # walking
                        walker_speed.append(walker_bp.get_attribute('speed').recommended_values[1])
                    else:
                        # running
                        walker_speed.append(walker_bp.get_attribute('speed').recommended_values[2])
                else:
                    print("Walker has no speed")
                    walker_speed.append(0.0)
                batch.append(SpawnActor(walker_bp, spawn_point))
            results = self.client.apply_batch_sync(batch, True)
            walker_speed2 = []
            for i in range(len(results)):
                if results[i].error:
                    logging.error(results[i].error)
                else:
                    self.walkers_list.append({"id": results[i].actor_id})
                    walker_speed2.append(walker_speed[i])
            walker_speed = walker_speed2
            # 3. we spawn the walker controller
            batch = []
            walker_controller_bp = world.get_blueprint_library().find('controller.ai.walker')
            for i in range(len(self.walkers_list)):
                batch.append(SpawnActor(walker_controller_bp, carla.Transform(), self.walkers_list[i]["id"]))
            results = self.client.apply_batch_sync(batch, True)
            for i in range(len(results)):
                if results[i].error:
                    logging.error(results[i].error)
                else:
                    self.walkers_list[i]["con"] = results[i].actor_id
            # 4. we put together the walkers and controllers id to get the objects from their id
            for i in range(len(self.walkers_list)):
                self.all_id.append(self.walkers_list[i]["con"])
                self.all_id.append(self.walkers_list[i]["id"])
            all_actors = world.get_actors(self.all_id)

            """
            # wait for a tick to ensure client receives the last transform of the walkers we have just created
            if self.asynch or not self.synchronous_master:
                world.wait_for_tick()
            else:
                world.tick()
            """

            # 5. initialize each controller and set target to walk to (list is [controller, actor, controller,
            # actor ...]) set how many pedestrians can cross the road
            world.set_pedestrians_cross_factor(percentagePedestriansCrossing)
            for i in range(0, len(self.all_id), 2):
                # start walker
                all_actors[i].start()
                # set walk to random point
                all_actors[i].go_to_location(world.get_random_location_from_navigation())
                # max speed
                all_actors[i].set_max_speed(float(walker_speed[int(i / 2)]))

            print('spawned %d vehicles and %d walkers, press Ctrl+C to exit.' % (len(self.vehicles_list),
                                                                                 len(self.walkers_list)))

            # Example of how to use Traffic Manager parameters
            self.traffic_manager.global_percentage_speed_difference(30.0)

            while not self.is_exit:
                continue

            """
            while True:
                if not self.asynch and self.synchronous_master:
                    world.tick()
                else:
                    world.wait_for_tick()
            """

        finally:

            if not self.asynch and self.synchronous_master:
                settings = world.get_settings()
                settings.synchronous_mode = False
                settings.no_rendering_mode = False
                settings.fixed_delta_seconds = None
                world.apply_settings(settings)

            print('\n Destroying %d vehicles' % len(self.vehicles_list))
            self.client.apply_batch([carla.command.DestroyActor(x) for x in self.vehicles_list])

            # stop walker controllers (list is [controller, actor, controller, actor ...])
            for i in range(0, len(self.all_id), 2):
                all_actors[i].stop()

            print('\n Destroying %d walkers' % len(self.walkers_list))
            self.client.apply_batch([carla.command.DestroyActor(x) for x in self.all_id])

            time.sleep(0.5)
