# Copyright 2019 The dm_control Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Soccer observables modules."""

import abc

from dm_control.composer.observation import observable as base_observable
from dm_control.locomotion.soccer import team as team_lib
import numpy as np


class ObservablesAdder(metaclass=abc.ABCMeta):
  """A callable that adds a set of per-player observables for a task."""

  @abc.abstractmethod
  def __call__(self, task, player):
    """Adds observables to a player for the given task.

    Args:
      task: A `soccer.Task` instance.
      player: A `Walker` instance to which observables will be added.
    """


class MultiObservablesAdder(ObservablesAdder):
  """Applies multiple `ObservablesAdder`s to a soccer task and player."""

  def __init__(self, observables):
    """Initializes a `MultiObservablesAdder` instance.

    Args:
      observables: A list of `ObservablesAdder` instances.
    """
    self._observables = observables

  def __call__(self, task, player):
    """Adds observables to a player for the given task.

    Args:
      task: A `soccer.Task` instance.
      player: A `Walker` instance to which observables will be added.
    """
    for observable in self._observables:
      observable(task, player)


class CoreObservablesAdder(ObservablesAdder):
  """Core set of per player observables."""

  def __call__(self, task, player):
    """Adds observables to a player for the given task.

    Args:
      task: A `soccer.Task` instance.
      player: A `Walker` instance to which observables will be added.
    """
    # Enable proprioceptive observables.
    self._add_player_proprio_observables(player)

    # Add egocentric observations of soccer ball.
    self._add_player_observables_on_ball(player, task.ball)

    # Add egocentric observations of others.
    teammate_id = 0
    opponent_id = 0
    for other in task.players:
      if other is player:
        continue
      # Infer team prefix for `other` conditioned on `player.team`.
      if player.team != other.team:
        prefix = 'opponent_{}'.format(opponent_id)
        opponent_id += 1
      else:
        prefix = 'teammate_{}'.format(teammate_id)
        teammate_id += 1

      self._add_player_observables_on_other(player, other, prefix)

    self._add_player_arena_observables(player, task.arena)

    # Add per player game statistics.
    self._add_player_stats_observables(task, player)

  def _add_player_observables_on_other(self, player, other, prefix):
    """Add observables of another player in this player's egocentric frame.

    Args:
      player: A `Walker` instance, the player we are adding observables to.
      other: A `Walker` instance corresponding to a different player.
      prefix: A string specifying a prefix to apply to the names of observables
        belonging to `player`.
    """
    if player is other:
      raise ValueError('Cannot add egocentric observables of player on itself.')

    sensors = []
    for effector in other.walker.end_effectors:
      name = effector.name + '_' + prefix + '_end_effector'
      sensors.append(player.walker.mjcf_model.sensor.add(
          'framepos', name=name,
          objtype=effector.tag, objname=effector,
          reftype='body', refname=player.walker.root_body))
    def _egocentric_end_effectors_xpos(physics):
      return np.reshape(physics.bind(sensors).sensordata, -1)
    # Adds end effectors of the other agents in the player's egocentric frame.
    name = '{}_ego_end_effectors_pos'.format(prefix)
    player.walker.obs_on_other[name] = sensors
    player.walker.observables.add_observable(
        name,
        base_observable.Generic(_egocentric_end_effectors_xpos))

    ego_linvel_name = '{}_ego_linear_velocity'.format(prefix)
    ego_linvel_sensor = player.walker.mjcf_model.sensor.add(
        'framelinvel', name=ego_linvel_name,
        objtype='body', objname=other.walker.root_body,
        reftype='body', refname=player.walker.root_body)
    player.walker.obs_on_other[ego_linvel_name] = [ego_linvel_sensor]
    player.walker.observables.add_observable(
        ego_linvel_name,
        base_observable.MJCFFeature('sensordata', ego_linvel_sensor))

    ego_pos_name = '{}_ego_position'.format(prefix)
    ego_pos_sensor = player.walker.mjcf_model.sensor.add(
        'framepos', name=ego_pos_name,
        objtype='body', objname=other.walker.root_body,
        reftype='body', refname=player.walker.root_body)
    player.walker.obs_on_other[ego_pos_name] = [ego_pos_sensor]
    player.walker.observables.add_observable(
        ego_pos_name,
        base_observable.MJCFFeature('sensordata', ego_pos_sensor))

    sensors_rot = []
    obsname = '{}_ego_orientation'.format(prefix)
    for direction in ['x', 'y', 'z']:
      sensorname = obsname + '_' + direction
      sensors_rot.append(player.walker.mjcf_model.sensor.add(
          'frame'+direction+'axis', name=sensorname,
          objtype='body', objname=other.walker.root_body,
          reftype='body', refname=player.walker.root_body))
    def _egocentric_orientation(physics):
      return np.reshape(physics.bind(sensors_rot).sensordata, -1)
    player.walker.obs_on_other[obsname] = sensors_rot
    player.walker.observables.add_observable(
        obsname,
        base_observable.Generic(_egocentric_orientation))

    # Adds end effectors of the other agents in the other's egocentric frame.
    # A is seeing B's hand extended to B's right.
    player.walker.observables.add_observable(
        '{}_end_effectors_pos'.format(prefix),
        other.walker.observables.end_effectors_pos)

  def _add_player_observables_on_ball(self, player, ball):
    """Add observables of the soccer ball in this player's egocentric frame.

    Args:
      player: A `Walker` instance, the player we are adding observations for.
      ball: A `SoccerBall` instance.
    """
    # Add egocentric ball observations.
    player.walker.ball_ego_angvel_sensor = player.walker.mjcf_model.sensor.add(
        'frameangvel', name='ball_ego_angvel',
        objtype='body', objname=ball.root_body,
        reftype='body', refname=player.walker.root_body)
    player.walker.observables.add_observable(
        'ball_ego_angular_velocity',
        base_observable.MJCFFeature('sensordata',
                                    player.walker.ball_ego_angvel_sensor))

    player.walker.ball_ego_pos_sensor = player.walker.mjcf_model.sensor.add(
        'framepos', name='ball_ego_pos',
        objtype='body', objname=ball.root_body,
        reftype='body', refname=player.walker.root_body)
    player.walker.observables.add_observable(
        'ball_ego_position',
        base_observable.MJCFFeature('sensordata',
                                    player.walker.ball_ego_pos_sensor))

    player.walker.ball_ego_linvel_sensor = player.walker.mjcf_model.sensor.add(
        'framelinvel', name='ball_ego_linvel',
        objtype='body', objname=ball.root_body,
        reftype='body', refname=player.walker.root_body)
    player.walker.observables.add_observable(
        'ball_ego_linear_velocity',
        base_observable.MJCFFeature('sensordata',
                                    player.walker.ball_ego_linvel_sensor))

  def _add_player_proprio_observables(self, player):
    """Add proprioceptive observables to the given player.

    Args:
      player: A `Walker` instance, the player we are adding observations for.
    """
    for observable in (player.walker.observables.proprioception +
                       player.walker.observables.kinematic_sensors):
      observable.enabled = True

    # Also enable previous action observable as part of proprioception.
    player.walker.observables.prev_action.enabled = True

  def _add_player_arena_observables(self, player, arena):
    """Add observables of the arena.

    Args:
      player: A `Walker` instance to which observables will be added.
      arena: A `Pitch` instance.
    """
    # Enable egocentric view of position detectors (goal, field).
    # Corners named according to walker *facing towards opponent goal*.
    clockwise_names = [
        'team_goal_back_right',
        'team_goal_mid',
        'team_goal_front_left',
        'field_front_left',
        'opponent_goal_back_left',
        'opponent_goal_mid',
        'opponent_goal_front_right',
        'field_back_right',
    ]
    clockwise_features = [
        lambda _: arena.home_goal.lower[:2],
        lambda _: arena.home_goal.mid,
        lambda _: arena.home_goal.upper[:2],
        lambda _: arena.field.upper,
        lambda _: arena.away_goal.upper[:2],
        lambda _: arena.away_goal.mid,
        lambda _: arena.away_goal.lower[:2],
        lambda _: arena.field.lower,
    ]
    xpos_xyz_callable = lambda p: p.bind(player.walker.root_body).xpos
    xpos_xy_callable = lambda p: p.bind(player.walker.root_body).xpos[:2]
    # A list of egocentric reference origin for each one of clockwise_features.
    clockwise_origins = [
        xpos_xy_callable,
        xpos_xyz_callable,
        xpos_xy_callable,
        xpos_xy_callable,
        xpos_xy_callable,
        xpos_xyz_callable,
        xpos_xy_callable,
        xpos_xy_callable,
    ]
    if player.team != team_lib.Team.HOME:
      half = len(clockwise_features) // 2
      clockwise_features = clockwise_features[half:] + clockwise_features[:half]
      clockwise_origins = clockwise_origins[half:] + clockwise_origins[:half]

    for name, feature, origin in zip(clockwise_names, clockwise_features,
                                     clockwise_origins):
      player.walker.observables.add_egocentric_vector(
          name, base_observable.Generic(feature), origin_callable=origin)

  def _add_player_stats_observables(self, task, player):
    """Add observables corresponding to game statistics.

    Args:
      task: A `soccer.Task` instance.
      player: A `Walker` instance to which observables will be added.
    """

    def _stats_vel_to_ball(physics):
      dir_ = (
          physics.bind(task.ball.geom).xpos -
          physics.bind(player.walker.root_body).xpos)
      vel_to_ball = np.dot(dir_[:2] / (np.linalg.norm(dir_[:2]) + 1e-7),
                           physics.bind(player.walker.root_body).cvel[3:5])
      return np.sum(vel_to_ball)

    player.walker.observables.add_observable(
        'stats_vel_to_ball', base_observable.Generic(_stats_vel_to_ball))

    def _stats_closest_vel_to_ball(physics):
      """Velocity to the ball if this walker is the team's closest."""
      closest = None
      min_team_dist_to_ball = np.inf
      for player_ in task.players:
        if player_.team == player.team:
          dist_to_ball = np.linalg.norm(
              physics.bind(task.ball.geom).xpos -
              physics.bind(player_.walker.root_body).xpos)
          if dist_to_ball < min_team_dist_to_ball:
            min_team_dist_to_ball = dist_to_ball
            closest = player_
      if closest is player:
        return _stats_vel_to_ball(physics)
      return 0.

    player.walker.observables.add_observable(
        'stats_closest_vel_to_ball',
        base_observable.Generic(_stats_closest_vel_to_ball))

    def _stats_veloc_forward(physics):
      """Player's forward velocity."""
      return player.walker.observables.veloc_forward(physics)

    player.walker.observables.add_observable(
        'stats_veloc_forward', base_observable.Generic(_stats_veloc_forward))

    def _stats_vel_ball_to_goal(physics):
      """Ball velocity towards opponents' goal."""
      if player.team == team_lib.Team.HOME:
        goal = task.arena.away_goal
      else:
        goal = task.arena.home_goal

      goal_center = (goal.upper + goal.lower) / 2.
      direction = goal_center - physics.bind(task.ball.geom).xpos
      ball_vel_observable = task.ball.observables.linear_velocity
      ball_vel = ball_vel_observable.observation_callable(physics)()

      norm_dir = np.linalg.norm(direction)
      normalized_dir = direction / norm_dir if norm_dir else direction
      return np.sum(np.dot(normalized_dir, ball_vel))

    player.walker.observables.add_observable(
        'stats_vel_ball_to_goal',
        base_observable.Generic(_stats_vel_ball_to_goal))

    def _stats_avg_teammate_dist(physics):
      """Compute average distance from `walker` to its teammates."""
      teammate_dists = []
      for other in task.players:
        if player is other:
          continue
        if other.team != player.team:
          continue
        dist = np.linalg.norm(
            physics.bind(player.walker.root_body).xpos -
            physics.bind(other.walker.root_body).xpos)
        teammate_dists.append(dist)
      return np.mean(teammate_dists) if teammate_dists else 0.

    player.walker.observables.add_observable(
        'stats_home_avg_teammate_dist',
        base_observable.Generic(_stats_avg_teammate_dist))

    def _stats_teammate_spread_out(physics):
      """Compute average distance from `walker` to its teammates."""
      return _stats_avg_teammate_dist(physics) > 5.

    player.walker.observables.add_observable(
        'stats_teammate_spread_out',
        base_observable.Generic(_stats_teammate_spread_out))

    def _stats_home_score(unused_physics):
      if (task.arena.detected_goal() and
          task.arena.detected_goal() == player.team):
        return 1.
      return 0.

    player.walker.observables.add_observable(
        'stats_home_score', base_observable.Generic(_stats_home_score))

    has_opponent = any([p.team != player.team for p in task.players])

    def _stats_away_score(unused_physics):
      if (has_opponent and task.arena.detected_goal() and
          task.arena.detected_goal() != player.team):
        return 1.
      return 0.

    player.walker.observables.add_observable(
        'stats_away_score', base_observable.Generic(_stats_away_score))


# TODO(b/124848293): add unit-test interception observables.
class InterceptionObservablesAdder(ObservablesAdder):
  """Adds obervables representing interception events.

  These observables represent events where this player received the ball from
  another player, or when an opponent intercepted the ball from this player's
  team. For each type of event there are three different thresholds applied to
  the distance travelled by the ball since it last made contact with a player
  (5, 10, or 15 meters).

  For example, on a given timestep `stats_i_received_ball_10m` will be 1 if
  * This player just made contact with the ball
  * The last player to have made contact with the ball was a different player
  * The ball travelled for at least 10 m since it last hit a player
  and 0 otherwise.

  Conversely, `stats_opponent_intercepted_ball_10m` will be 1 if:
  * An opponent just made contact with the ball
  * The last player to have made contact with the ball was on this player's team
  * The ball travelled for at least 10 m since it last hit a player
  """

  def __call__(self, task, player):
    """Adds observables to a player for the given task.

    Args:
      task: A `soccer.Task` instance.
      player: A `Walker` instance to which observables will be added.
    """

    def _stats_i_received_ball(unused_physics):
      if (task.ball.hit and task.ball.repossessed and
          task.ball.last_hit is player):
        return 1.
      return 0.

    player.walker.observables.add_observable(
        'stats_i_received_ball',
        base_observable.Generic(_stats_i_received_ball))

    def _stats_opponent_intercepted_ball(unused_physics):
      """Indicator on if an opponent intercepted the ball."""
      if (task.ball.hit and task.ball.intercepted and
          task.ball.last_hit.team != player.team):
        return 1.
      return 0.

    player.walker.observables.add_observable(
        'stats_opponent_intercepted_ball',
        base_observable.Generic(_stats_opponent_intercepted_ball))

    for dist in [5, 10, 15]:

      def _stats_i_received_ball_dist(physics, dist=dist):
        if (_stats_i_received_ball(physics) and
            task.ball.dist_between_last_hits is not None and
            task.ball.dist_between_last_hits > dist):
          return 1.
        return 0.

      player.walker.observables.add_observable(
          'stats_i_received_ball_%dm' % dist,
          base_observable.Generic(_stats_i_received_ball_dist))

      def _stats_opponent_intercepted_ball_dist(physics, dist=dist):
        if (_stats_opponent_intercepted_ball(physics) and
            task.ball.dist_between_last_hits is not None and
            task.ball.dist_between_last_hits > dist):
          return 1.
        return 0.

      player.walker.observables.add_observable(
          'stats_opponent_intercepted_ball_%dm' % dist,
          base_observable.Generic(_stats_opponent_intercepted_ball_dist))
