#!/usr/bin/env python

# Copyright (c) 2018-2020 Intel Corporation
#
# This work is licensed under the terms of the MIT license.
# For a copy, see <https://opensource.org/licenses/MIT>.

"""
This module provides all frequently used data from CARLA via
local buffers to avoid blocking calls to CARLA
"""

from __future__ import print_function

import math
import re
import numpy.random as random
from six import iteritems

import carla

def calculate_velocity(actor):
    """
    Method to calculate the velocity of a actor
    """
    velocity_squared = actor.get_velocity().x**2
    velocity_squared += actor.get_velocity().y**2
    return math.sqrt(velocity_squared)

class CarlaDataProvider(object):  # pylint: disable=too-many-public-methods

    """
    This class provides access to various data of all registered actors
    It buffers the data and updates it on every CARLA tick

    Currently available data:
    - Absolute velocity
    - Location
    - Transform

    Potential additions:
    - Acceleration

    In addition it provides access to the map and the transform of all traffic lights
    """

    _actor_velocity_map = dict()
    _actor_location_map = dict()
    _actor_transform_map = dict()
    _traffic_light_map = dict()
    _carla_actor_pool = dict()
    _client = None
    _world = None
    _map = None
    _sync_flag = False
    _spawn_points = None
    _spawn_index = 0
    _blueprint_library = None
    _ego_vehicle_route = None
    _traffic_manager_port = 8000
    _random_seed = 2000
    _rng = random.RandomState(_random_seed)

    @staticmethod
    def register_actor(actor):
        """
        Add new actor to dictionaries
        If actor already exists, throw an exception
        """
        if actor in CarlaDataProvider._actor_velocity_map:
            raise KeyError(
                "Vehicle '{}' already registered. Cannot register twice!".format(actor.id))
        else:
            CarlaDataProvider._actor_velocity_map[actor] = 0.0

        if actor in CarlaDataProvider._actor_location_map:
            raise KeyError(
                "Vehicle '{}' already registered. Cannot register twice!".format(actor.id))
        else:
            CarlaDataProvider._actor_location_map[actor] = None

        if actor in CarlaDataProvider._actor_transform_map:
            raise KeyError(
                "Vehicle '{}' already registered. Cannot register twice!".format(actor.id))
        else:
            CarlaDataProvider._actor_transform_map[actor] = None

    @staticmethod
    def register_actors(actors):
        """
        Add new set of actors to dictionaries
        """
        for actor in actors:
            CarlaDataProvider.register_actor(actor)

    @staticmethod
    def on_carla_tick():
        """
        Callback from CARLA
        """
        for actor in CarlaDataProvider._actor_velocity_map:
            if actor is not None and actor.is_alive:
                CarlaDataProvider._actor_velocity_map[actor] = calculate_velocity(actor)

        for actor in CarlaDataProvider._actor_location_map:
            if actor is not None and actor.is_alive:
                CarlaDataProvider._actor_location_map[actor] = actor.get_location()

        for actor in CarlaDataProvider._actor_transform_map:
            if actor is not None and actor.is_alive:
                CarlaDataProvider._actor_transform_map[actor] = actor.get_transform()

        world = CarlaDataProvider._world
        if world is None:
            print("WARNING: CarlaDataProvider couldn't find the world")

    @staticmethod
    def get_velocity(actor):
        """
        returns the absolute velocity for the given actor
        """
        for key in CarlaDataProvider._actor_velocity_map:
            if key.id == actor.id:
                return CarlaDataProvider._actor_velocity_map[key]

        # We are intentionally not throwing here
        # This may cause exception loops in py_trees
        print('{}.get_velocity: {} not found!' .format(__name__, actor))
        return 0.0

    @staticmethod
    def get_location(actor):
        """
        returns the location for the given actor
        """
        for key in CarlaDataProvider._actor_location_map:
            if key.id == actor.id:
                return CarlaDataProvider._actor_location_map[key]

        # We are intentionally not throwing here
        # This may cause exception loops in py_trees
        print('{}.get_location: {} not found!' .format(__name__, actor))
        return None

    @staticmethod
    def get_transform(actor):
        """
        returns the transform for the given actor
        """
        for key in CarlaDataProvider._actor_transform_map:
            if key.id == actor.id:
                return CarlaDataProvider._actor_transform_map[key]

        # We are intentionally not throwing here
        # This may cause exception loops in py_trees
        print('{}.get_transform: {} not found!' .format(__name__, actor))
        return None

    @staticmethod
    def set_client(client):
        """
        Set the CARLA client
        """
        CarlaDataProvider._client = client

    @staticmethod
    def get_client():
        """
        Get the CARLA client
        """
        return CarlaDataProvider._client

    @staticmethod
    def set_world(world):
        """
        Set the world and world settings
        """
        CarlaDataProvider._world = world
        CarlaDataProvider._sync_flag = world.get_settings().synchronous_mode
        CarlaDataProvider._map = world.get_map()
        CarlaDataProvider._blueprint_library = world.get_blueprint_library()
        CarlaDataProvider.generate_spawn_points()
        CarlaDataProvider.prepare_map()

    @staticmethod
    def get_world():
        """
        Return world
        """
        return CarlaDataProvider._world

    @staticmethod
    def get_map(world=None):
        """
        Get the current map
        """
        if CarlaDataProvider._map is None:
            if world is None:
                if CarlaDataProvider._world is None:
                    raise ValueError("class member \'world'\' not initialized yet")
                else:
                    CarlaDataProvider._map = CarlaDataProvider._world.get_map()
            else:
                CarlaDataProvider._map = world.get_map()

        return CarlaDataProvider._map

    @staticmethod
    def is_sync_mode():
        """
        @return true if syncronuous mode is used
        """
        return CarlaDataProvider._sync_flag

    @staticmethod
    def find_weather_presets():
        """
        Get weather presets from CARLA
        """
        rgx = re.compile('.+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)')
        name = lambda x: ' '.join(m.group(0) for m in rgx.finditer(x))
        presets = [x for x in dir(carla.WeatherParameters) if re.match('[A-Z].+', x)]
        return [(getattr(carla.WeatherParameters, x), name(x)) for x in presets]

    @staticmethod
    def prepare_map():
        """
        This function set the current map and loads all traffic lights for this map to
        _traffic_light_map
        """
        if CarlaDataProvider._map is None:
            CarlaDataProvider._map = CarlaDataProvider._world.get_map()

        # Parse all traffic lights
        CarlaDataProvider._traffic_light_map.clear()
        for traffic_light in CarlaDataProvider._world.get_actors().filter('*traffic_light*'):
            if traffic_light not in CarlaDataProvider._traffic_light_map.keys():
                CarlaDataProvider._traffic_light_map[traffic_light] = traffic_light.get_transform()
            else:
                raise KeyError(
                    "Traffic light '{}' already registered. Cannot register twice!".format(traffic_light.id))

    @staticmethod
    def annotate_trafficlight_in_group(traffic_light):
        """
        Get dictionary with traffic light group info for a given traffic light
        """
        dict_annotations = {'ref': [], 'opposite': [], 'left': [], 'right': []}

        # Get the waypoints
        ref_location = CarlaDataProvider.get_trafficlight_trigger_location(traffic_light)
        ref_waypoint = CarlaDataProvider.get_map().get_waypoint(ref_location)
        ref_yaw = ref_waypoint.transform.rotation.yaw

        group_tl = traffic_light.get_group_traffic_lights()

        for target_tl in group_tl:
            if traffic_light.id == target_tl.id:
                dict_annotations['ref'].append(target_tl)
            else:
                # Get the angle between yaws
                target_location = CarlaDataProvider.get_trafficlight_trigger_location(target_tl)
                target_waypoint = CarlaDataProvider.get_map().get_waypoint(target_location)
                target_yaw = target_waypoint.transform.rotation.yaw

                diff = (target_yaw - ref_yaw) % 360

                if diff > 330:
                    continue
                elif diff > 225:
                    dict_annotations['right'].append(target_tl)
                elif diff > 135.0:
                    dict_annotations['opposite'].append(target_tl)
                elif diff > 30:
                    dict_annotations['left'].append(target_tl)

        return dict_annotations

    @staticmethod
    def get_trafficlight_trigger_location(traffic_light):    # pylint: disable=invalid-name
        """
        Calculates the yaw of the waypoint that represents the trigger volume of the traffic light
        """
        def rotate_point(point, angle):
            """
            rotate a given point by a given angle
            """
            x_ = math.cos(math.radians(angle)) * point.x - math.sin(math.radians(angle)) * point.y
            y_ = math.sin(math.radians(angle)) * point.x - math.cos(math.radians(angle)) * point.y

            return carla.Vector3D(x_, y_, point.z)

        base_transform = traffic_light.get_transform()
        base_rot = base_transform.rotation.yaw
        area_loc = base_transform.transform(traffic_light.trigger_volume.location)
        area_ext = traffic_light.trigger_volume.extent

        point = rotate_point(carla.Vector3D(0, 0, area_ext.z), base_rot)
        point_location = area_loc + carla.Location(x=point.x, y=point.y)

        return carla.Location(point_location.x, point_location.y, point_location.z)

    @staticmethod
    def update_light_states(ego_light, annotations, states, freeze=False, timeout=1000000000):
        """
        Update traffic light states
        """
        reset_params = []

        for state in states:
            relevant_lights = []
            if state == 'ego':
                relevant_lights = [ego_light]
            else:
                relevant_lights = annotations[state]
            for light in relevant_lights:
                prev_state = light.get_state()
                prev_green_time = light.get_green_time()
                prev_red_time = light.get_red_time()
                prev_yellow_time = light.get_yellow_time()
                reset_params.append({'light': light,
                                     'state': prev_state,
                                     'green_time': prev_green_time,
                                     'red_time': prev_red_time,
                                     'yellow_time': prev_yellow_time})

                light.set_state(states[state])
                if freeze:
                    light.set_green_time(timeout)
                    light.set_red_time(timeout)
                    light.set_yellow_time(timeout)

        return reset_params

    @staticmethod
    def reset_lights(reset_params):
        """
        Reset traffic lights
        """
        for param in reset_params:
            param['light'].set_state(param['state'])
            param['light'].set_green_time(param['green_time'])
            param['light'].set_red_time(param['red_time'])
            param['light'].set_yellow_time(param['yellow_time'])

    @staticmethod
    def get_next_traffic_light(actor, use_cached_location=True):
        """
        returns the next relevant traffic light for the provided actor
        """

        if not use_cached_location:
            location = actor.get_transform().location
        else:
            location = CarlaDataProvider.get_location(actor)

        waypoint = CarlaDataProvider.get_map().get_waypoint(location)
        # Create list of all waypoints until next intersection
        list_of_waypoints = []
        while waypoint and not waypoint.is_intersection:
            list_of_waypoints.append(waypoint)
            waypoint = waypoint.next(2.0)[0]

        # If the list is empty, the actor is in an intersection
        if not list_of_waypoints:
            return None

        relevant_traffic_light = None
        distance_to_relevant_traffic_light = float("inf")

        for traffic_light in CarlaDataProvider._traffic_light_map:
            if hasattr(traffic_light, 'trigger_volume'):
                tl_t = CarlaDataProvider._traffic_light_map[traffic_light]
                transformed_tv = tl_t.transform(traffic_light.trigger_volume.location)
                distance = carla.Location(transformed_tv).distance(list_of_waypoints[-1].transform.location)

                if distance < distance_to_relevant_traffic_light:
                    relevant_traffic_light = traffic_light
                    distance_to_relevant_traffic_light = distance

        return relevant_traffic_light

    @staticmethod
    def set_ego_vehicle_route(route):
        """
        Set the route of the ego vehicle

        @todo extend ego_vehicle_route concept to support multi ego_vehicle scenarios
        """
        CarlaDataProvider._ego_vehicle_route = route

    @staticmethod
    def get_ego_vehicle_route():
        """
        returns the currently set route of the ego vehicle
        Note: Can be None
        """
        return CarlaDataProvider._ego_vehicle_route

    @staticmethod
    def generate_spawn_points():
        """
        Generate spawn points for the current map
        """
        spawn_points = list(CarlaDataProvider.get_map(CarlaDataProvider._world).get_spawn_points())
        CarlaDataProvider._rng.shuffle(spawn_points)
        CarlaDataProvider._spawn_points = spawn_points
        CarlaDataProvider._spawn_index = 0

    @staticmethod
    def create_blueprint(model, rolename='scenario', color=None, actor_category="car"):
        """
        Function to setup the blueprint of an actor given its model and other relevant parameters
        """

        _actor_blueprint_categories = {
            'car': 'vehicle.tesla.model3',
            'van': 'vehicle.volkswagen.t2',
            'truck': 'vehicle.carlamotors.carlacola',
            'trailer': '',
            'semitrailer': '',
            'bus': 'vehicle.volkswagen.t2',
            'motorbike': 'vehicle.kawasaki.ninja',
            'bicycle': 'vehicle.diamondback.century',
            'train': '',
            'tram': '',
            'pedestrian': 'walker.pedestrian.0001',
        }

        # Set the model
        try:
            blueprint = CarlaDataProvider._rng.choice(CarlaDataProvider._blueprint_library.filter(model))
        except ValueError:
            # The model is not part of the blueprint library. Let's take a default one for the given category
            if actor_category == "car":
                bp_filter = "vehicle.*"
            elif actor_category == "pedestrian":
                bp_filter = "walker.pedestrian.*"
            new_model = _actor_blueprint_categories[actor_category]
            if new_model != '':
                bp_filter = new_model
            print("WARNING: Actor model {} not available. Using instead {}".format(model, new_model))
            blueprint = CarlaDataProvider._rng.choice(CarlaDataProvider._blueprint_library.filter(bp_filter))

        # Set the color
        if color:
            if not blueprint.has_attribute('color'):
                print(
                    "WARNING: Cannot set Color ({}) for actor {} due to missing blueprint attribute".format(
                        color, blueprint.id))
            else:
                default_color_rgba = blueprint.get_attribute('color').as_color()
                default_color = '({}, {}, {})'.format(default_color_rgba.r, default_color_rgba.g, default_color_rgba.b)
                try:
                    blueprint.set_attribute('color', color)
                except ValueError:
                    # Color can't be set for this vehicle
                    print("WARNING: Color ({}) cannot be set for actor {}. Using instead: ({})".format(
                        color, blueprint.id, default_color))
                    blueprint.set_attribute('color', default_color)
        else:
            if blueprint.has_attribute('color') and rolename != 'hero':
                color = CarlaDataProvider._rng.choice(blueprint.get_attribute('color').recommended_values)
                blueprint.set_attribute('color', color)

        # Make pedestrians mortal
        if blueprint.has_attribute('is_invincible'):
            blueprint.set_attribute('is_invincible', 'false')

        # Set the rolename
        if blueprint.has_attribute('role_name'):
            blueprint.set_attribute('role_name', rolename)

        return blueprint

    @staticmethod
    def handle_actor_batch(batch):
        """
        Forward a CARLA command batch to spawn actors to CARLA, and gather the responses.
        Returns list of actors on success, none otherwise
        """

        actors = []

        sync_mode = CarlaDataProvider.is_sync_mode()

        if CarlaDataProvider._client and batch is not None:
            responses = CarlaDataProvider._client.apply_batch_sync(batch, sync_mode)
        else:
            return None

        # wait for the actors to be spawned properly before we do anything
        if sync_mode:
            CarlaDataProvider._world.tick()
        else:
            CarlaDataProvider._world.wait_for_tick()

        actor_ids = []
        if responses:
            for response in responses:
                if not response.error:
                    actor_ids.append(response.actor_id)

        carla_actors = CarlaDataProvider._world.get_actors(actor_ids)
        for actor in carla_actors:
            actors.append(actor)

        return actors

    @staticmethod
    def request_new_actor(model, spawn_point, rolename='scenario', autopilot=False,
                          random_location=False, color=None, actor_category="car"):
        """
        This method tries to create a new actor, returning it if successful (None otherwise).
        """
        blueprint = CarlaDataProvider.create_blueprint(model, rolename, color, actor_category)

        if random_location:
            actor = None
            while not actor:
                spawn_point = CarlaDataProvider._rng.choice(CarlaDataProvider._spawn_points)
                actor = CarlaDataProvider._world.try_spawn_actor(blueprint, spawn_point)

        else:
            # slightly lift the actor to avoid collisions with ground when spawning the actor
            # DO NOT USE spawn_point directly, as this will modify spawn_point permanently
            _spawn_point = carla.Transform(carla.Location(), spawn_point.rotation)
            _spawn_point.location.x = spawn_point.location.x
            _spawn_point.location.y = spawn_point.location.y
            _spawn_point.location.z = spawn_point.location.z + 0.2
            actor = CarlaDataProvider._world.try_spawn_actor(blueprint, _spawn_point)

        if actor is None:
            raise RuntimeError(
                "Error: Unable to spawn vehicle {} at {}".format(blueprint.id, spawn_point))
        else:
            # Let's deactivate the autopilot of the actor if it belongs to vehicle
            if actor in CarlaDataProvider._blueprint_library.filter('vehicle.*'):
                actor.set_autopilot(autopilot)
            else:
                pass

        # wait for the actor to be spawned properly before we do anything
        if CarlaDataProvider.is_sync_mode():
            CarlaDataProvider._world.tick()
        else:
            CarlaDataProvider._world.wait_for_tick()

        if actor is None:
            return None

        CarlaDataProvider._carla_actor_pool[actor.id] = actor
        CarlaDataProvider.register_actor(actor)
        return actor

    @staticmethod
    def request_new_actors(actor_list):
        """
        This method tries to series of actor in batch. If this was successful,
        the new actors are returned, None otherwise.

        param:
        - actor_list: list of ActorConfigurationData
        """

        SpawnActor = carla.command.SpawnActor               # pylint: disable=invalid-name
        PhysicsCommand = carla.command.SetSimulatePhysics   # pylint: disable=invalid-name
        FutureActor = carla.command.FutureActor             # pylint: disable=invalid-name
        ApplyTransform = carla.command.ApplyTransform       # pylint: disable=invalid-name
        SetAutopilot = carla.command.SetAutopilot           # pylint: disable=invalid-name

        batch = []
        actors = []

        CarlaDataProvider.generate_spawn_points()

        for actor in actor_list:

            # Get the blueprint
            blueprint = CarlaDataProvider.create_blueprint(actor.model, actor.rolename, actor.color, actor.category)

            # Get the spawn point
            transform = actor.transform
            if actor.random_location:
                if CarlaDataProvider._spawn_index >= len(CarlaDataProvider._spawn_points):
                    print("No more spawn points to use")
                    break
                else:
                    _spawn_point = CarlaDataProvider._spawn_points[CarlaDataProvider._spawn_index]
                    CarlaDataProvider._spawn_index += 1

            else:
                _spawn_point = carla.Transform()
                _spawn_point.rotation = transform.rotation
                _spawn_point.location.x = transform.location.x
                _spawn_point.location.y = transform.location.y
                _spawn_point.location.z = transform.location.z + 0.2

            # Get the command
            command = SpawnActor(blueprint, _spawn_point)
            command.then(SetAutopilot(FutureActor, actor.autopilot,
                                      CarlaDataProvider._traffic_manager_port))

            if actor.category == 'misc':
                command.then(PhysicsCommand(FutureActor, True))
            elif actor.args is not None and 'physics' in actor.args and actor.args['physics'] == "off":
                command.then(ApplyTransform(FutureActor, _spawn_point)).then(PhysicsCommand(FutureActor, False))

            batch.append(command)

        actors = CarlaDataProvider.handle_actor_batch(batch)

        if not actors:
            return None

        for actor in actors:
            if actor is None:
                continue
            CarlaDataProvider._carla_actor_pool[actor.id] = actor
            CarlaDataProvider.register_actor(actor)
        return actors

    @staticmethod
    def request_new_batch_actors(model, amount, spawn_points, autopilot=False,
                                 random_location=False, rolename='scenario', walker=False):
        """
        Simplified version of "request_new_actors". This method also create several actors in batch.

        Instead of needing a list of ActorConfigurationData, an "amount" parameter is used.
        This makes actor spawning easier but reduces the amount of configurability.

        Some parameters are the same for all actors (rolename, autopilot and random location)
        while others are randomized (color)
        """

        SpawnActor = carla.command.SpawnActor      # pylint: disable=invalid-name
        SetAutopilot = carla.command.SetAutopilot  # pylint: disable=invalid-name
        FutureActor = carla.command.FutureActor    # pylint: disable=invalid-name

        CarlaDataProvider.generate_spawn_points()
        batch = []
        for i in range(amount):
            # Get vehicle or pedestrian by model
            if walker:
                blueprint = CarlaDataProvider.create_blueprint(model, rolename,  actor_category="pedestrian")
            else:
                blueprint = CarlaDataProvider.create_blueprint(model, rolename)

            if random_location:
                if CarlaDataProvider._spawn_index >= len(CarlaDataProvider._spawn_points):
                    print("No more spawn points to use. Spawned {} actors out of {}".format(i + 1, amount))
                    break
                else:
                    spawn_point = CarlaDataProvider._spawn_points[CarlaDataProvider._spawn_index]
                    CarlaDataProvider._spawn_index += 1
            else:
                try:
                    spawn_point = spawn_points[i]
                except IndexError:
                    print("The amount of spawn points is lower than the amount of vehicles spawned")
                    break

            if walker:
                # generate walker actor
                if spawn_point:
                    spawn_point = carla.Transform()
                    spawn_point.location = CarlaDataProvider.get_world().get_random_location_from_navigation()
                    walker_actor_command = SpawnActor(blueprint, spawn_point)
                    batch.append(walker_actor_command)

            elif spawn_point:
                # generate vehicles
                batch.append(SpawnActor(blueprint, spawn_point).then(
                    SetAutopilot(FutureActor, autopilot,
                                 CarlaDataProvider._traffic_manager_port)))

        actors = CarlaDataProvider.handle_actor_batch(batch)

        if walker:
            # generate walker ai controller
            walker_batch = []
            for i in range(len(actors)):
                # Get vehicle or pedestrian by model
                blueprint_ai = CarlaDataProvider.create_blueprint("controller.ai.walker", rolename,
                                                                  actor_category="pedestrian")
                spawn_point = actors[i].get_transform()
                walker_controller_actor_command = SpawnActor(blueprint_ai, spawn_point, actors[i])
                walker_batch.append(walker_controller_actor_command)

            walker_controller_actors = CarlaDataProvider.handle_actor_batch(walker_batch)

            assert len(actors) == len(walker_controller_actors)

            # start all the ai controllers
            for walker_controller_actor in walker_controller_actors:
                walker_controller_actor.start()
                collision_flag = 1
                while collision_flag:
                    try:
                        walker_controller_actor.go_to_location(
                            CarlaDataProvider.get_world().get_random_location_from_navigation())
                        walker_controller_actor.set_max_speed(1 + random.random()/2)
                        collision_flag = 0
                    except:
                        collision_flag = 1
                        print("collision occurs... try again...")
            sync_mode = CarlaDataProvider.is_sync_mode()
            if sync_mode:
                CarlaDataProvider._world.tick()
            else:
                CarlaDataProvider._world.wait_for_tick()
            actors += walker_controller_actors

        if actors is None:
            return None

        for actor in actors:
            if actor is None:
                continue
            CarlaDataProvider._carla_actor_pool[actor.id] = actor
            CarlaDataProvider.register_actor(actor)
        return actors

    @staticmethod
    def get_actors():
        """
        Return list of actors and their ids

        Note: iteritems from six is used to allow compatibility with Python 2 and 3
        """
        return iteritems(CarlaDataProvider._carla_actor_pool)

    @staticmethod
    def actor_id_exists(actor_id):
        """
        Check if a certain id is still at the simulation
        """
        if actor_id in CarlaDataProvider._carla_actor_pool:
            return True

        return False

    @staticmethod
    def get_hero_actor():
        """
        Get the actor object of the hero actor if it exists, returns none otherwise.
        """
        for actor_id in CarlaDataProvider._carla_actor_pool:
            if CarlaDataProvider._carla_actor_pool[actor_id].attributes['role_name'] == 'hero':
                return CarlaDataProvider._carla_actor_pool[actor_id]
        return None

    @staticmethod
    def get_actor_by_id(actor_id):
        """
        Get an actor from the pool by using its ID. If the actor
        does not exist, None is returned.
        """
        if actor_id in CarlaDataProvider._carla_actor_pool:
            return CarlaDataProvider._carla_actor_pool[actor_id]

        print("Non-existing actor id {}".format(actor_id))
        return None

    @staticmethod
    def remove_actor_by_id(actor_id):
        """
        Remove an actor from the pool using its ID
        """
        if actor_id in CarlaDataProvider._carla_actor_pool:
            CarlaDataProvider._carla_actor_pool[actor_id].destroy()
            CarlaDataProvider._carla_actor_pool[actor_id] = None
            CarlaDataProvider._carla_actor_pool.pop(actor_id)
        else:
            print("Trying to remove a non-existing actor id {}".format(actor_id))

    @staticmethod
    def remove_actors_in_surrounding(location, distance):
        """
        Remove all actors from the pool that are closer than distance to the
        provided location
        """
        for actor_id in CarlaDataProvider._carla_actor_pool.copy():
            if CarlaDataProvider._carla_actor_pool[actor_id].get_location().distance(location) < distance:
                CarlaDataProvider._carla_actor_pool[actor_id].destroy()
                CarlaDataProvider._carla_actor_pool.pop(actor_id)

        # Remove all keys with None values
        CarlaDataProvider._carla_actor_pool = dict({k: v for k, v in CarlaDataProvider._carla_actor_pool.items() if v})

    @staticmethod
    def get_traffic_manager_port():
        """
        Get the port of the traffic manager.
        """
        return CarlaDataProvider._traffic_manager_port

    @staticmethod
    def set_traffic_manager_port(tm_port):
        """
        Set the port to use for the traffic manager.
        """
        CarlaDataProvider._traffic_manager_port = tm_port

    @staticmethod
    def cleanup():
        """
        Cleanup and remove all entries from all dictionaries
        """
        DestroyActor = carla.command.DestroyActor       # pylint: disable=invalid-name
        batch = []

        for actor_id in CarlaDataProvider._carla_actor_pool.copy():
            actor = CarlaDataProvider._carla_actor_pool[actor_id]
            if actor.is_alive:
                batch.append(DestroyActor(actor))

        if CarlaDataProvider._client:
            try:
                CarlaDataProvider._client.apply_batch_sync(batch)
            except RuntimeError as e:
                if "time-out" in str(e):
                    pass
                else:
                    raise e

        CarlaDataProvider._actor_velocity_map.clear()
        CarlaDataProvider._actor_location_map.clear()
        CarlaDataProvider._actor_transform_map.clear()
        CarlaDataProvider._traffic_light_map.clear()
        CarlaDataProvider._map = None
        CarlaDataProvider._world = None
        CarlaDataProvider._sync_flag = False
        CarlaDataProvider._ego_vehicle_route = None
        CarlaDataProvider._carla_actor_pool = dict()
        CarlaDataProvider._client = None
        CarlaDataProvider._spawn_points = None
        CarlaDataProvider._spawn_index = 0
        CarlaDataProvider._rng = random.RandomState(CarlaDataProvider._random_seed)


  from concurrent.futures import ProcessPoolExecutor
from functools import partial
import numpy as np
import os
import audio
from nnmnkwii.datasets import cmu_arctic
from nnmnkwii.io import hts
from nnmnkwii import preprocessing as P
from hparams import hparams
from os.path import exists
import librosa
import sys

from wavenet_vocoder.util import is_mulaw_quantize, is_mulaw, is_raw
from hparams import hparams

linguistic_feats_path = '/home2/srallaba/challenges/blizzard2018/data/input_files_npy/'
ccoeffs_feats_path = 'ccoefs_ascii/'

g = open('logfile','w')
g.close()

def build_from_path(in_dir, out_dir, num_workers=1, tqdm=lambda x: x):
    executor = ProcessPoolExecutor(max_workers=num_workers)
    futures = []

    speakers = cmu_arctic.available_speakers
    speakers = ['awb']

    wd = cmu_arctic.WavFileDataSource(in_dir, speakers=speakers)
    wav_paths = wd.collect_files()
    speaker_ids = wd.labels

    for index, (speaker_id, wav_path) in enumerate(
            zip(speaker_ids, wav_paths)):
        futures.append(executor.submit(
            partial(_process_utterance, out_dir, index + 1, speaker_id, wav_path, "N/A")))
    return [future.result() for future in tqdm(futures)]

def start_at(labels):
    has_silence = labels[0][-1] == "pau"
    if not has_silence:
        return labels[0][0]
    for i in range(1, len(labels)):
        if labels[i][-1] != "pau":
            return labels[i][0]
    assert False

def end_at(labels):
    has_silence = labels[-1][-1] == "pau"
    if not has_silence:
        return labels[-1][1]
    for i in range(len(labels) - 2, 0, -1):
        if labels[i][-1] != "pau":
            return labels[i][1]
    assert False

def ensure_divisible(sig, length):
    l = len(sig)
    if l % length == 0:
         return sig
    else:
       difference = l % length
       #for k in range(difference):
       #    np.append(sig, sig[-1])
       return sig[:len(sig)-difference]

def ensure_frameperiod(sig, mel):
   length = mel.shape[0]
   l = len(sig)
   g = open('logfile','a')
   g.write('Original length of signal: ' + str(len(sig)) + ' original number of frames ' + str(mel.shape[0]) + '\n')

   if float(80 * length) == l:
      g.write('Not changing anything ' + '\n')
      g.close()
      return sig, mel
   else:
      num_samples = 80 * length
      if num_samples > l:
         difference = int((num_samples - l))
         g.write("Signal is shorter"+ '\n')
         g.close()
         for k in range(difference):
           sig =  np.append(sig, sig[-1])
         return sig, mel

      elif num_samples < l:
         difference = int((l - num_samples))
         g.write("Signal is longer" + '\n')
         g.close()
         return sig[:len(sig)-difference], mel

      else:
         print("This is hard")
         sys.exit()

def _process_utterance(out_dir, index, speaker_id, wav_path, text):
    sr = hparams.sample_rate

    # Load the audio to a numpy array. Resampled if needed
    wav = audio.load_wav(wav_path)

    if hparams.rescaling:
        wav = wav / np.abs(wav).max() * hparams.rescaling_max

    # Mu-law quantize
    if is_mulaw_quantize(hparams.input_type):
        # [0, quantize_channels)
        out = P.mulaw_quantize(wav, hparams.quantize_channels)

        # Trim silences
        start, end = audio.start_and_end_indices(out, hparams.silence_threshold)
        wav = wav[start:end]
        out = out[start:end]
        constant_values = P.mulaw_quantize(0, hparams.quantize_channels)
        out_dtype = np.int16
    elif is_mulaw(hparams.input_type):
        # [-1, 1]
        out = P.mulaw(wav, hparams.quantize_channels)
        constant_values = P.mulaw(0.0, hparams.quantize_channels)
        out_dtype = np.float32
    else:
        # [-1, 1]
        out = wav
        constant_values = 0.0
        out_dtype = np.float32

    #print("Wavepath is ", wav_path)
    filename = wav_path.split('/wav/')[-1].split('.wav')[0]
    fname = filename
    filename = ccoeffs_feats_path + '/' + filename + '.mcep'
    mel_spectrogram = np.loadtxt(filename)
    #print("Shape of mel scptrogram is ", mel_spectrogram.shape)
    # Compute a mel-scale spectrogram from the trimmed wav:
    # (N, D)
    #mel_spectrogram = audio.melspectrogram(wav).astype(np.float32).T
    # lws pads zeros internally before performing stft
    # this is needed to adjust time resolution between audio and mel-spectrogram
    #l, r = audio.lws_pad_lr(wav, hparams.fft_size, audio.get_hop_size())

    # zero pad for quantized signal
    #out = np.pad(out, (l, r), mode="constant", constant_values=constant_values)
    N = mel_spectrogram.shape[0]
    #out = ensure_divisible(out, N)
    #print("Length of out: ", len(out), "N ", N)

    #print("Out and N: ", len(out), N)
    #if len(out) < N * audio.get_hop_size():
        #print("Out and N: ", filename, len(out), N, N * audio.get_hop_size())
    #    sys.exit()
    #assert len(out) >= N * audio.get_hop_size()

    # time resolution adjustment
    # ensure length of raw audio is multiple of hop_size so that we can use
    # transposed convolution to upsample
    #out = out[:N * 80]
    #out = ensure_divisible(out, N)
    g = open('logfile','a')
    g.write("Processing " + fname + '\n')
    g.close()

    out,mel_spectrogram = ensure_frameperiod(out,mel_spectrogram)
    #out = ensure_divisible(out, audio.get_hop_size())
    #assert len(out) % audio.get_hop_size() == 0
    #assert len(out) % N == 0
    timesteps = len(out)
    g = open('logfile','a')
    g.write(fname + ' ' + str(len(out)) + ' ' + str(N) + ' ' + str(len(out) % N) + '\n')
    g.write('\n')
    g.close()

    # Write the spectrograms to disk:
    audio_filename = fname + '-audio-%05d.npy' % index
    mel_filename = fname + '-mel-%05d.npy' % index
    np.save(os.path.join(out_dir, audio_filename),
            out.astype(out_dtype), allow_pickle=False)
    np.save(os.path.join(out_dir, mel_filename),
            mel_spectrogram.astype(np.float32), allow_pickle=False)

    # Return a tuple describing this training example:
    return (audio_filename, mel_filename, timesteps, text, speaker_id)


#!/usr/bin/env python

"""
md_one
========

This script is a part of ForceBalance and runs a single simulation
that may be combined with others to calculate general thermodynamic
properties.

This script is meant to be launched automatically by ForceBalance.

"""

#==================#
#| Global Imports |#
#==================#

import os, sys, re
import argparse
import numpy as np
import importlib as il

from forcebalance.nifty import click
from forcebalance.nifty import lp_dump, lp_load, wopen
from forcebalance.nifty import printcool, printcool_dictionary
from forcebalance.molecule import Molecule
from forcebalance.thermo import energy_derivatives

from collections import OrderedDict

from forcebalance.output import getLogger
logger = getLogger(__name__)

#========================================================#
#| Global, user-tunable variables (simulation settings) |#
#========================================================#

# Note: Only the simulation settings that vary across different
# simulations in a target may be specified on the command line.
parser = argparse.ArgumentParser()
parser.add_argument('-T', '--temp', '--temperature', dest='temperature', type=float,
                    help='Simulation temperature, leave blank for constant energy')
parser.add_argument('-P', '--pres', '--pressure', dest='pressure', type=float,
                    help='Simulation pressure, leave blank for constant volume')
parser.add_argument('-g', '--grad', '--gradient', dest='gradient', action='store_true',
                    help='Calculate gradients for output time series')

# These settings may be specified for debugging purposes (i.e. they
# will override what we read from simulation.p)
parser.add_argument('-eq', '--nequil', dest='nequil', type=int,
                    help='Number of steps for equilibration run (leave blank to use default from simulation.p)')
parser.add_argument('-md', '--nsteps', dest='nsteps', type=int,
                    help='Number of steps for production run (leave blank to use default from simulation.p)')
parser.add_argument('-dt', '--timestep', dest='timestep', type=float,
                    help='Time step in femtoseconds (leave blank to use default from simulation.p)')
parser.add_argument('-sp', '--sample', dest='sample', type=float,
                    help='Sampling interval in picoseconds (leave blank to use default from simulation.p)')
parser.add_argument('-nt', '--threads', dest='threads', type=int,
                    help='Number of threads for running parallel simulations (GMX, TINKER)')
parser.add_argument('-min', '--minimize', dest='minimize', action='store_true',
                    help='Whether to minimize the energy before starting the simulation')
parser.add_argument('-o', '-out', '--output', dest='output', type=str, nargs='+',
                    help='Specify the time series which are written to disk')

# Parse the command line options and save as a dictionary (don't save NoneTypes)
parsed = parser.parse_args()
args = OrderedDict([(i, j) for i, j in vars(parsed).items() if j is not None])

#----
# Load the ForceBalance pickle file which contains:
#----
# - Force field object
# - Optimization parameters
# - Options loaded from file
FF, mvals = lp_load('forcefield.p')
#----
# Load the simulation pickle file which contains:
#----
# - Target options
# - Engine options
# - MD simulation options
TgtOpts, EngOpts, MDOpts = lp_load('simulation.p')
FF.ffdir = '.'

# Engine name.
engname = TgtOpts['engname']

# Import modules and create the correct Engine object.
if engname == "openmm":
    try:
        from simtk.unit import *
        from simtk.openmm import *
        from simtk.openmm.app import *
    except:
        traceback.print_exc()
        raise Exception("Cannot import OpenMM modules")
    from forcebalance.openmmio import *
    EngineClass = OpenMM
elif engname == "gromacs" or engname == "gmx":
    from forcebalance.gmxio import *
    EngineClass = GMX
elif engname == "tinker":
    from forcebalance.tinkerio import *
    EngineClass = TINKER
else:
    raise Exception('OpenMM, GROMACS, and TINKER are supported at this time.')

def main():

    """Usage:

    (prefix.sh) md_one.py -T, --temperature <temperature in kelvin>
                          -P, --pressure <pressure in atm>
                          -g, --grad (if gradients of output timeseries are desired)
                          -eq, --nequil <number of equilibration MD steps>
                          -md, --nsteps <number of production MD steps>
                          -dt, --timestep <number of production MD steps>
                          -sp, --sample <number of production MD steps>
                          -nt, --threads <number of CPU threads to use>
                          -min, --minimize <minimize the energy>

    This program is meant to be called automatically by ForceBalance because
    force field options are loaded from the 'forcefield.p' file, and
    simulation options are loaded from the 'simulation.p' file.
    The files are separated because the same force field file
    may be used for many simulations.

    """

    # Write the force field file.
    FF.make(mvals)

    # Read the command line options (they may override the options from file.)
    AGrad = args['gradient']
    for i in ['temperature', 'pressure', 'nequil', 'nsteps', 'timestep', 'sample', 'threads', 'minimize']:
        if i in args:
            MDOpts[i] = args[i]
    MDOpts['nsave'] = int(1000.0*MDOpts['sample']/MDOpts['timestep'])
    if 'save_traj' in TgtOpts:
        MDOpts['save_traj'] = TgtOpts['save_traj']

    #----
    # Print some options.
    # At this point, engine and MD options should be SET!
    #----
    printcool("ForceBalance simulation using engine: %s" % engname.upper(),
              color=4, bold=True)
    printcool_dictionary(args, title="Options from command line")
    printcool_dictionary(EngOpts, title="Engine options")
    printcool_dictionary(MDOpts, title="Molecular dynamics options")

    #----
    # For convenience, assign some local variables.
    #----
    # Finite difference step size
    h = TgtOpts['h']
    # Active parameters to differentiate
    pgrad = TgtOpts['pgrad']
    # Create instances of the MD Engine objects.
    Engine = EngineClass(**EngOpts)
    click() # Start timer.
    # This line runs the condensed phase simulation.
    #----
    # The molecular dynamics simulation returns a dictionary of properties
    # In the future, the properties will be stored as data inside the object
    Results = Engine.molecular_dynamics(**MDOpts)
    if AGrad:
        Results['Potential_Derivatives'] = energy_derivatives(Engine, FF, mvals, h, pgrad, dipole=False)['potential']
    # Set up engine and calculate the potential in the other phase.
    EngOpts_ = deepcopy(EngOpts)
    EngOpts_['implicit_solvent'] = not EngOpts['implicit_solvent']
    Engine_ = EngineClass(**EngOpts_)
    Engine_.xyz_omms = Engine.xyz_omms
    Energy_ = Engine_.energy()
    Results_ = {'Potentials' : Energy_}
    if AGrad:
        Derivs_ = energy_derivatives(Engine_, FF, mvals, h, pgrad, dipole=False)['potential']
        Results_['Potential_Derivatives'] = Derivs_
    # Calculate the hydration energy of each snapshot and its parametric derivatives.
    if EngOpts['implicit_solvent']:
        Energy_liq = Results['Potentials']
        Energy_gas = Results_['Potentials']
        if AGrad:
            Derivs_liq = Results['Potential_Derivatives']
            Derivs_gas = Results_['Potential_Derivatives']
    else:
        Energy_gas = Results['Potentials']
        Energy_liq = Results_['Potentials']
        if AGrad:
            Derivs_gas = Results['Potential_Derivatives']
            Derivs_liq = Results_['Potential_Derivatives']
    Results['Hydration'] = Energy_liq - Energy_gas
    if AGrad:
        Results['Hydration_Derivatives'] = Derivs_liq - Derivs_gas
    # Code of the future!
    # Don't know how to use it yet though.
    # Engine.molecular_dynamics(**MDOpts)
    # logger.info("MD simulation took %.3f seconds\n" % click())
    # # Extract properties.
    # Results = Engine.md_extract(OrderedDict([(i, {}) for i in Tgt.timeseries.keys()]))
    # potential = properties['Potential']
    # Calculate energy and dipole derivatives if needed.
    # if AGrad:
    #     Results['derivatives'] = energy_derivatives(Engine, FF, mvals, h, pgrad, dipole='dipole' in Tgt.timeseries.keys())
    # Dump results to file
    logger.info("Writing final force field.\n")
    pvals = FF.make(mvals)
    logger.info("Writing all simulation data to disk.\n")
    lp_dump(Results, 'md_result.p')

if __name__ == "__main__":
    main()


# -*- coding: utf-8 -*-

"""EphysTraces class."""

#------------------------------------------------------------------------------
# Imports
#------------------------------------------------------------------------------

import logging

import numpy as np
import dask.array as da
import mtscomp

from .array import _clip

logger = logging.getLogger(__name__)

#------------------------------------------------------------------------------
# Utility functions
#------------------------------------------------------------------------------

def _extract_waveforms(traces, sample, channel_ids=None, n_samples_waveforms=None):
    nsw = n_samples_waveforms
    dur = traces.shape[0]
    a = nsw // 2
    b = nsw - a
    assert traces.ndim == 2
    assert nsw > 0
    assert a + b == nsw
    if channel_ids is None:  # pragma: no cover
        channel_ids = slice(None, None, None)
    t0, t1 = int(sample - a), int(sample + b)
    t0, t1 = _clip(t0, 0, dur), _clip(t1, 0, dur)
    # Extract the waveforms.
    w = traces[t0:t1][:, channel_ids]
    # Pad with zeros.
    bef = aft = 0
    if t0 == 0:  # pragma: no cover
        bef = nsw - w.shape[0]
    if t1 == dur:  # pragma: no cover
        aft = nsw - w.shape[0]
    assert bef + w.shape[0] + aft == nsw
    if bef > 0 or aft > 0:  # pragma: no cover
        w = np.pad(w, ((bef, aft), (0, 0)), 'constant')
    assert w.shape[0] == nsw
    return w

#------------------------------------------------------------------------------
# EphysTraces
#------------------------------------------------------------------------------

class EphysTraces(da.Array):
    """A class representing raw data, an arbitrarily large array (n_samples, n_channels).
    Derives from dask.array.Array, has a few convenient methods and properties.
    Can be used as a NumPy array, but one needs to call `.compute()` eventually to get
    actual values (lazy evaluation).

    Properties
    ----------

    - format: for now, either 'flat' or 'mtscomp
    - sample_rate: float
    - duration: float

    Methods
    -------

    - extract_waveforms(spike_times, channel_ids): with mtscomp, load chunks in parallel.
      Most efficient when the spikes are selected from a limited number of chunks.
    - subset_time_range(t0, t1): return a view of the object, with a smaller time interval

    """

    def __new__(cls, *args, **kwargs):
        format = kwargs.pop('format', 'flat')
        sample_rate = kwargs.pop('sample_rate', None)
        assert sample_rate > 0
        self = super(EphysTraces, cls).__new__(cls, *args, **kwargs)
        self.format = format
        self.sample_rate = sample_rate
        self.duration = self.shape[0] / self.sample_rate
        return self

    @property
    def chunk_bounds(self):
        return [0] + np.cumsum([self.chunks[0]]).tolist()

    def _get_time_chunks(self, spike_times):
        """Return the time chunk indices of every spike."""
        return np.searchsorted(self.chunk_bounds, spike_times, side='right') - 1

    def extract_waveforms(self, spike_times, channel_ids, n_samples_waveforms=None):
        """Extract waveforms for a given set of spikes, on certain channels."""
        # Create the output array.
        ns = len(spike_times)
        nsw = n_samples_waveforms
        assert nsw > 0, "Please specify n_samples_waveforms > 0"
        nc = len(channel_ids)
        out = np.zeros((ns, nsw, nc), dtype=np.float64)

        # NOTE: put those chunks in the mtscomp cache in parallel.
        if self.format == 'mtscomp':
            # Find the time chunks where the spikes belong.
            chunks = np.unique(self._get_time_chunks(spike_times))
            pool = self.reader.start_thread_pool()
            self.reader.decompress_chunks(chunks, pool=pool)
            self.reader.stop_thread_pool()

        # Extract the spike waveforms.
        for i, ts in enumerate(spike_times):
            # Extract waveforms on the fly from raw data.
            out[i, ...] = _extract_waveforms(
                self, ts, channel_ids=channel_ids, n_samples_waveforms=nsw)

        return out

    def subset_time_range(self, t0, t1):
        """Return a new EphysTraces instance for a subset of the time range."""
        i0, i1 = int(round(t0 * self.sample_rate)), int(round(t1 * self.sample_rate))
        return from_dask(self[i0:i1, :], sample_rate=self.sample_rate, format=self.format)

def from_dask(arr, **kwargs):
    """From dask.array.Array instance to EphysTraces instance."""
    return EphysTraces(arr.dask, arr.name, arr.chunks, dtype=arr.dtype, shape=arr.shape, **kwargs)

def from_mtscomp(reader):
    """From a mtscomp-compressed array to EphysTraces."""
    name = 'mtscomp_traces'
    # Use the compression chunks as dask chunks.
    chunks = (tuple(np.diff(reader.chunk_bounds)), (reader.n_channels,))

    def f(i):
        return reader._decompress_chunk(i)[1]
    dask = {(name, i, 0): (f, i) for i in range(reader.n_chunks)}
    assert len(chunks[0]) == len(dask)
    out = EphysTraces(
        dask, name, chunks, dtype=reader.dtype, shape=reader.shape,
        format='mtscomp', sample_rate=reader.sample_rate)
    out.reader = reader
    return out

def from_array(arr, sample_rate):
    """From any NumPy-like array to EphysTraces instance."""
    n_samples, n_channels = arr.shape
    chunk_shape = (int(sample_rate), n_channels)
    dask_arr = da.from_array(arr, chunk_shape, name='traces')
    return from_dask(dask_arr, sample_rate=sample_rate)

def get_ephys_traces(obj, sample_rate=None):
    """Get an EphysTraces instance."""
    if isinstance(obj, mtscomp.Reader):
        return from_mtscomp(obj)
    else:
        assert sample_rate, "Please specify a sample rate."
        assert sample_rate > 0, "Please specify a sample rate as a strictly positive number."
        return from_array(obj, sample_rate)


import numpy as np
# import matplotlib.pyplot as plt
import csv
# import glob

rootPath = 'D:/rongshidata/experiment_data_1'
periodStr = 'monthly'
usedFeature = 'V2_V22'
indexCalcMethod = 'Index1'
modelParams = 'MP1'
usedAlgo = 'Xgb'

predMonthsList = []
predMonthsList.extend(['201502', '201503', '201504', '201505', '201506',
                       '201507', '201508', '201509', '201510', '201511', '201512'])
predMonthsList.extend(['201601', '201602', '201603', '201604', '201605', '201606',
                       '201607', '201608', '201609', '201610', '201611', '201612'])
predMonthsList.extend(['201701', '201702', '201703', '201704', '201705', '201706',
                       '201707', '201708', '201709', '201710', '201711', '201712'])
predMonthsList.extend(['201801', '201802'])

dropSPnum = 40

s = '{0}/{1}/ReportDF_{2}_{3}_{4}'
dfPath = s.format(rootPath, periodStr, usedFeature, indexCalcMethod, modelParams)
algoPath = dfPath.replace('DF', usedAlgo)

histData = []
for predMonth in predMonthsList:
    spDFPath = '{0}/for{1}/StockPool_{1}.txt'.format(dfPath, predMonth)
    spInitPath = '{0}/for{1}/StockPool_{1}.txt'.format(algoPath, predMonth)
    spFilterPath = spInitPath[:-4] + '_filter' + str(dropSPnum) + spInitPath[-4:]
    spDF = np.loadtxt(spDFPath, dtype=bytes).astype(str)
    spInit = np.loadtxt(spInitPath, dtype=bytes).astype(str)
    spIndexList = []
    for sp in spInit:
        index = np.where(spDF == sp)
        if not list(index[0]):
            i = len(spDF)
        else:
            i = index[0][0]

        spIndexList.append(i)

    histData.extend(spIndexList)

    spIndexArray = np.array(spIndexList)
    spIndex = np.argsort(spIndexArray)[::-1]
    spDroped = np.delete(spInit, spIndex[:dropSPnum])
    spIntegrated = np.append(spDroped, spDF[:dropSPnum])
    spIntegrated = spIntegrated.reshape((spIntegrated.shape[0], 1))
    with open(spFilterPath, 'w', newline='') as f:
        cw = csv.writer(f)
        cw.writerows(spIntegrated)

# plt.hist(histData)
# plt.show()


import os
import time

import matplotlib.pyplot as plt
import numpy as np
import torch

from hardware.camera import RealSenseCamera
from hardware.device import get_device
from inference.post_process import post_process_output
from utils.data.camera_data import CameraData
from utils.dataset_processing.grasp import detect_grasps
from utils.visualisation.plot import plot_grasp

class GraspGenerator:
    def __init__(self, saved_model_path, cam_id, visualize=False):
        self.saved_model_path = saved_model_path
        self.camera = RealSenseCamera(device_id=cam_id)

        self.saved_model_path = saved_model_path
        self.model = None
        self.device = None

        self.cam_data = CameraData(include_depth=True, include_rgb=True)

        # Connect to camera
        self.camera.connect()

        # Load camera pose and depth scale (from running calibration)
        self.cam_pose = np.loadtxt('saved_data/camera_pose.txt', delimiter=' ')
        self.cam_depth_scale = np.loadtxt('saved_data/camera_depth_scale.txt', delimiter=' ')

        homedir = os.path.join(os.path.expanduser('~'), "grasp-comms")
        self.grasp_request = os.path.join(homedir, "grasp_request.npy")
        self.grasp_available = os.path.join(homedir, "grasp_available.npy")
        self.grasp_pose = os.path.join(homedir, "grasp_pose.npy")

        if visualize:
            self.fig = plt.figure(figsize=(10, 10))
        else:
            self.fig = None

    def load_model(self):
        print('Loading model... ')
        self.model = torch.load(self.saved_model_path)
        # Get the compute device
        self.device = get_device(force_cpu=False)

    def generate(self):
        # Get RGB-D image from camera
        image_bundle = self.camera.get_image_bundle()
        rgb = image_bundle['rgb']
        depth = image_bundle['aligned_depth']
        x, depth_img, rgb_img = self.cam_data.get_data(rgb=rgb, depth=depth)

        # Predict the grasp pose using the saved model
        with torch.no_grad():
            xc = x.to(self.device)
            pred = self.model.predict(xc)

        q_img, ang_img, width_img = post_process_output(pred['pos'], pred['cos'], pred['sin'], pred['width'])
        grasps = detect_grasps(q_img, ang_img, width_img)

        # Get grasp position from model output
        pos_z = depth[grasps[0].center[0], grasps[0].center[1]] * self.cam_depth_scale - 0.04
        pos_x = np.multiply(grasps[0].center[1] + self.cam_data.top_left[1] - self.camera.intrinsics.ppx,
                            pos_z / self.camera.intrinsics.fx)
        pos_y = np.multiply(grasps[0].center[0] + self.cam_data.top_left[0] - self.camera.intrinsics.ppy,
                            pos_z / self.camera.intrinsics.fy)

        if pos_z == 0:
            return

        target = np.asarray([pos_x, pos_y, pos_z])
        target.shape = (3, 1)
        print('target: ', target)

        # Convert camera to robot coordinates
        camera2robot = self.cam_pose
        target_position = np.dot(camera2robot[0:3, 0:3], target) + camera2robot[0:3, 3:]
        target_position = target_position[0:3, 0]

        # Convert camera to robot angle
        angle = np.asarray([0, 0, grasps[0].angle])
        angle.shape = (3, 1)
        target_angle = np.dot(camera2robot[0:3, 0:3], angle)

        # Concatenate grasp pose with grasp angle
        grasp_pose = np.append(target_position, target_angle[2])

        print('grasp_pose: ', grasp_pose)

        np.save(self.grasp_pose, grasp_pose)

        if self.fig:
            plot_grasp(fig=self.fig, rgb_img=self.cam_data.get_rgb(rgb, False), grasps=grasps, save=True)

    def run(self):
        while True:
            if np.load(self.grasp_request):
                self.generate()
                np.save(self.grasp_request, 0)
                np.save(self.grasp_available, 1)
            else:
                time.sleep(0.1)


import gym
import sys
import os
import itertools
import collections
import unittest
import numpy as np
import tensorflow as tf

from inspect import getsourcefile
current_path = os.path.dirname(os.path.abspath(getsourcefile(lambda: 0)))
import_path = os.path.abspath(os.path.join(current_path, "../.."))

if import_path not in sys.path:
    sys.path.append(import_path)

# from lib import plotting
from lib.atari.state_processor import StateProcessor
from lib.atari import helpers as atari_helpers
from worker import Worker
from estimators import ValueEstimator, PolicyEstimator

def make_env():
    return gym.envs.make("Breakout-v0")

VALID_ACTIONS = [0, 1, 2, 3]

class WorkerTest(tf.test.TestCase):

    def setUp(self):
        super(WorkerTest, self).setUp()

        self.env = make_env()
        self.discount_factor = 0.99
        self.global_step = tf.Variable(0, name="global_step", trainable=False)
        self.global_counter = itertools.count()
        self.sp = StateProcessor()

        with tf.variable_scope("global") as vs:
            self.global_policy_net = PolicyEstimator(len(VALID_ACTIONS))
            self.global_value_net = ValueEstimator(reuse=True)

    def testPolicyNetPredict(self):
        w = Worker(
            name="test",
            env=make_env(),
            policy_net=self.global_policy_net,
            value_net=self.global_value_net,
            global_counter=self.global_counter,
            discount_factor=self.discount_factor)

        with self.test_session() as sess:
            sess.run(tf.initialize_all_variables())
            state = self.sp.process(self.env.reset())
            processed_state = atari_helpers.atari_make_initial_state(state)
            action_values = w._policy_net_predict(processed_state, sess)
            self.assertEqual(action_values.shape, (4,))

    def testValueNetPredict(self):
        w = Worker(
            name="test",
            env=make_env(),
            policy_net=self.global_policy_net,
            value_net=self.global_value_net,
            global_counter=self.global_counter,
            discount_factor=self.discount_factor)

        with self.test_session() as sess:
            sess.run(tf.initialize_all_variables())
            state = self.sp.process(self.env.reset())
            processed_state = atari_helpers.atari_make_initial_state(state)
            state_value = w._value_net_predict(processed_state, sess)
            self.assertEqual(state_value.shape, ())

    def testRunNStepsAndUpdate(self):
        w = Worker(
            name="test",
            env=make_env(),
            policy_net=self.global_policy_net,
            value_net=self.global_value_net,
            global_counter=self.global_counter,
            discount_factor=self.discount_factor)

        with self.test_session() as sess:
            sess.run(tf.initialize_all_variables())
            state = self.sp.process(self.env.reset())
            processed_state = atari_helpers.atari_make_initial_state(state)
            w.state = processed_state
            transitions, local_t, global_t = w.run_n_steps(10, sess)
            policy_net_loss, value_net_loss, policy_net_summaries, value_net_summaries = w.update(
                transitions, sess)

        self.assertEqual(len(transitions), 10)
        self.assertIsNotNone(policy_net_loss)
        self.assertIsNotNone(value_net_loss)
        self.assertIsNotNone(policy_net_summaries)
        self.assertIsNotNone(value_net_summaries)

if __name__ == '__main__':
    unittest.main()


import cv2 as cv  # opencv
import copy  # for deepcopy on images

import numpy as np  # numpy
from random import randint  # for random values

import threading  # for deamon processing

from pathlib import Path  # for directory information
import os  # for directory information

from constants import constants  # constants

class PlantDetector:
    """Dynamically apply detection algorithms to source images

    All images are sourced from and follow naming standard from
    the KOMATSUNA dataset
    http://limu.ait.kyushu-u.ac.jp/~agri/komatsuna/

    METHODS

    __init__(self, src='multi_plant', labels='multi_label') [void]
        prepares the images and labels for display
        initializes windows and trackbars
        runs background subtraction on plant group images

    on_low_H_thresh_trackbar(self, val)
    on_high_H_thresh_trackbar(self, val)
    on_low_S_thresh_trackbar(self, val)
    on_high_S_thresh_trackbar(self, val)
    on_low_V_thresh_trackbar(self, val)
    on_high_V_thresh_trackbar(self, val)
        HSV trackbar triggers

    prepare_plant_collection(self, src, labelsrc)
        returns [plants, plant_groups, labels]
        constructor helper function for loading plant images

    parse(self, auto_inc=False, mode=0) [void]
        main function
        dynamically applies
            HSV inRange filters
            watershed algorithm
        to the currently displayed image
        based on selected HSV trackbar values

        six modes are displayable:

        mode: window1               + window2
        0   : original (fallback)   + original
        1   : HSV filter range      + original
        2   : bare watershed masks  + labels
        3   : watershed masks w/ bg + original
        4   : sequential bg sub     + original
        5   : seq bg sub w/ watersh + original

        additionally, the user is allowed control
        key | function
        m   | next image
        n   | prev image
        s   | save selected image in the selected mode
        z   | save all images in selected mode
        esc | exit the program
        d   | dynamically calculate dice
        f   | show dice data based on saved images
        1-5 | select the respective mode

        parse is also used for saving all images
        parse is run for all images in the given mode
        either in parrallel or in place

    save_one(self, mode, image, filename) [void]
        saves the image in the appropriate mode folder with filename

    HSV_filtering_and_watershed(self, input_im) [mask, input_im, im_threshold]
        image is filtered through HSV inRange according to trackbar values
        image is prepared (threshold) for watershed algorithm
        watershed algorithm is applied
        markers are applied to image

    dicify_wrapper(self, image_id) [void]
        runs dice summary in background

    dicify_summary(self, image_id) [void]
        prints summary of dice values for image, plant, dateset
        note: based on saved images

    dicify_one(self, image_id) [dice]
        returns the dice value for the given image_id
        based on saved segmentation and label images

    dicify_one_dynamic(self, mask, image_id) [dice]
        returns dice value for the given image_id
        based on given mask (current) and saved label image

    dicify_plant(self, plant_id) [mean, min, max]
        returns mean, min and max dice values for images in plant group

    dicify_all(self) [mean, min, max]
        returns mean, min and max dice values for images in dataset
        and for each plant
    """

    def __init__(self, src='multi_plant', labels='multi_label'):
        self.c = constants()
        self.window1 = self.c.window.window1
        self.window2 = self.c.window.window2

        cv.namedWindow(self.window1)
        cv.namedWindow(self.window2)

        cv.moveWindow(self.window2, 550, 90)

        cv.createTrackbar(
            self.c.HSV.low_H_name, self.window1, self.c.HSV.low_H,
            self.c.HSV.max_value_H, self.on_low_H_thresh_trackbar)
        cv.createTrackbar(
            self.c.HSV.high_H_name, self.window1, self.c.HSV.high_H,
            self.c.HSV.max_value_H, self.on_high_H_thresh_trackbar)
        cv.createTrackbar(
            self.c.HSV.low_S_name, self.window1, self.c.HSV.low_S,
            self.c.HSV.max_value, self.on_low_S_thresh_trackbar)
        cv.createTrackbar(
            self.c.HSV.high_S_name, self.window1, self.c.HSV.high_S,
            self.c.HSV.max_value, self.on_high_S_thresh_trackbar)
        cv.createTrackbar(
            self.c.HSV.low_V_name, self.window1, self.c.HSV.low_V,
            self.c.HSV.max_value, self.on_low_V_thresh_trackbar)
        cv.createTrackbar(
            self.c.HSV.high_V_name, self.window1, self.c.HSV.high_V,
            self.c.HSV.max_value, self.on_high_V_thresh_trackbar)

        self.plants, self.plant_groups, self.labels = self.prepare_plant_collection(src, labels)

        # source https://docs.opencv.org/3.4/d1/dc5/tutorial_background_subtraction.html

        for key in self.plant_groups:
            if self.c.bgsub.mod == 'MOG2':
                backSub = cv.createBackgroundSubtractorMOG2(history=60, detectShadows=True)
            elif self.c.bgsub.mod == 'KNN':
                backSub = cv.createBackgroundSubtractorKNN()
            fgMask = None
            for i, image in enumerate(self.plant_groups[key]):
                fgMask = backSub.apply(image)
                self.plant_groups[key][i] = fgMask

    def on_low_H_thresh_trackbar(self, val):
        self.c.HSV.low_H = val
        self.c.HSV.low_H = min(self.c.HSV.high_H-1, self.c.HSV.low_H)
        cv.setTrackbarPos(
            self.c.HSV.low_H_name, self.window1, self.c.HSV.low_H)

    def on_high_H_thresh_trackbar(self, val):
        self.c.HSV.high_H = val
        self.c.HSV.high_H = max(self.c.HSV.high_H, self.c.HSV.low_H+1)
        cv.setTrackbarPos(
            self.c.HSV.high_H_name, self.window1, self.c.HSV.high_H)

    def on_low_S_thresh_trackbar(self, val):
        self.c.HSV.low_S = val
        self.c.HSV.low_S = min(self.c.HSV.high_S-1, self.c.HSV.low_S)
        cv.setTrackbarPos(
            self.c.HSV.low_S_name, self.window1, self.c.HSV.low_S)

    def on_high_S_thresh_trackbar(self, val):
        self.c.HSV.high_S = val
        self.c.HSV.high_S = max(self.c.HSV.high_S, self.c.HSV.low_S+1)
        cv.setTrackbarPos(
            self.c.HSV.high_S_name, self.window1, self.c.HSV.high_S)

    def on_low_V_thresh_trackbar(self, val):
        self.c.HSV.low_V = val
        self.c.HSV.low_V = min(self.c.HSV.high_V-1, self.c.HSV.low_V)
        cv.setTrackbarPos(
            self.c.HSV.low_V_name, self.window1, self.c.HSV.low_V)

    def on_high_V_thresh_trackbar(self, val):
        self.c.HSV.high_V = val
        self.c.HSV.high_V = max(self.c.HSV.high_V, self.c.HSV.low_V+1)
        cv.setTrackbarPos(
            self.c.HSV.high_V_name, self.window1, self.c.HSV.high_V)

    def prepare_plant_collection(self, src, labelsrc):
        plants = []
        plant_groups = dict()
        files = os.listdir(src)
        files.sort()
        for fl in files:
            input_im = cv.imread(src + '/' + fl, cv.IMREAD_COLOR)
            if (input_im is None):
                exit()

            plants.append({
                'p': input_im,
                'n': fl
            })
            group_id = f'{fl.split("_")[1]}{fl.split("_")[2]}'
            if group_id not in plant_groups:
                plant_groups[group_id] = []
            plant_groups[group_id].append(input_im)

        labels = []
        files = os.listdir(labelsrc)
        files.sort()
        for fl in files:
            input_im = cv.imread(labelsrc + '/' + fl)
            if (input_im is None):
                exit()
            labels.append(input_im)

        return plants, plant_groups, labels

    def parse(self, auto_inc=False, mode=0):
        key = 0
        i = 0
        l_tog = False

        while key != self.c.cntr.exit_k:

            if auto_inc and i == len(self.plants):
                break

            image = copy.deepcopy(self.plants[i]['p'])
            group_id = f'{self.plants[i]["n"].split("_")[1]}{self.plants[i]["n"].split("_")[2]}'

            mask, markers, im_threshold = self.HSV_filtering_and_watershed(image)

            _, bgfgSegMarkers, _ = self.HSV_filtering_and_watershed(
                        cv.cvtColor(self.plant_groups[group_id][i % 60], cv.COLOR_GRAY2BGR)
                        )

            if mode == 5:
                alt = bgfgSegMarkers
                text = f'Watershed new areas w/ fg/bg segm. {self.plants[i]["n"]}'
                tcol = (255, 255, 255)
            elif mode == 4:
                alt = copy.deepcopy(self.plant_groups[group_id][i % 60])
                text = f'FG/BG segmentation {self.plants[i]["n"]}'
                tcol = (255, 255, 255)
            elif mode == 3:
                alt = markers
                text = f'Watershed algorithm areas w/ bg {self.plants[i]["n"]}'
                tcol = (0, 0, 0)
            elif mode == 2:
                alt = mask
                text = f'Watershed algorithm areas bare {self.plants[i]["n"]}'
                tcol = (255, 255, 255)
            elif mode == 1:
                alt = im_threshold
                text = f'HSV inRange threshold {self.plants[i]["n"]}'
                tcol = (255, 255, 255)
            else:
                alt = copy.deepcopy(self.plants[i]['p'])
                text = f'Original {self.plants[i]["n"]}'
                tcol = (0, 0, 0)

            if self.c.asth.text:
                cv.putText(alt, text, (0, 20), self.c.asth.font, .5, tcol, 1)

            cv.imshow(self.window1, alt)
            if l_tog:
                cv.imshow(self.window2, self.labels[i])
            else:
                cv.imshow(self.window2, self.plants[i]['p'])

            key = cv.waitKey(10)

            if key == self.c.cntr.prev_k and i > 0:
                i -= 1
            if key == self.c.cntr.next_k and i < len(self.plants) - 1:
                i += 1
            if key == self.c.cntr.save or auto_inc:
                self.save_one(mode, alt, self.plants[i]["n"])
            if key == self.c.cntr.save_all:
                self.parse(True, mode)
            if key == self.c.cntr.dice:
                print(self.dicify_one_dynamic(mask, self.plants[i]['n']))
            if key == self.c.cntr.dice_more:
                self.dicify_wrapper(self.plants[i]['n'])

            if key == self.c.cntr.m1_k:
                mode = 1
                l_tog = False
            elif key == self.c.cntr.m2_k:
                mode = 2
                l_tog = True
            elif key == self.c.cntr.m3_k:
                mode = 3
                l_tog = False
            elif key == self.c.cntr.m4_k:
                mode = 4
                l_tog = False
            elif key == self.c.cntr.m5_k:
                mode = 5
                l_tog = False

            if auto_inc:
                i += 1

    def save_one(self, mode, image, filename):
        Path(f'formatted/{self.c.cntr.modes[mode]}').mkdir(parents=True, exist_ok=True)
        cv.imwrite(f'formatted/{self.c.cntr.modes[mode]}/{filename}', image)

    def HSV_filtering_and_watershed(self, input_im):

        im_threshold = cv.inRange(
            cv.cvtColor(input_im, cv.COLOR_BGR2HSV),
            (self.c.HSV.low_H, self.c.HSV.low_S, self.c.HSV.low_V),
            (self.c.HSV.high_H, self.c.HSV.high_S, self.c.HSV.high_V)
            )

        # source https://docs.opencv.org/master/d3/db4/tutorial_py_watershed.html

        kernel = np.ones((3, 3), np.uint8)
        opening = cv.morphologyEx(im_threshold, cv.MORPH_OPEN, kernel, iterations=5)

        sure_bg = cv.dilate(opening, kernel, iterations=7)

        dist_transform = cv.distanceTransform(opening, cv.DIST_L2, 5)
        _, sure_fg = cv.threshold(dist_transform, 0.3*dist_transform.max(), 255, 0)

        sure_fg = np.uint8(sure_fg)
        unknown = cv.subtract(sure_bg, sure_fg)

        _, markers = cv.connectedComponents(sure_fg)
        markers = markers + 1
        markers[unknown == 255] = 0

        markers = cv.watershed(input_im, markers)

        input_im[markers == -1] = [255, 0, 0]
        for i in range(2, markers.max() + 1):
            input_im[markers == i] = [
                randint(0, 255), randint(0, 255), randint(0, 255)
                ] if self.c.xtra.disco else [
                    (40 + i * 40) % 255, (i * 40) % 255, (50 + i * 40) % 255
                ]

        mask = copy.deepcopy(input_im)
        mask[markers < 2] = [0, 0, 0]

        return mask, input_im, im_threshold

    def dicify_wrapper(self, image_id):
        thread = threading.Thread(target=self.dicify_summary, args=(image_id,), daemon=True)
        thread.start()

    def dicify_summary(self, image_id):
        print(self.dicify_all())

    def dicify_one(self, image_id):

        # Source: https://github.com/Kornelos/CV_MINI_1/blob/master/process_plants.py

        img = cv.imread(f'multi_label/label_{image_id.split("_", 1)[1]}')
        img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
        _, gt = cv.threshold(img, 1, 255, cv.THRESH_BINARY)

        img = cv.imread(f'formatted/ws_mask/{image_id}')
        img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
        _, rt = cv.threshold(img, 1, 255, cv.THRESH_BINARY)

        k = 255

        dice = np.sum(rt[gt == k]) * 2.0 / (np.sum(rt[rt == k]) + np.sum(gt[gt == k]))

        return dice

    def dicify_one_dynamic(self, mask, image_id):
        img = cv.imread(f'multi_label/label_{image_id.split("_", 1)[1]}')
        img = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
        _, gt = cv.threshold(img, 1, 255, cv.THRESH_BINARY)

        img = cv.cvtColor(mask, cv.COLOR_BGR2GRAY)
        _, rt = cv.threshold(img, 1, 255, cv.THRESH_BINARY)

        k = 255

        dice = np.sum(rt[gt == k]) * 2.0 / (np.sum(rt[rt == k]) + np.sum(gt[gt == k]))

        return dice

    def dicify_plant(self, plant_id):
        vals = []
        for im_data in [
            t for t in self.plants
            if t['n'].split('_')[2] == plant_id
                ]:
            vals.append(self.dicify_one(im_data['n']))
        return [np.mean(vals), min(vals), max(vals)]

    def dicify_all(self):
        means = []
        mins = []
        maxs = []
        summ = "id | mean  | min   | max"
        for i in range(0, 5):
            plant = self.dicify_plant(f'0{str(i)}')
            means.append(plant[0])
            mins.append(plant[1])
            maxs.append(plant[2])
            summ += f'\n0{str(i)} | {round(plant[0], 3)} | {round(plant[1], 3)} | {round(plant[2], 3)}'
        summ += f'\nsm | {round(np.mean(means), 3)} | {round(min(mins), 3)} | {round(max(maxs), 3)}'
        return summ

# Main
plDt = PlantDetector()
plDt.parse()


#!/usr/bin/env python
#
# File: run_con_hostage.py
#
# Created: Sunday, August 14 2016 by rejuvyesh <mail@rejuvyesh.com>
#
from __future__ import absolute_import, print_function

import argparse
import json

import numpy as np
import tensorflow as tf

from gym import spaces
import rltools.algos.policyopt
import rltools.log
import rltools.util
from rltools.samplers.serial import SimpleSampler, ImportanceWeightedSampler, DecSampler
from rltools.samplers.parallel import ThreadedSampler, ParallelSampler
from madrl_environments import ObservationBuffer
from madrl_environments.hostage import ContinuousHostageWorld
from rltools.baselines.linear import LinearFeatureBaseline
from rltools.baselines.mlp import MLPBaseline
from rltools.baselines.zero import ZeroBaseline
from rltools.policy.gaussian import GaussianMLPPolicy

from runners.archs import *

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--discount', type=float, default=0.95)
    parser.add_argument('--gae_lambda', type=float, default=0.99)

    parser.add_argument('--interp_alpha', type=float, default=0.5)
    parser.add_argument('--policy_avg_weights', type=str, default='0.3333333,0.3333333,0.3333333')

    parser.add_argument('--n_iter', type=int, default=250)
    parser.add_argument('--sampler', type=str, default='simple')
    parser.add_argument('--sampler_workers', type=int, default=4)
    parser.add_argument('--max_traj_len', type=int, default=500)
    parser.add_argument('--adaptive_batch', action='store_true', default=False)

    parser.add_argument('--n_timesteps', type=int, default=8000)
    parser.add_argument('--n_timesteps_min', type=int, default=1000)
    parser.add_argument('--n_timesteps_max', type=int, default=64000)
    parser.add_argument('--timestep_rate', type=int, default=20)

    parser.add_argument('--is_n_backtrack', type=int, default=1)
    parser.add_argument('--is_randomize_draw', action='store_true', default=False)
    parser.add_argument('--is_n_pretrain', type=int, default=0)
    parser.add_argument('--is_skip_is', action='store_true', default=False)
    parser.add_argument('--is_max_is_ratio', type=float, default=0)

    parser.add_argument('--buffer_size', type=int, default=1)
    parser.add_argument('--n_good', type=int, default=3)
    parser.add_argument('--n_hostage', type=int, default=5)
    parser.add_argument('--n_bad', type=int, default=5)
    parser.add_argument('--n_coop_save', type=int, default=2)
    parser.add_argument('--n_coop_avoid', type=int, default=2)
    parser.add_argument('--n_sensors', type=int, default=20)
    parser.add_argument('--sensor_range', type=float, default=0.2)
    parser.add_argument('--save_reward', type=float, default=3)
    parser.add_argument('--hit_reward', type=float, default=-1)
    parser.add_argument('--encounter_reward', type=float, default=0.01)
    parser.add_argument('--bomb_reward', type=float, default=-10.)

    parser.add_argument('--policy_hidden_spec', type=str, default=GAE_ARCH)
    parser.add_argument('--min_std', type=float, default=0)
    parser.add_argument('--blend_freq', type=int, default=20)

    parser.add_argument('--baseline_type', type=str, default='mlp')
    parser.add_argument('--baseline_hidden_spec', type=str, default=GAE_ARCH)

    parser.add_argument('--max_kl', type=float, default=0.01)
    parser.add_argument('--vf_max_kl', type=float, default=0.01)
    parser.add_argument('--vf_cg_damping', type=float, default=0.01)

    parser.add_argument('--save_freq', type=int, default=20)
    parser.add_argument('--log', type=str, required=False)
    parser.add_argument('--tblog', type=str, default='/tmp/madrl_tb')
    parser.add_argument('--debug', dest='debug', action='store_true')
    parser.add_argument('--no-debug', dest='debug', action='store_false')
    parser.set_defaults(debug=True)

    args = parser.parse_args()

    policy_avg_weights = np.array(map(float, args.policy_avg_weights.split(',')))
    assert len(policy_avg_weights) == args.n_good

    env = ContinuousHostageWorld(args.n_good, args.n_hostage, args.n_bad, args.n_coop_save,
                                 args.n_coop_avoid, n_sensors=args.n_sensors,
                                 sensor_range=args.sensor_range, save_reward=args.save_reward,
                                 hit_reward=args.hit_reward, encounter_reward=args.encounter_reward,
                                 bomb_reward=args.bomb_reward)

    if args.buffer_size > 1:
        env = ObservationBuffer(env, args.buffer_size)

    policies = [GaussianMLPPolicy(agent.observation_space, agent.action_space,
                                  hidden_spec=args.policy_hidden_spec, enable_obsnorm=True,
                                  min_stdev=args.min_std, init_logstdev=0., tblog=args.tblog,
                                  varscope_name='gaussmlp_policy_{}'.format(agid))
                for agid, agent in enumerate(env.agents)]
    if args.blend_freq:
        assert all(
            [agent.observation_space == env.agents[0].observation_space for agent in env.agents])
        target_policy = GaussianMLPPolicy(env.agents[0].observation_space,
                                          env.agents[0].action_space,
                                          hidden_spec=args.policy_hidden_spec, enable_obsnorm=True,
                                          min_stdev=0., init_logstdev=0., tblog=args.tblog,
                                          varscope_name='targetgaussmlp_policy')
    else:
        target_policy = None

    if args.baseline_type == 'linear':
        baselines = [LinearFeatureBaseline(agent.observation_space, enable_obsnorm=True,
                                           varscope_name='linear_baseline_{}'.format(agid))
                     for agid, agent in enumerate(env.agents)]
    elif args.baseline_type == 'mlp':
        baselines = [MLPBaseline(agent.observation_space, args.baseline_hidden_spec,
                                 enable_obsnorm=True, enable_vnorm=True, max_kl=args.vf_max_kl,
                                 damping=args.vf_cg_damping, time_scale=1. / args.max_traj_len,
                                 varscope_name='mlp_baseline_{}'.format(agid))
                     for agid, agent in enumerate(env.agents)]
    else:
        baselines = [ZeroBaseline(agent.observation_space) for agent in env.agents]

    if args.sampler == 'parallel':
        sampler_cls = ParallelSampler
        sampler_args = dict(max_traj_len=args.max_traj_len, n_timesteps=args.n_timesteps,
                            n_timesteps_min=args.n_timesteps_min,
                            n_timesteps_max=args.n_timesteps_max, timestep_rate=args.timestep_rate,
                            adaptive=args.adaptive_batch, enable_rewnorm=True,
                            n_workers=args.sampler_workers, mode='concurrent')

    else:
        raise NotImplementedError()

    step_func = rltools.algos.policyopt.TRPO(max_kl=args.max_kl)

    popt = rltools.algos.policyopt.ConcurrentPolicyOptimizer(
        env=env, policies=policies, baselines=baselines, step_func=step_func,
        discount=args.discount, gae_lambda=args.gae_lambda, sampler_cls=sampler_cls,
        sampler_args=sampler_args, n_iter=args.n_iter, target_policy=target_policy,
        weights=policy_avg_weights, interp_alpha=args.interp_alpha)

    argstr = json.dumps(vars(args), separators=(',', ':'), indent=2)
    rltools.util.header(argstr)
    log_f = rltools.log.TrainingLog(args.log, [('args', argstr)], debug=args.debug)

    with tf.Session() as sess:
        sess.run(tf.initialize_all_variables())
        popt.train(sess, log_f, args.blend_freq, args.save_freq)

if __name__ == '__main__':
    main()


import lightspeed as ls
import psiw
from .base_lot import Lot
import numpy as np
from . import manage_xyz
from .units import *
from rhf_lot import RHF_LOT
from casci_lot_svd import CASCI_LOT_SVD
import sys
from .nifty import custom_redirection

#TODO get rid of get_energy, get_gradient
class PyTC(Lot):
    """
    Level of theory is a wrapper object to do DFT and CASCI calculations
    Inherits from Lot. Requires a PSIW object
    """

    @property
    def psiw(self):
        return self.options['job_data']['psiw']

    @psiw.setter
    def psiw(self,value):
        self.options['job_data']['psiw']=value

    def get_energy(self,coords,multiplicity,state):
        if self.hasRanForCurrentCoords==False or (coords != self.currentCoords).all():
            self.currentCoords = coords.copy()
            geom = manage_xyz.np_to_xyz(self.geom,self.currentCoords)
            self.run(geom)
        tmp = self.search_tuple(self.E,multiplicity)
        return tmp[state][1]*KCAL_MOL_PER_AU

    def get_mm_energy(self,coords):
        if self.hasRanForCurrentCoords==False or (coords != self.currentCoords).all():
            self.currentCoords = coords.copy()
            self.psiw.update_qmmm(coords*ANGSTROM_TO_AU)
        if self.psiw.__class__.__name__=="CASCI_LOT" or self.psiw.__class__.__name__=="CASCI_LOT_SVD":
            return self.psiw.casci.ref.geometry.qmmm.mm_energy
        else:
            return self.psiw.rhf.geometry.qmmm.mm_energy

    def get_mm_gradient(self,coords):
        #TODO need diff variable for hasRan MM energy
        if self.hasRanForCurrentCoords==False or (coords != self.currentCoords).all():
            self.currentCoords = coords.copy()
            self.psiw.update_qmmm(coords*ANGSTROM_TO_AU)
        if self.psiw.__class__.__name__=="CASCI_LOT" or self.psiw.__class__.__name__=="CASCI_LOT_SVD":
            return self.psiw.casci.ref.geometry.qmmm.mm_gradient
        else:
            return self.psiw.rhf.geometry.qmmm.mm_gradient

    def run(self,geom,verbose=False):
        self.E=[]
        self.grada=[]
        #normal update
        coords = manage_xyz.xyz_to_np(geom)
        T = ls.Tensor.array(coords*ANGSTROM_TO_AU)
        def run_code(T):
            self.psiw = self.psiw.update_xyz(T)
            for state in self.states:
                multiplicity=state[0]
                ad_idx=state[1]
                S=multiplicity-1
                if self.psiw.__class__.__name__=="CASCI_LOT" or self.psiw.__class__.__name__=="CASCI_LOT_SVD":
                    self.E.append((multiplicity,self.psiw.compute_energy(S=S,index=ad_idx)))
                    tmp = self.psiw.compute_gradient(S=S,index=ad_idx)
                elif self.psiw.__class__.__name__=="RHF_LOT":
                    self.E.append((multiplicity,self.psiw.compute_energy()))
                    tmp = self.psiw.compute_gradient()
                self.grada.append((multiplicity,tmp[...]))
            if self.do_coupling==True:
                state1=self.states[0][1]
                state2=self.states[1][1]
                tmp = self.psiw.compute_coupling(S=S,indexA=state1,indexB=state2)
                self.coup = tmp[...]

        if not verbose:
            with open('psiw_jobs.txt','a') as out:
                with custom_redirection(out):
                    run_code(T)
        else:
            run_code(T)
                #filename="{}_rhf_update.molden".format(self.node_id)
                #self.psiw.casci.reference.save_molden_file(filename)

        self.hasRanForCurrentCoords=True
        return

    def get_gradient(self,coords,multiplicity,state):
        if self.hasRanForCurrentCoords==False or (coords != self.currentCoords).all():
            self.currentCoords = coords.copy()
            geom = manage_xyz.np_to_xyz(self.geom,self.currentCoords)
            self.run(geom)
        tmp = self.search_tuple(self.grada,multiplicity)
        return np.asarray(tmp[state][1])*ANGSTROM_TO_AU

    def get_coupling(self,coords,multiplicity,state1,state2):
        if self.hasRanForCurrentCoords==False or (coords != self.currentCoords).all():
            self.currentCoords = coords.copy()
            geom = manage_xyz.np_to_xyz(self.geom,self.currentCoords)
            self.run(geom)
        return np.reshape(self.coup,(3*len(self.coup),1))*ANGSTROM_TO_AU