"""
Environments for training vehicles to reduce congestion in a merge.

This environment was used in:
"""

from flow.envs.base_env import Env
from flow.core import rewards

from gym.spaces.box import Box

import numpy as np
import collections

ADDITIONAL_ENV_PARAMS = {
    # maximum acceleration for autonomous vehicles, in m/s^2
    "max_accel": 3,
    # maximum deceleration for autonomous vehicles, in m/s^2
    "max_decel": 3,
    # desired velocity for all vehicles in the network, in m/s
    "target_velocity": 25,
    # maximum number of controllable vehicles in the network
    "num_rl": 5,  # TODO: get good number
}


class WaveAttenuationMergePOEnv(Env):
    """Partially observable merge environment.

    This environment is used to train autonomous vehicles to attenuate the
    formation and propagation of waves in an open merge network.

    Required from env_params:

    * max_accel: maximum acceleration for autonomous vehicles, in m/s^2
    * max_decel: maximum deceleration for autonomous vehicles, in m/s^2
    * target_velocity: desired velocity for all vehicles in the network, in m/s
    * num_rl: maximum number of controllable vehicles in the network

    States
        The observation consists of the speeds and bumper-to-bumper headways of
        the vehicles immediately preceding and following autonomous vehicle, as
        well as the ego speed of the autonomous vehicles.

        In order to maintain a fixed observation size, when the number of AVs
        in the network is less than "num_rl", the extra entries are filled in
        with zeros. Conversely, if the number of autonomous vehicles is greater
        than "num_rl", the observations from the additional vehicles are not
        included in the state space.

    Actions
        The action space consists of a vector of bounded accelerations for each
        autonomous vehicle $i$. In order to ensure safety, these actions are
        bounded by failsafes provided by the simulator at every time step.

        In order to account for variability in the number of autonomous
        vehicles, if n_AV < "num_rl" the additional actions provided by the
        agent are not assigned to any vehicle. Moreover, if n_AV > "num_rl",
        the additional vehicles are not provided with actions from the learning
        agent, and instead act as human-driven vehicles as well.

    Rewards
        The reward function encourages proximity of the system-level velocity
        to a desired velocity, while slightly penalizing small time headways
        among autonomous vehicles.

    Termination
        A rollout is terminated if the time horizon is reached or if two
        vehicles collide into one another.
    """

    def __init__(self, env_params, sumo_params, scenario):
        for p in ADDITIONAL_ENV_PARAMS.keys():
            if p not in env_params.additional_params:
                raise KeyError(
                    'Environment parameter "{}" not supplied'.format(p))

        # maximum number of controlled vehicles
        self.num_rl = env_params.additional_params["num_rl"]
        # queue of rl vehicles waiting to be controlled
        self.rl_queue = collections.deque()
        # names of the rl vehicles controlled at any step
        self.rl_veh = []
        # used for visualization
        self.leader = []
        self.follower = []
        self.name={}
        self.count=0
        self.communication = [0 for _ in range(self.num_rl)]

        super().__init__(env_params, sumo_params, scenario)

    @property
    def action_space(self):
        """See class definition."""
        # TODO: bind by 0->1, and denormalize based on what each
        # TODO: of acceleration, lane change, and communication want
        return Box(
            low=0,
            high=1,
            shape=(3,),
            dtype=np.float32)

    def denormalize(self, value, action_type):
        '''Helper function used to denormalized the 0-1 value of action space
        value: value between 0 and 1
        type: type of the action: 1 for acceleration, 2 for lane change and 3 for communication'''
        if action_type == 1:
            return np.array(value) * (self.env_params.additional_params["max_accel"] +
                                      abs(self.env_params.additional_params["max_decel"])) + \
                   -abs(self.env_params.additional_params["max_decel"])
        elif action_type == 2:
            value[np.where(value >= 2 / 3)] = 1
            value[np.where(value <= 1 / 3)] = -1
            value[np.where((value >= 1 / 3) & (value <= 2 / 3))] = 0
            return value
        elif action_type == 3:
            # TODO: define communication variable range
            return 0

    @property
    def observation_space(self):
        """See class definition."""
        return Box(low=0, high=1, shape=(6,), dtype=np.float32)

    def _apply_rl_actions(self, rl_actions):
        """See class definition."""
        rl_actions = np.array(rl_actions) / 2 + 0.5
        # ok this is the problem. self.vehicles.num_rl_vehicles is 20, which matches the size,
        # but rl_actions[:self.vehicles.num_rl_vehicles] somehow only has size 15
        # 20 rl vehicles but only 15 actions for all (not only acceleration, but all)??
        self.apply_acceleration(self.vehicles.get_rl_ids(),
                                self.denormalize(rl_actions[:self.vehicles.num_rl_vehicles], 1))
        self.apply_lane_change(self.vehicles.get_rl_ids(),
                               self.denormalize(rl_actions[self.num_rl: self.num_rl + self.vehicles.num_rl_vehicles],
                                                2))

        # self.communication = []
        # for i, rl_id in enumerate(self.rl_veh):
        # self.communication.append(self.denormalize(rl_actions[2 * self.num_rl + i], 3))

    def get_state(self, rl_id=None, **kwargs):
        """See class definition."""
        self.leader = []
        self.follower = []

        # normalizing constants
        max_speed = self.scenario.max_speed
        max_length = self.scenario.length

        observation = [0 for _ in range(7 * self.num_rl)]
        for i, rl_id in enumerate(self.rl_veh):
            this_speed = self.vehicles.get_speed(rl_id)
            lead_id = self.vehicles.get_leader(rl_id)
            follower = self.vehicles.get_follower(rl_id)

            if lead_id in ["", None]:
                # in case leader is not visible
                lead_speed = max_speed
                lead_head = max_length
            else:
                self.leader.append(lead_id)
                lead_speed = self.vehicles.get_speed(lead_id)
                lead_head = self.get_x_by_id(lead_id) \
                            - self.get_x_by_id(rl_id) - self.vehicles.get_length(rl_id)

            if follower in ["", None]:
                # in case follower is not visible
                follow_speed = 0
                follow_head = max_length
            else:
                self.follower.append(follower)
                follow_speed = self.vehicles.get_speed(follower)
                follow_head = self.vehicles.get_headway(follower)

            # TODO: try local communication (so maybe nearest neighbors in each lane in front and behind)
            comm = (np.sum(self.communication) - self.communication[i]) / (len(self.communication) - 1)

            observation[7 * i + 0] = this_speed / max_speed
            observation[7 * i + 1] = (lead_speed - this_speed) / max_speed
            observation[7 * i + 2] = lead_head / max_length
            observation[7 * i + 3] = (this_speed - follow_speed) / max_speed
            observation[7 * i + 4] = follow_head / max_length
            observation[7 * i + 5] = comm
            # TODO: try position, lane, lane leaders and followers, route (1, 2, 3, 4)
            if rl_id not in self.name:
                self.name[rl_id]=self.count
                self.count=self.count+1
            observation[7 * i + 6] = self.name[rl_id]  # self.only_numerics(rl_id) #
        return observation

    @staticmethod
    def only_numerics(seq):
        seq_type = type(seq)
        return seq_type().join(filter(seq_type.isdigit, seq))

    def compute_reward(self, rl_actions, **kwargs):
        """See class definition."""
        if kwargs["fail"]:
            return -5000

        return -1

    def sort_by_position(self):
        """See parent class.

        Sorting occurs by the ``get_x_by_id`` method instead of
        ``get_absolute_position``.
        """
        # vehicles are sorted by their get_x_by_id value
        sorted_ids = sorted(self.vehicles.get_ids(), key=self.get_x_by_id)
        return sorted_ids, None

    def additional_command(self):
        """See parent class.

        This method performs to auxiliary tasks:

        * Define which vehicles are observed for visualization purposes.
        * Maintains the "rl_veh" and "rl_queue" variables to ensure the RL
        vehicles that are represented in the state space does not change
        until one of the vehicles in the state space leaves the network.
        Then, the next vehicle in the queue is added to the state space and
        provided with actions from the policy.         """
        # add rl vehicles that just entered the network into the rl queue         
        for veh_id in self.vehicles.get_rl_ids():
            if veh_id not in list(self.rl_queue) + self.rl_veh:
                self.rl_queue.append(veh_id)

        # remove rl vehicles that exited the network
        for veh_id in list(self.rl_queue):
            if veh_id not in self.vehicles.get_rl_ids():
                self.rl_queue.remove(veh_id)
        for veh_id in self.rl_veh:
            if veh_id not in self.vehicles.get_rl_ids():
                self.rl_veh.remove(veh_id)

        # fil up rl_veh until they are enough controlled vehicles
        while len(self.rl_queue) > 0 and len(self.rl_veh) < self.num_rl:
            rl_id = self.rl_queue.popleft()
            self.rl_veh.append(rl_id)

        # specify observed vehicles
        for veh_id in self.leader + self.follower:
            self.vehicles.set_observed(veh_id)

    def reset(self):
        """See parent class.

        In addition, a few variables that are specific to this class are
        emptied before they are used by the new rollout.
        """
        self.leader = []
        self.follower = []
        return super().reset()
