# 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.
# ============================================================================
"""Maze-based arenas."""

import string

from absl import logging
from dm_control import composer
from dm_control.composer.observation import observable
from dm_control.locomotion.arenas import assets as locomotion_arenas_assets
from dm_control.locomotion.arenas import covering
import labmaze
import numpy as np


# Put all "actual" wall geoms in a separate group since they are not rendered.
_WALL_GEOM_GROUP = 3

_TOP_CAMERA_DISTANCE = 100
_TOP_CAMERA_Y_PADDING_FACTOR = 1.1

_DEFAULT_WALL_CHAR = '*'
_DEFAULT_FLOOR_CHAR = '.'


class MazeWithTargets(composer.Arena):
  """A 2D maze with target positions specified by a LabMaze-style text maze."""

  def _build(self, maze, xy_scale=2.0, z_height=2.0,
             skybox_texture=None, wall_textures=None, floor_textures=None,
             aesthetic='default', name='maze'):
    """Initializes this maze arena.

    Args:
      maze: A `labmaze.BaseMaze` instance.
      xy_scale: The size of each maze cell in metres.
      z_height: The z-height of the maze in metres.
      skybox_texture: (optional) A `composer.Entity` that provides a texture
        asset for the skybox.
      wall_textures: (optional) Either a `composer.Entity` that provides texture
        assets for the maze walls, or a dict mapping printable characters to
        such Entities. In the former case, the maze walls are assumed to be
        represented by '*' in the maze's entity layer. In the latter case,
        the dict's keys specify the different characters that can be present
        in the maze's entity layer, and the dict's values are the corresponding
        texture providers.
      floor_textures: (optional) A `composer.Entity` that provides texture
        assets for the maze floor. Unlike with walls, we do not currently
        support per-variation floor texture. Instead, we sample textures from
        the same texture provider for each variation in the variations layer.
      aesthetic: option to adjust the material properties and skybox
      name: (optional) A string, the name of this arena.
    """
    super()._build(name)
    self._maze = maze
    self._xy_scale = xy_scale
    self._z_height = z_height

    self._x_offset = (self._maze.width - 1) / 2
    self._y_offset = (self._maze.height - 1) / 2

    self._mjcf_root.default.geom.rgba = [1, 1, 1, 1]

    if aesthetic != 'default':
      sky_info = locomotion_arenas_assets.get_sky_texture_info(aesthetic)
      texturedir = locomotion_arenas_assets.get_texturedir(aesthetic)
      self._mjcf_root.compiler.texturedir = texturedir
      self._skybox = self._mjcf_root.asset.add(
          'texture', name='aesthetic_skybox', file=sky_info.file,
          type='skybox', gridsize=sky_info.gridsize,
          gridlayout=sky_info.gridlayout)
    elif skybox_texture:
      self._skybox_texture = skybox_texture.texture
      self.attach(skybox_texture)
    else:
      self._skybox_texture = self._mjcf_root.asset.add(
          'texture', type='skybox', name='skybox', builtin='gradient',
          rgb1=[.4, .6, .8], rgb2=[0, 0, 0], width=100, height=100)

    self._texturing_geom_names = []
    self._texturing_material_names = []
    if wall_textures:
      if isinstance(wall_textures, dict):
        for texture_provider in set(wall_textures.values()):
          self.attach(texture_provider)
        self._wall_textures = {
            wall_char: texture_provider.textures
            for wall_char, texture_provider in wall_textures.items()
        }
      else:
        self.attach(wall_textures)
        self._wall_textures = {_DEFAULT_WALL_CHAR: wall_textures.textures}
    else:
      self._wall_textures = {_DEFAULT_WALL_CHAR: [self._mjcf_root.asset.add(
          'texture', type='2d', name='wall', builtin='flat',
          rgb1=[.8, .8, .8], width=100, height=100)]}

    if aesthetic != 'default':
      ground_info = locomotion_arenas_assets.get_ground_texture_info(aesthetic)
      self._floor_textures = [
          self._mjcf_root.asset.add(
              'texture',
              name='aesthetic_texture_main',
              file=ground_info.file,
              type=ground_info.type),
          self._mjcf_root.asset.add(
              'texture',
              name='aesthetic_texture',
              file=ground_info.file,
              type=ground_info.type)
      ]
    elif floor_textures:
      self._floor_textures = floor_textures.textures
      self.attach(floor_textures)
    else:
      self._floor_textures = [self._mjcf_root.asset.add(
          'texture', type='2d', name='floor', builtin='flat',
          rgb1=[.2, .2, .2], width=100, height=100)]

    ground_x = ((self._maze.width - 1) + 1) * (xy_scale / 2)
    ground_y = ((self._maze.height - 1) + 1) * (xy_scale / 2)
    self._mjcf_root.worldbody.add(
        'geom', name='ground', type='plane',
        pos=[0, 0, 0], size=[ground_x, ground_y, 1], rgba=[0, 0, 0, 0])

    self._maze_body = self._mjcf_root.worldbody.add('body', name='maze_body')

    self._mjcf_root.visual.map.znear = 0.0005

    # Choose the FOV so that the maze always fits nicely within the frame
    # irrespective of actual maze size.
    maze_size = max(self._maze.width, self._maze.height)
    top_camera_fovy = (360 / np.pi) * np.arctan2(
        _TOP_CAMERA_Y_PADDING_FACTOR * maze_size * self._xy_scale / 2,
        _TOP_CAMERA_DISTANCE)
    self._top_camera = self._mjcf_root.worldbody.add(
        'camera', name='top_camera',
        pos=[0, 0, _TOP_CAMERA_DISTANCE], zaxis=[0, 0, 1], fovy=top_camera_fovy)

    self._target_positions = ()
    self._spawn_positions = ()

    self._text_maze_regenerated_hook = None
    self._tile_geom_names = {}

  def _build_observables(self):
    return MazeObservables(self)

  @property
  def top_camera(self):
    return self._top_camera

  @property
  def xy_scale(self):
    return self._xy_scale

  @property
  def z_height(self):
    return self._z_height

  @property
  def maze(self):
    return self._maze

  @property
  def text_maze_regenerated_hook(self):
    """A callback that is executed after the LabMaze object is regenerated."""
    return self._text_maze_modifier

  @text_maze_regenerated_hook.setter
  def text_maze_regenerated_hook(self, hook):
    self._text_maze_regenerated_hook = hook

  @property
  def target_positions(self):
    """A tuple of Cartesian target positions generated for the current maze."""
    return self._target_positions

  @property
  def spawn_positions(self):
    """The Cartesian position at which the agent should be spawned."""
    return self._spawn_positions

  @property
  def target_grid_positions(self):
    """A tuple of grid coordinates of targets generated for the current maze."""
    return self._target_grid_positions

  @property
  def spawn_grid_positions(self):
    """The grid-coordinate position at which the agent should be spawned."""
    return self._spawn_grid_positions

  def regenerate(self, random_state=np.random.RandomState()):
    """Generates a new maze layout."""
    del random_state
    self._maze.regenerate()
    logging.debug('GENERATED MAZE:\n%s', self._maze.entity_layer)
    self._find_spawn_and_target_positions()

    if self._text_maze_regenerated_hook:
      self._text_maze_regenerated_hook()

    # Remove old texturing planes.
    for geom_name in self._texturing_geom_names:
      del self._mjcf_root.worldbody.geom[geom_name]
    self._texturing_geom_names = []

    # Remove old texturing materials.
    for material_name in self._texturing_material_names:
      del self._mjcf_root.asset.material[material_name]
    self._texturing_material_names = []

    # Remove old actual-wall geoms.
    self._maze_body.geom.clear()

    self._current_wall_texture = {
        wall_char: np.random.choice(wall_textures)
        for wall_char, wall_textures in self._wall_textures.items()
    }

    for wall_char in self._wall_textures:
      self._make_wall_geoms(wall_char)
    self._make_floor_variations()

  def _make_wall_geoms(self, wall_char):
    walls = covering.make_walls(
        self._maze.entity_layer, wall_char=wall_char, make_odd_sized_walls=True)
    for i, wall in enumerate(walls):
      wall_mid = covering.GridCoordinates(
          (wall.start.y + wall.end.y - 1) / 2,
          (wall.start.x + wall.end.x - 1) / 2)
      wall_pos = np.array([(wall_mid.x - self._x_offset) * self._xy_scale,
                           -(wall_mid.y - self._y_offset) * self._xy_scale,
                           self._z_height / 2])
      wall_size = np.array([(wall.end.x - wall_mid.x - 0.5) * self._xy_scale,
                            (wall.end.y - wall_mid.y - 0.5) * self._xy_scale,
                            self._z_height / 2])
      self._maze_body.add('geom', name='wall{}_{}'.format(wall_char, i),
                          type='box', pos=wall_pos, size=wall_size,
                          group=_WALL_GEOM_GROUP)
      self._make_wall_texturing_planes(wall_char, i, wall_pos, wall_size)

  def _make_wall_texturing_planes(self, wall_char, wall_id,
                                  wall_pos, wall_size):
    xyaxes = {
        'x': {-1: [0, -1, 0, 0, 0, 1], 1: [0, 1, 0, 0, 0, 1]},
        'y': {-1: [1, 0, 0, 0, 0, 1], 1: [-1, 0, 0, 0, 0, 1]},
        'z': {-1: [-1, 0, 0, 0, 1, 0], 1: [1, 0, 0, 0, 1, 0]}
    }
    for direction_index, direction in enumerate(('x', 'y', 'z')):
      index = list(i for i in range(3) if i != direction_index)
      delta_vector = np.array([int(i == direction_index) for i in range(3)])
      material_name = 'wall{}_{}_{}'.format(wall_char, wall_id, direction)
      self._texturing_material_names.append(material_name)
      mat = self._mjcf_root.asset.add(
          'material', name=material_name,
          texture=self._current_wall_texture[wall_char],
          texrepeat=(2 * wall_size[index] / self._xy_scale))
      for sign, sign_name in zip((-1, 1), ('neg', 'pos')):
        if direction == 'z' and sign == -1:
          continue
        geom_name = (
            'wall{}_{}_texturing_{}_{}'.format(
                wall_char, wall_id, sign_name, direction))
        self._texturing_geom_names.append(geom_name)
        self._mjcf_root.worldbody.add(
            'geom', type='plane', name=geom_name,
            pos=(wall_pos + sign * delta_vector * wall_size),
            size=np.concatenate([wall_size[index], [self._xy_scale]]),
            xyaxes=xyaxes[direction][sign], material=mat,
            contype=0, conaffinity=0)

  def _make_floor_variations(self, build_tile_geoms_fn=None):
    """Builds the floor tiles.

    Args:
      build_tile_geoms_fn: An optional callable returning floor tile geoms.
        If not passed, the floor will be built using a default covering method.
        Takes a kwarg `wall_char` that can be used control how active floor
        tiles are selected.
    """
    main_floor_texture = np.random.choice(self._floor_textures)
    for variation in _DEFAULT_FLOOR_CHAR + string.ascii_uppercase:
      if variation not in self._maze.variations_layer:
        break

      if build_tile_geoms_fn is None:
        # Break the floor variation down to odd-sized tiles.
        tiles = covering.make_walls(self._maze.variations_layer,
                                    wall_char=variation,
                                    make_odd_sized_walls=True)
      else:
        tiles = build_tile_geoms_fn(wall_char=variation)

      # Sample a texture that's not the same as the main floor texture.
      variation_texture = main_floor_texture
      if variation != _DEFAULT_FLOOR_CHAR:
        if len(self._floor_textures) == 1:
          return
        else:
          while variation_texture is main_floor_texture:
            variation_texture = np.random.choice(self._floor_textures)

      for i, tile in enumerate(tiles):
        tile_mid = covering.GridCoordinates(
            (tile.start.y + tile.end.y - 1) / 2,
            (tile.start.x + tile.end.x - 1) / 2)
        tile_pos = np.array([(tile_mid.x - self._x_offset) * self._xy_scale,
                             -(tile_mid.y - self._y_offset) * self._xy_scale,
                             0.0])
        tile_size = np.array([(tile.end.x - tile_mid.x - 0.5) * self._xy_scale,
                              (tile.end.y - tile_mid.y - 0.5) * self._xy_scale,
                              self._xy_scale])
        if variation == _DEFAULT_FLOOR_CHAR:
          tile_name = 'floor_{}'.format(i)
        else:
          tile_name = 'floor_{}_{}'.format(variation, i)
        self._tile_geom_names[tile.start] = tile_name
        self._texturing_material_names.append(tile_name)
        self._texturing_geom_names.append(tile_name)
        material = self._mjcf_root.asset.add(
            'material', name=tile_name, texture=variation_texture,
            texrepeat=(2 * tile_size[[0, 1]] / self._xy_scale))
        self._mjcf_root.worldbody.add(
            'geom', name=tile_name, type='plane', material=material,
            pos=tile_pos, size=tile_size, contype=0, conaffinity=0)

  @property
  def ground_geoms(self):
    return tuple([
        geom for geom in self.mjcf_model.find_all('geom')
        if 'ground' in geom.name
    ])

  def find_token_grid_positions(self, tokens):
    out = {token: [] for token in tokens}
    for y in range(self._maze.entity_layer.shape[0]):
      for x in range(self._maze.entity_layer.shape[1]):
        for token in tokens:
          if self._maze.entity_layer[y, x] == token:
            out[token].append((y, x))
    return out

  def grid_to_world_positions(self, grid_positions):
    out = []
    for y, x in grid_positions:
      out.append(np.array([(x - self._x_offset) * self._xy_scale,
                           -(y - self._y_offset) * self._xy_scale,
                           0.0]))
    return out

  def world_to_grid_positions(self, world_positions):
    out = []
    # the order of x, y is reverse between grid positions format and
    # world positions format.
    for x, y, _ in world_positions:
      out.append(np.array([self._y_offset - y / self._xy_scale,
                           self._x_offset + x / self._xy_scale]))
    return out

  def _find_spawn_and_target_positions(self):
    grid_positions = self.find_token_grid_positions([
        labmaze.defaults.OBJECT_TOKEN, labmaze.defaults.SPAWN_TOKEN])
    self._target_grid_positions = tuple(
        grid_positions[labmaze.defaults.OBJECT_TOKEN])
    self._spawn_grid_positions = tuple(
        grid_positions[labmaze.defaults.SPAWN_TOKEN])
    self._target_positions = tuple(
        self.grid_to_world_positions(self._target_grid_positions))
    self._spawn_positions = tuple(
        self.grid_to_world_positions(self._spawn_grid_positions))


class MazeObservables(composer.Observables):

  @composer.observable
  def top_camera(self):
    return observable.MJCFCamera(self._entity.top_camera)


class RandomMazeWithTargets(MazeWithTargets):
  """A randomly generated 2D maze with target positions."""

  def _build(self,
             x_cells,
             y_cells,
             xy_scale=2.0,
             z_height=2.0,
             max_rooms=labmaze.defaults.MAX_ROOMS,
             room_min_size=labmaze.defaults.ROOM_MIN_SIZE,
             room_max_size=labmaze.defaults.ROOM_MAX_SIZE,
             spawns_per_room=labmaze.defaults.SPAWN_COUNT,
             targets_per_room=labmaze.defaults.OBJECT_COUNT,
             max_variations=labmaze.defaults.MAX_VARIATIONS,
             simplify=labmaze.defaults.SIMPLIFY,
             skybox_texture=None,
             wall_textures=None,
             floor_textures=None,
             aesthetic='default',
             name='random_maze'):
    """Initializes this random maze arena.

    Args:
      x_cells: The number of cells along the x-direction of the maze. Must be
        an odd integer.
      y_cells: The number of cells along the y-direction of the maze. Must be
        an odd integer.
      xy_scale: The size of each maze cell in metres.
      z_height: The z-height of the maze in metres.
      max_rooms: (optional) The maximum number of rooms in each generated maze.
      room_min_size: (optional) The minimum size of each room generated.
      room_max_size: (optional) The maximum size of each room generated.
      spawns_per_room: (optional) Number of spawn points
        to generate in each room.
      targets_per_room: (optional) Number of targets to generate in each room.
      max_variations: (optional) Maximum number of variations to generate
        in the variations layer.
      simplify: (optional) flag to simplify the maze.
      skybox_texture: (optional) A `composer.Entity` that provides a texture
        asset for the skybox.
      wall_textures: (optional) A `composer.Entity` that provides texture
        assets for the maze walls.
      floor_textures: (optional) A `composer.Entity` that provides texture
        assets for the maze floor.
      aesthetic: option to adjust the material properties and skybox
      name: (optional) A string, the name of this arena.
    """
    random_seed = np.random.randint(2147483648)  # 2**31
    super()._build(
        maze=labmaze.RandomMaze(
            height=y_cells,
            width=x_cells,
            max_rooms=max_rooms,
            room_min_size=room_min_size,
            room_max_size=room_max_size,
            max_variations=max_variations,
            spawns_per_room=spawns_per_room,
            objects_per_room=targets_per_room,
            simplify=simplify,
            random_seed=random_seed),
        xy_scale=xy_scale,
        z_height=z_height,
        skybox_texture=skybox_texture,
        wall_textures=wall_textures,
        floor_textures=floor_textures,
        aesthetic=aesthetic,
        name=name)
