# 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
"""Python Image Library (PIL/Pillow) renderer."""

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

from dm_env import specs
import numpy as np
from PIL import Image
from PIL import ImageDraw
from spriteworld.renderers import abstract_renderer


class PILRenderer(abstract_renderer.AbstractRenderer):
  """Render using Python Image Library (PIL/Pillow)."""

  def __init__(self,
               image_size=(64, 64),
               anti_aliasing=1,
               bg_color=None,
               color_to_rgb=None):
    """Construct PIL renderer.

    Args:
      image_size: Int tuple (height, width). Size of output of .render().
      anti_aliasing: Int. Anti-aliasing factor. Linearly scales the size of the
        internal canvas.
      bg_color: None or 3-tuple of ints in [0, 255]. Background color. If None,
        background is (0, 0, 0).
      color_to_rgb: Callable converting a tuple (c1, c2, c3) to a uint8 tuple
        (r, g, b) in [0, 255].
    """
    self._image_size = image_size
    self._anti_aliasing = anti_aliasing
    self._canvas_size = (anti_aliasing * image_size[0],
                         anti_aliasing * image_size[1])

    if color_to_rgb is None:
      color_to_rgb = lambda x: x
    self._color_to_rgb = color_to_rgb

    if bg_color is None:
      bg_color = (0, 0, 0)
    self._canvas_bg = Image.new('RGB', self._canvas_size, bg_color)

    self._observation_spec = specs.Array(
        shape=self._image_size + (3,), dtype=np.uint8)

    self._canvas = Image.new('RGB', self._canvas_size)
    self._draw = ImageDraw.Draw(self._canvas)

  def render(self, sprites=(), global_state=None):
    """Render sprites.

    Sprites are ordered from background to foreground.

    Args:
      sprites: Iterable of sprite.Sprite instances.
      global_state: Unused global state.

    Returns:
      Numpy uint8 RGB array of size self._image_size + (3,).
    """
    self._canvas.paste(self._canvas_bg)
    for obj in sprites:
      vertices = self._canvas_size * obj.vertices
      color = self._color_to_rgb(obj.color)
      self._draw.polygon([tuple(v) for v in vertices], fill=color)
    image = self._canvas.resize(self._image_size, resample=Image.ANTIALIAS)

    # PIL uses a coordinate system with the origin (0, 0) at the upper-left, but
    # our environment uses an origin at the bottom-left (i.e. mathematical
    # convention). Hence we need to flip the render vertically to correct for
    # that.
    image = np.flipud(np.array(image))
    return image

  def observation_spec(self):
    return self._observation_spec
