# Copyright 2019 DeepMind Technologies Limited.
#
# 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
#
#    https://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.
# ============================================================================
# python2 python3
"""Interactive GUI for Spriteworld.

Be aware that this UI overrides the action space and renderer for ease of
playing, so those will be different from what are specified in the task config.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import logging as log
import sys
from absl import logging
from matplotlib import gridspec
import matplotlib.pylab as plt
import numpy as np

from spriteworld import action_spaces
from spriteworld import environment
from spriteworld import renderers


class MatplotlibUI(object):
  """Class for visualising the environment based on Matplotlib."""

  def __init__(self):
    self.rewards = 10 * [np.nan]
    self.rewards_bounds = [-10, 10]
    self.last_success = None

    plt.ion()
    self._fig = plt.figure(
        figsize=(9, 12), num='Spriteworld', facecolor='white')
    gs = gridspec.GridSpec(2, 1, height_ratios=[3, 1])
    self._ax_image = plt.subplot(gs[0])
    self._ax_image.axis('off')

    self._ax_scalar = plt.subplot(gs[1])
    self._ax_scalar.spines['right'].set_visible(False)
    self._ax_scalar.spines['top'].set_visible(False)
    self._ax_scalar.xaxis.set_ticks_position('bottom')
    self._ax_scalar.yaxis.set_ticks_position('left')
    self._setup_callbacks()

  @property
  def ax_image(self):
    return self._ax_image

  def _setup_callbacks(self):
    """Default callbacks for the UI."""

    # Pressing escape should stop the UI
    def _onkeypress(event):
      if event.key == 'escape':
        # Stop UI
        logging.info('Pressed escape, stopping UI.')
        plt.close(self._fig)
        sys.exit()

    self._fig.canvas.mpl_connect('key_release_event', _onkeypress)

    # Disable default keyboard shortcuts
    for key in ('keymap.fullscreen', 'keymap.home', 'keymap.back',
                'keymap.forward', 'keymap.pan', 'keymap.zoom', 'keymap.save',
                'keymap.quit', 'keymap.grid', 'keymap.yscale', 'keymap.xscale',
                'keymap.all_axes'):
      plt.rcParams[key] = ''

    # Disable logging of some matplotlib events
    log.getLogger('matplotlib').setLevel('WARNING')

  def _draw_observation(self, image, action):
    """Draw the latest observation."""
    self._ax_image.clear()
    self._ax_image.imshow(image, interpolation='none')
    self._ax_image.set_xticks([])
    self._ax_image.set_yticks([])
    if action is not None:
      self._ax_image.annotate(
          '',
          xycoords='axes fraction',
          xy=action[:2],  # Start of arrow
          xytext=action[2:],  # End of arrow
          arrowprops={
              'arrowstyle': '<|-',
              'color': 'red',
              'lw': 4,
          })

    # Indicate success
    linewidth = 1
    color = 'black'
    if np.isnan(self.rewards[-1]):
      linewidth = 8
      color = 'green' if self.last_success else 'red'

    for sp in self._ax_image.spines.values():
      sp.set_color(color)
      sp.set_linewidth(linewidth)

  def _draw_rewards(self):
    """Draw the past rewards plot."""
    self._ax_scalar.clear()
    self._ax_scalar.set_ylabel('Rewards')
    self._ax_scalar.set_xlabel('Timestep')
    xs = np.arange(-len(self.rewards), 0)
    self._ax_scalar.set_xticks(xs)
    self._ax_scalar.axhline(y=0.0, color='lightgrey', linestyle='--')
    self._ax_scalar.stem(xs, self.rewards, basefmt=' ')

    self._ax_scalar.set_xlim((xs[0] - 1.0, xs[-1] + 1.0))
    self._ax_scalar.set_ylim(
        (self.rewards_bounds[0] - 1.0, self.rewards_bounds[1] + 1.0))

  def register_callback(self, event_name, callback):
    """Register a callback for the given event."""
    self._fig.canvas.mpl_connect(event_name, callback)

  def update(self, timestep, action):
    """Update the visualisation with the latest timestep and action."""
    reward = timestep.reward
    if reward is None:
      reward = np.nan
    self.rewards = self.rewards[1:] + [reward]
    self.rewards_bounds[0] = np.nanmin(
        [np.nanmin(self.rewards), self.rewards_bounds[0]])
    self.rewards_bounds[1] = np.nanmax(
        [np.nanmax(self.rewards), self.rewards_bounds[1]])
    self._draw_observation(timestep.observation['image'], action)
    self._draw_rewards()
    plt.show(block=False)

    self.last_success = timestep.observation['success']


class HumanDragAndDropAgent(object):
  """Demo agent for mouse-clicking interface with DragAndDrop action space."""

  def __init__(self, action_space, timeout=600):
    self._action_space = action_space
    self._click = None
    self._timeout = timeout

  def help(self):
    logging.info('Click to select an object, then click again to select where '
                 'to move it.')

  def register_callbacks(self, ui):
    """Register the matplotlib callbacks required by the agent."""

    def _onclick(event):
      if event.inaxes and event.inaxes == ui.ax_image:
        # Map the click into axis-fraction positions (origin at bottom-left).
        self._click = event.inaxes.transAxes.inverted().transform(
            (event.x, event.y))
      else:
        self._click = None
      return

    ui.register_callback('button_press_event', _onclick)

  def begin_episode(self):
    logging.info('Starting episode')

  def step(self, timestep):
    """Take a step."""
    del timestep  # Unused

    def _get_click():
      """Get mouse click."""
      click = None
      while click is None:
        x = plt.waitforbuttonpress(timeout=self._timeout)
        if x is None:
          logging.info('Timed out. You took longer than %d seconds to click.',
                       self._timeout)
        elif x:
          logging.info('You pressed a key, but were supposed to click with the '
                       'mouse.')
          self.help()
        else:
          click = self._click
      return click

    def _get_action():
      """Get action from user."""
      logging.info('Select sprite')
      click_from = _get_click()
      logging.info('Select target')
      click_to = _get_click()
      try:
        action = np.concatenate((click_from, click_to)).astype(np.float32)
        if any(np.isnan(action)):
          raise ValueError
        self._action_space.action_spec().validate(action)
        return action
      except (ValueError, TypeError):
        logging.info('Select a valid action')
        return _get_action()

    action = _get_action()
    return action


class HumanEmbodiedAgent(object):
  """Demo agent for keyboard interface with Embodied action space."""

  MOTION_KEY_TO_ACTION = {
      'up': 0,
      'left': 1,
      'down': 2,
      'right': 3,
      'w': 0,
      'a': 1,
      's': 2,
      'd': 3
  }

  def __init__(self, action_space, timeout=600):
    self._action_space = action_space
    self._key_press = None
    self._carry = False
    self._movement = None
    self._timeout = timeout

  def help(self):
    logging.info('Use WASD/arrow keys to move, hold Space to carry.')

  def register_callbacks(self, ui):
    """Register the matplotlib callbacks required by the agent."""

    def _onkeypress(event):
      if event.key in HumanEmbodiedAgent.MOTION_KEY_TO_ACTION:
        self._movement = HumanEmbodiedAgent.MOTION_KEY_TO_ACTION[event.key]
      elif event.key == ' ':
        self._carry = True
      else:
        self.help()

    ui.register_callback('key_press_event', _onkeypress)

    def _onkeyrelease(event):
      if event.key == ' ':
        self._carry = False
      elif event.key in HumanEmbodiedAgent.MOTION_KEY_TO_ACTION:
        self._movement = None

    ui.register_callback('key_release_event', _onkeyrelease)

  def begin_episode(self):
    logging.info('Starting episode')

  def step(self, timestep):
    """Take a step."""
    del timestep  # Unused

    def _wait_for_movement_key_press():
      """Get key press."""
      ready = False
      while not ready:
        x = plt.waitforbuttonpress(timeout=self._timeout)
        if x is None:
          logging.info('Timed out. You took longer than %d seconds to click.',
                       self._timeout)
        elif not x:
          logging.info('You clicked, but you are supposed to use the Keyboard.')
          self.help()
        elif self._movement is not None:
          ready = True

    def _get_action():
      """Get action from user."""
      _wait_for_movement_key_press()

      action = (int(self._carry), self._movement)
      for spec, a in zip(self._action_space.action_spec(), action):
        spec.validate(a)
      return action

    return _get_action()


def setup_run_ui(env_config, render_size, task_hsv_colors, anti_aliasing):
  """Start a Demo UI given an env_config."""
  if isinstance(env_config['action_space'], action_spaces.SelectMove):
    # DragAndDrop is a bit easier to demo than the SelectMove action space
    env_config['action_space'] = action_spaces.DragAndDrop(scale=0.5)
    agent = HumanDragAndDropAgent(env_config['action_space'])
  elif isinstance(env_config['action_space'], action_spaces.Embodied):
    agent = HumanEmbodiedAgent(env_config['action_space'])
  else:
    raise ValueError(
        'Demo is not configured to run with action space {}.'.format(
            env_config['action_space']))
  env_config['renderers'] = {
      'image':
          renderers.PILRenderer(
              image_size=(render_size, render_size),
              color_to_rgb=renderers.color_maps.hsv_to_rgb
              if task_hsv_colors else None,
              anti_aliasing=anti_aliasing),
      'success':
          renderers.Success()
  }
  env = environment.Environment(**env_config)
  ui = MatplotlibUI()
  agent.register_callbacks(ui)

  # Start RL loop
  timestep = env.reset()
  ui.update(timestep, action=None)

  while True:
    action = agent.step(timestep)
    timestep = env.step(action)
    if isinstance(env_config['action_space'], action_spaces.DragAndDrop):
      ui.update(timestep, action)
    else:
      ui.update(timestep, None)
