# Copyright (c) 2020 fortiss GmbH
#
# Authors: Marco Oliva, Silvan Wimmer
#
# This work is licensed under the terms of the MIT license.
# For a copy, see <https://opensource.org/licenses/MIT>.

import numpy as np
import logging
import tensorflow as tf
from gym import spaces
from typing import Dict

from bark.core.models.dynamic import StateDefinition
from bark.core.geometry import Distance, Point2d
from bark.runtime.commons.parameters import ParameterServer

from bark_ml.observers.observer import StateObserver

class GraphObserver(StateObserver):
  """
  This observer converts an `ObservsedWorld` instance into
  a graph-structured observation, consisting of nodes (with
  features), an adjacency matrix expressing the (directed)
  connections between nodes in the graph, and edges (with
  features), expressing the relative information between
  the agents that constitute these connections.
  """

  def __init__(self, params=ParameterServer()):
    """
    Creates an instance of `GraphObserver`.

    Args:
      params: A `ParameterServer` instance that enables further
        configuration of the observer. Defaults to a new instance.
    """
    StateObserver.__init__(self, params)
    self._logger = logging.getLogger()

    self._normalize_observations = \
      self._params["ML"]["GraphObserver"]["NormalizationEnabled",
      "Whether normalization of features should be performed.",
      True]

    self._num_agents =\
      params["ML"]["GraphObserver"]["AgentLimit", 
      "The maximum number of agents that is observed.",
      4]
    self._visibility_radius =\
      params["ML"]["GraphObserver"]["VisibilityRadius",
      "The radius in which agent can 'see', i.e. detect other agents.", 
      550]
    self._add_self_loops =\
      params["ML"]["GraphObserver"]["SelfLoops", 
      "Whether nodes have self connections.", 
      True]
    requested_node_attribute_keys =\
      params["ML"]["GraphObserver"]["EnabledNodeFeatures",
      "The list of available node features, given by their string key that \
       the observer should extract from the world and insert into the \
       observation. For a list of available features, refer to the list \
       returned by `GraphObserver.available_node_attributes`.",
      self.available_node_attributes()
    ]
    self._enabled_node_attribute_keys = self._filter_requested_attributes(
      requested_keys=requested_node_attribute_keys,
      available_keys=self.available_node_attributes(),
      context="node")
    self._logger.info(
      "GraphObserver configured with node attributes: " +
      f"{self._enabled_node_attribute_keys}")
    requested_edge_attribute_keys =\
      params["ML"]["GraphObserver"]["EnabledEdgeFeatures",
      "The list of available edge features, given by their string key that \
       the observer should extract from the world and insert into the \
       observation. For a list of available features, refer to the list \
       returned by `GraphObserver.available_edge_attributes`.",
      self.available_edge_attributes()
    ]
    self._enabled_edge_attribute_keys = self._filter_requested_attributes(
      requested_keys=requested_edge_attribute_keys,
      available_keys=self.available_edge_attributes(),
      context="edge")
    self._logger.info(
      "GraphObserver configured with edge attributes: " +
      f"{self._enabled_edge_attribute_keys}")

    # the number of features of a node in the graph
    self.feature_len = len(self._enabled_node_attribute_keys)
    # the number of features of an edge between two nodes
    self.edge_feature_len = len(self._enabled_edge_attribute_keys)

  def Observe(self, world):
    """see base class"""
    agents = self._preprocess_agents(world)

    # placeholder for the output observation
    obs = np.zeros(self._len_state)

    # insert node features for each agent
    for i, agent in agents:
      start_index = i * self.feature_len
      end_index = start_index + self.feature_len
      obs[start_index:end_index] = self._extract_node_features(agent)

    edge_features = np.zeros((self._num_agents, self._num_agents, self.edge_feature_len))
    adjacency_matrix = np.zeros((self._num_agents, self._num_agents))

    # add edges to all visible agents
    for index, agent in agents:
      nearby_agents = self._nearby_agents(
        center_agent=agent, 
        agents=agents, 
        radius=self._visibility_radius)

      if self._add_self_loops:
        adjacency_matrix[index, index] = 1

      for target_index, nearby_agent in nearby_agents:
        edge_features[index, target_index, :] = self._extract_edge_features(agent, nearby_agent)
        adjacency_matrix[index, target_index] = 1
        adjacency_matrix[target_index, index] = 1
    
    # insert adjacency list
    adj_start_index = self._num_agents * self.feature_len
    adj_end_index = adj_start_index + self._num_agents ** 2
    obs[adj_start_index:adj_end_index] = adjacency_matrix.reshape(-1)

    # insert edge features
    obs[adj_end_index:] = edge_features.reshape(-1)
    
    return tf.convert_to_tensor(obs, dtype=tf.float32, name='observation')

  @classmethod
  def graph(cls, observations, graph_dims, dense=False):
    """
    Maps the given batch of observations into a
    graph representation.

    Args:
    observations: The batch of observations as a tf.Tensor
      of shape (batch_size, observation_size).
    graph_dims: A tuple containing the dimensions of the 
      graph as (num_nodes, num_features, num_edge_features).
      If `dense` is set to True, num_edge_features is ignored.
    dense: Specifies the format of the returned graph representation. 
      If set to `True`, the edges are returned as a list of pairs 
      of nodes indices (relative to the flattened batch) and an additional
      mapping of each node to a graph is returned. If set to 
      False (default), edges are returned as sparse adjacency 
      matrix and an edge feature matrix is additionally returned.

    Returns:
      X: Node features of (batch_size, num_nodes, num_features)

      If `dense` is True:
      A: Dense representation of edges as a list of node index pairs,
        shape (num_total_edges, 2).
      node_to_graph_map: A 1d list, where each element is the mapping
      of the node (indexed relative to the batch) to a graph.
      
      If `dense`is False:
      A: Spare binary adjacency matrix of shape 
        (batch_size num_nodes, num_nodes).
      E: Edge features of shape 
        (batch_size, num_nodes, num_edge_features, num_edge_features).
    """
    obs = observations # for brevity

    if not tf.is_tensor(obs):
      obs = tf.convert_to_tensor(obs)
    
    n_nodes, n_features = graph_dims[0:2]
    batch_size = tf.shape(observations)[0]
    
    # extract node features F
    F = tf.reshape(obs[:, :n_nodes * n_features], [batch_size, n_nodes, n_features])

    # extract adjacency matrix A
    adj_start_idx = n_nodes * n_features
    adj_end_idx = adj_start_idx + n_nodes ** 2
    A = tf.reshape(obs[:, adj_start_idx:adj_end_idx], [batch_size, n_nodes, n_nodes])

  
    if dense:
      # in dense mode, the nodes of all graphs are 
      # concatenated into one list that contains all 
      # feature vectors of the whole batch.
      # the assignment to graphs is handled by the
      # node_to_graph_map.
      F = tf.reshape(F, [batch_size * n_nodes, n_features])

      # find non-zero elements in the adjacency matrix (edges)
      # and collect their indices
      A = tf.where(tf.greater(A, 0))

      # we need the indices of the source and target nodes to
      # be represented as their indices in the whole batch,
      # in other words: each node index must be the index
      # of the graph in the batch plus the index of the node
      # in the graph. E.g. if each graph has 5 nodes, the 
      # node indices are: graph 0: 0-4, graph 1: 5-9, etc.

      # compute a tensor where each element
      # is the graph index of the node that is represented
      # by the same index in A
      graph_indices = tf.reshape(A[:, 0], [1, -1])
      graph_indices = tf.scalar_mul(n_nodes, graph_indices)
      graph_indices = tf.transpose(tf.tile(graph_indices, [2, 1]))
      
      A = A[:, 1:] + graph_indices
      A = tf.cast(A, tf.int32)

      # since the node feature vectors are all stored in a 2D array,
      # we need to map them to the graphs that they belong to.

      # construct a mapping from the node feature vectors in F
      # to the graphs in the batch, i.e a list where each element 
      # contains the index of the graph that the feature vector at
      # the same index in F belongs to, e.g. if each graph contains
      # two nodes, then the node_to_graph_map is [0, 0, 1, 1, 2, 2, ...]
      node_to_graph_map = tf.reshape(tf.range(batch_size), [-1, 1])
      node_to_graph_map = tf.tile(node_to_graph_map, [1, n_nodes])
      node_to_graph_map = tf.reshape(node_to_graph_map, [-1])
      
      # extract edge features
      n_edge_features = graph_dims[2]
      E_shape = [-1, n_edge_features]
      E = tf.reshape(obs[:, adj_end_idx:], E_shape)
      return F, A, node_to_graph_map, E

    # extract edge features
    n_edge_features = graph_dims[2]
    E_shape = [batch_size, n_nodes, n_nodes, n_edge_features]
    E = tf.reshape(obs[:, adj_end_idx:], E_shape)
    return F, A, E

  def _preprocess_agents(self, world):
    """
    Preproccesses the agents for constructing an
    observation by resorting.

    Args:
      world: A `bark.core.ObservedWorld` instance.

    Returns:
      A list of tuples, consisting of an index and an 
      agent object element.

      The first element always represents the ego agent.
      The remaining elements resemble other agents, up 
      to the limit defined by `self._num_agents`, sorted 
      in ascending order with respect to the agents'
      distance in the world to the ego agent.
    """
    ego_agent = world.ego_agent
    agents = list(world.agents.values())
    agents.remove(ego_agent)
    agents = self._agents_sorted_by_distance(ego_agent, agents)
    agents.insert(0, ego_agent)
    return list(enumerate(agents))[:self._num_agents]

  def _agents_sorted_by_distance(self, ego_agent, agents):
    """
    Returns the given list of `agents`, sorted in ascending
    order by their relative distance to the `ego_agent`.
    """
    def distance(agent):
      return Distance(
        self._position(ego_agent), 
        self._position(agent))
    if len(agents) > 2:
      agents.sort(key=distance)
    return agents

  def _nearby_agents(self, center_agent, agents, radius: float):
    """
    Returns all elements from 'agents' whose position, defined 
    in x and y coordinates is within the specified `radius` of 
    the 'center_agent's position.
    """
    center_agent_pos = self._position(center_agent)
    other_agents = filter(lambda a: a[1].id != center_agent.id, agents)
    nearby_agents = []
    for index, agent in other_agents:
      agent_pos = self._position(agent)
      distance = Distance(center_agent_pos, agent_pos)
      if distance <= radius:
        nearby_agents.append((index, agent))
    return nearby_agents

  def _extract_node_features(self, agent, as_dict=False):
    res = {}

    # Init data (to keep ordering always equal for reading and writing)
    for label in self._enabled_node_attribute_keys:
      res[label] = "inf"
    
    n = self.normalization_data
    def add_feature(key, value, norm_range):
      if self._normalize_observations:
        value = self._normalize_value(value, norm_range)
      res[key] = value

    try:
      state = agent.state
      if "x" in res:
        add_feature("x", state[int(StateDefinition.X_POSITION)], n["x"])
      if "y" in res:
        add_feature("y", state[int(StateDefinition.Y_POSITION)], n["y"])
      if "theta" in res:
        add_feature("theta", state[int(StateDefinition.THETA_POSITION)], n["theta"])
      if "vel" in res:
        add_feature("vel", state[int(StateDefinition.VEL_POSITION)], n["vel"])

      # get information related to goal
      goal_center = agent.goal_definition.goal_shape.center[0:2]
      goal_dx = goal_center[0] - state[int(StateDefinition.X_POSITION)]
      goal_dy = goal_center[1] - state[int(StateDefinition.Y_POSITION)]

      if "goal_x" in res:
        add_feature("goal_x", goal_center[0], n["x"])
      if "goal_y" in res:
        add_feature("goal_y", goal_center[1], n["y"])
      if "goal_dx" in res:
        add_feature("goal_dx", goal_dx, n["dx"])
      if "goal_dy" in res:
        add_feature("goal_dy", goal_dy, n["dy"])
      if "goal_theta" in res:
        add_feature("goal_theta", np.arctan2(goal_dy, goal_dx), n["theta"])
      if "goal_d" in res:
        add_feature("goal_d", np.sqrt(goal_dx**2 + goal_dy**2), n["distance"])
      if "goal_vel" in res:
        goal_velocity = np.mean(agent.goal_definition.velocity_range)
        add_feature("goal_vel", goal_velocity, n["vel"])
    except:
      raise AttributeError(
        "A problem occured during node feature extraction. Possibly " +
        "a node feature that's specifed in the 'EnabledNodeFeatures' " +
        "parameter is not supported by the current BARK-ML environment.")
    
    # remove disabled attributes
    res = {key: res[key] for key in self._enabled_node_attribute_keys}
    assert list(res.keys()) == self._enabled_node_attribute_keys

    if not as_dict:
      res = list(res.values())

    return res

  def _extract_edge_features(self, source_agent, target_agent):
    """
    Encodes the relation between the given agents as the
    features of a directed edge in the graph.

    Args:
      source_agent: the agent that defines the source of the edge.
      target_agent: the agent that defines the target of the edge.

    Returns:
      An np.array containing the differences in the agents' 
      x and y position, velocities and orientations.
    """
    source_features =\
      self._extract_node_features(source_agent, as_dict=True)
    target_features =\
      self._extract_node_features(target_agent, as_dict=True)

    features = {
      "dx": source_features["x"] - target_features["x"],
      "dy": source_features["y"] - target_features["y"],
      "dvel": source_features["vel"] - target_features["vel"],
      "dtheta": source_features["theta"] - target_features["theta"]
    }

    features = {key: features[key] for key in self._enabled_edge_attribute_keys}
    features = list(features.values())
    assert len(features) == len(self._enabled_edge_attribute_keys)

    return features

  def _filter_requested_attributes(self, 
                                   requested_keys, 
                                   available_keys, 
                                   context=""):
    """
    Filters the requested node/edge attributes by the available 
    attributes and returns the intersection of the requested
    and available attributes.
    """
    if not isinstance(requested_keys, list):
      raise ValueError(
        f"Requested {context} feature list must be a list of " +
        f"strings, not {type(requested_keys)}.")

    if len(requested_keys) == 0:
      self._logger.warning(
        "The `GraphObserver` received an empty list of requested " +
        f"{context} attributes.")
      return []

    valid_keys = []
    invalid_keys = []

    for key in requested_keys:
      if key in available_keys:
        valid_keys.append(key)
      else:
        invalid_keys.append(key)
    
    if len(invalid_keys) > 0:
      self._logger.warning(
        f"The following {context} attributes requested from the GraphObserver " +
        f"are not supported and will be ignored: {invalid_keys}")

    return valid_keys

  def _normalize_value(self, value, range):
    """
    Normalizes the `value` with can take on values in
    into the range of [-1, 1].
    If the `value` is outside the given range, it's clamped 
    to the bound of [-1, 1]
    """
    normed = (value - range[0]) / (range[1] - range[0])
    # normed = max(-1, normed) # values lower -1 clipped
    # normed = min(1, normed) # values bigger 1 clipped
    return normed

  def reset(self, world):
    return world

  def _position(self, agent) -> Point2d:
    return Point2d(
      agent.state[int(StateDefinition.X_POSITION)],
      agent.state[int(StateDefinition.Y_POSITION)]
    )

  @property
  def normalization_data(self) -> Dict[str, list]:
    """
    The reference ranges of how certain attributes are normalized.
    Use this info to scale normalized values back to real values.
    
    E.g. the value for key 'x' returns the possible range of 
    the x-element of positions. A normalized value of 1.0 
    corresponds to the maximum of this range (and vice versa).

    Note: This dictionary does not include a range for each
    possible attribute, but rather for each _kind_ of attribute,
    like distances, velocities, angles, etc.
    E.g. all distances (between agents, between objects, etc.)
    are scaled relative to the 'distance' range. 
    """ 
    x_range = self._world_x_range[1] - self._world_x_range[0]
    y_range = self._world_y_range[1] - self._world_y_range[0]
    max_dist = np.linalg.norm([x_range, y_range])

    d = {}
    d['x'] = self._world_x_range
    d['y'] = self._world_y_range
    d['theta'] = self._theta_range
    d['vel'] = self._velocity_range
    d['distance'] = [0, max_dist]
    d['dx'] = [-x_range, x_range]
    d['dy'] = [-y_range, y_range]
    return d

  def sample(self):
    raise NotImplementedError

  @classmethod
  def available_node_attributes(cls, with_descriptions=False):
    attributes = {
      "x": "The x-components of the agent's position.",
      "y": "The y-components of the agent's position.", 
      "theta": "The current heading angle of tha agent.",
      "vel": "The current velocity of the agent.", 
      "goal_x": "The x-component of the goal's position.", 
      "goal_y": "The y-component of the goal's position.",
      "goal_dx": "The difference in the x-component of the agent's and the goal's position.", 
      "goal_dy": "The difference in the y-component of the agent's and the goal's position.", 
      "goal_theta": "The goal heading angle.", 
      "goal_d": "The euclidian distance of the agent to the goal.", 
      "goal_vel": "The goal velocity."
    }

    if not with_descriptions:
      attributes = list(attributes.keys())

    return attributes

  @classmethod
  def available_edge_attributes(cls, with_descriptions=False):
    attributes = {
      "dx": "The difference in the x-position of the two agents.", 
      "dy": "The difference in the y-position of the two agents.", 
      "dvel": "The difference in the velocity of the two agents.", 
      "dtheta": "The difference in the heading angle of the two agents."
    }

    if not with_descriptions:
      attributes = list(attributes.keys())

    return attributes

  @property
  def observation_space(self):    
    # 0 ... 1   for all node attributes
    # 0 ... 1   for the adjacency list
    # 0 ... 1   for the edge attributes
    return spaces.Box(
      low=np.concatenate((
        np.zeros(self._num_agents*self.feature_len),
        np.zeros(self._num_agents**2),
        np.zeros(self._num_agents**2*self.edge_feature_len))),
      high=np.ones(self._len_state))

  @property
  def _len_state(self):
    len_node_features = self._num_agents*self.feature_len
    len_adjacency = self._num_agents**2
    len_edge_features = len_adjacency*self.edge_feature_len
    return len_node_features + len_adjacency + len_edge_features

  @property
  def graph_dimensions(self):
    """
    Returns a three-element tuple that defines the
    dimensions of the graph as returned by the
    `Observe` method, consisting of
    - the number of nodes
    - the number of node features
    - the number of edge features
    """
    return (
      self._num_agents,
      self.feature_len,
      self.edge_feature_len
    )