# 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.
# ============================================================================

"""Multi-agent MuJoCo soccer environment."""

import enum

from dm_control import composer
from dm_control.locomotion import walkers
from dm_control.locomotion.soccer.boxhead import BoxHead
from dm_control.locomotion.soccer.humanoid import Humanoid
from dm_control.locomotion.soccer.initializers import Initializer
from dm_control.locomotion.soccer.initializers import UniformInitializer
from dm_control.locomotion.soccer.observables import CoreObservablesAdder
from dm_control.locomotion.soccer.observables import InterceptionObservablesAdder
from dm_control.locomotion.soccer.observables import MultiObservablesAdder
from dm_control.locomotion.soccer.observables import ObservablesAdder
from dm_control.locomotion.soccer.pitch import MINI_FOOTBALL_GOAL_SIZE
from dm_control.locomotion.soccer.pitch import MINI_FOOTBALL_MAX_AREA_PER_HUMANOID
from dm_control.locomotion.soccer.pitch import MINI_FOOTBALL_MIN_AREA_PER_HUMANOID
from dm_control.locomotion.soccer.pitch import Pitch
from dm_control.locomotion.soccer.pitch import RandomizedPitch
from dm_control.locomotion.soccer.soccer_ball import regulation_soccer_ball
from dm_control.locomotion.soccer.soccer_ball import SoccerBall
from dm_control.locomotion.soccer.task import MultiturnTask
from dm_control.locomotion.soccer.task import Task
from dm_control.locomotion.soccer.team import Player
from dm_control.locomotion.soccer.team import RGBA_BLUE
from dm_control.locomotion.soccer.team import RGBA_RED
from dm_control.locomotion.soccer.team import Team
from dm_control.locomotion.walkers.initializers import mocap
import numpy as np


class WalkerType(enum.Enum):
  BOXHEAD = 0
  ANT = 1
  HUMANOID = 2


def _make_walker(name, walker_id, marker_rgba, walker_type=WalkerType.BOXHEAD):
  """Construct a BoxHead walker."""
  if walker_type == WalkerType.BOXHEAD:
    return BoxHead(
        name=name,
        walker_id=walker_id,
        marker_rgba=marker_rgba,
    )
  if walker_type == WalkerType.ANT:
    return walkers.Ant(name=name, marker_rgba=marker_rgba)
  if walker_type == WalkerType.HUMANOID:
    initializer = mocap.CMUMocapInitializer()
    return Humanoid(
        name=name,
        marker_rgba=marker_rgba,
        walker_id=walker_id,
        visual=Humanoid.Visual.JERSEY,
        initializer=initializer)
  raise ValueError("Unrecognized walker type: %s" % walker_type)


def _make_players(team_size, walker_type):
  """Construct home and away teams each of `team_size` players."""
  home_players = []
  away_players = []
  for i in range(team_size):
    home_walker = _make_walker("home%d" % i, i, RGBA_BLUE, walker_type)
    home_players.append(Player(Team.HOME, home_walker))

    away_walker = _make_walker("away%d" % i, i, RGBA_RED, walker_type)
    away_players.append(Player(Team.AWAY, away_walker))
  return home_players + away_players


def _area_to_size(area, aspect_ratio=0.75):
  """Convert from area and aspect_ratio to (width, height)."""
  return np.sqrt([area / aspect_ratio, area * aspect_ratio]) / 2.


def load(team_size,
         time_limit=45.,
         random_state=None,
         disable_walker_contacts=False,
         enable_field_box=False,
         keep_aspect_ratio=False,
         terminate_on_goal=True,
         walker_type=WalkerType.BOXHEAD):
  """Construct `team_size`-vs-`team_size` soccer environment.

  Args:
    team_size: Integer, the number of players per team. Must be between 1 and
      11.
    time_limit: Float, the maximum duration of each episode in seconds.
    random_state: (optional) an int seed or `np.random.RandomState` instance.
    disable_walker_contacts: (optional) if `True`, disable physical contacts
      between walkers.
    enable_field_box: (optional) if `True`, enable physical bounding box for
      the soccer ball (but not the players).
    keep_aspect_ratio: (optional) if `True`, maintain constant pitch aspect
      ratio.
    terminate_on_goal: (optional) if `False`, continuous game play across
      scoring events.
    walker_type: the type of walker to instantiate in the environment.

  Returns:
    A `composer.Environment` instance.

  Raises:
    ValueError: If `team_size` is not between 1 and 11.
    ValueError: If `walker_type` is not recognized.
  """
  goal_size = None
  min_size = (32, 24)
  max_size = (48, 36)
  ball = SoccerBall()

  if walker_type == WalkerType.HUMANOID:
    goal_size = MINI_FOOTBALL_GOAL_SIZE
    num_walkers = team_size * 2
    min_size = _area_to_size(MINI_FOOTBALL_MIN_AREA_PER_HUMANOID * num_walkers)
    max_size = _area_to_size(MINI_FOOTBALL_MAX_AREA_PER_HUMANOID * num_walkers)
    ball = regulation_soccer_ball()

  task_factory = Task
  if not terminate_on_goal:
    task_factory = MultiturnTask

  return composer.Environment(
      task=task_factory(
          players=_make_players(team_size, walker_type),
          arena=RandomizedPitch(
              min_size=min_size,
              max_size=max_size,
              keep_aspect_ratio=keep_aspect_ratio,
              field_box=enable_field_box,
              goal_size=goal_size),
          ball=ball,
          disable_walker_contacts=disable_walker_contacts),
      time_limit=time_limit,
      random_state=random_state)
