# Path: spriteworld/constants.py
# SHAPES = {
#     'triangle': shapes.polygon(num_sides=3, theta_0=np.pi/2),
#     'square': shapes.polygon(num_sides=4, theta_0=np.pi/4),
#     'pentagon': shapes.polygon(num_sides=5, theta_0=np.pi/2),
#     'hexagon': shapes.polygon(num_sides=6),
#     'octagon': shapes.polygon(num_sides=8),
#     'circle': shapes.polygon(num_sides=30),
#     'star_4': shapes.star(num_sides=4, theta_0=np.pi/4),
#     'star_5': shapes.star(num_sides=5, theta_0=np.pi + np.pi/10),
#     'star_6': shapes.star(num_sides=6),
#     'spoke_4': shapes.spokes(num_sides=4, theta_0=np.pi/4),
#     'spoke_5': shapes.spokes(num_sides=5, theta_0=np.pi + np.pi/10),
#     'spoke_6': shapes.spokes(num_sides=6),
# }
# class ShapeType(enum.IntEnum):
# 
# Path: spriteworld/sprite.py
# FACTOR_NAMES = (
#     'x',  # x-position of sprite center-of-mass (float)
#     'y',  # y-position of sprite center-of-mass (float)
#     'shape',  # shape (string)
#     'angle',  # angle in degrees (scalar)
#     'scale',  # size of sprite (float)
#     'c0',  # first color component (scalar)
#     'c1',  # second color component (scalar)
#     'c2',  # third color component (scalar)
#     'x_vel',  # x-component of velocity (float)
#     'y_vel',  # y-component of velocity (float)
# )
# _MAX_TRIES = int(1e6)
# class Sprite(object):
#   def __init__(self,
#                x=0.5,
#                y=0.5,
#                shape='square',
#                angle=0,
#                scale=0.1,
#                c0=0,
#                c1=0,
#                c2=0,
#                x_vel=0.0,
#                y_vel=0.0):
#   def _reset_centered_path(self):
#   def move(self, motion, keep_in_frame=False):
#   def update_position(self, keep_in_frame=False):
#   def contains_point(self, point):
#   def sample_contained_position(self):
#   def vertices(self):
#   def out_of_frame(self):
#   def x(self):
#   def y(self):
#   def shape(self):
#   def shape(self, s):
#   def angle(self):
#   def angle(self, a):
#   def scale(self):
#   def scale(self, s):
#   def c0(self):
#   def c1(self):
#   def c2(self):
#   def x_vel(self):
#   def y_vel(self):
#   def color(self):
#   def position(self):
#   def velocity(self):
#   def factors(self):
# 
# Path: spriteworld/renderers/handcrafted.py
# class SpriteFactors(abstract_renderer.AbstractRenderer):
# class SpritePassthrough(abstract_renderer.AbstractRenderer):
# class Success(abstract_renderer.AbstractRenderer):
#   def __init__(self, factors=sprite_lib.FACTOR_NAMES):
#   def render(self, sprites=(), global_state=None):
#     def _process_factor(name, value):
#     def _sprite_to_factors(sprite):
#   def observation_spec(self):
#   def __init__(self):
#   def render(self, sprites=(), global_state=None):
#   def observation_spec(self):
#   def render(self, sprites=(), global_state=None):
#   def observation_spec(self):


from absl.testing import absltest
from absl.testing import parameterized
from six.moves import range
from spriteworld import constants as const
from spriteworld import sprite as sprite_lib
from spriteworld.renderers import handcrafted
import numpy as np

  @parameterized.parameters((1, ('x', 'y', 'scale', 'c0', 'c1', 'c2', 'shape',
                                 'angle', 'x_vel', 'y_vel')),
                            (1, ('x', 'y', 'scale', 'c0', 'c1', 'c2', 'shape')),
                            (2, ('x', 'y', 'scale', 'c0', 'c1', 'c2', 'shape')),
                            (5, ('x', 'y', 'angle', 'x_vel', 'y_vel')))
  def testObservationSpec(self, num_sprites, factors):
    sprites = [sprite_lib.Sprite() for _ in range(num_sprites)]
    renderer = handcrafted.SpriteFactors(factors=factors)
    renderer.render(sprites=sprites)
    obs_spec = renderer.observation_spec()

    for v in obs_spec[0].values():
      self.assertEqual(v.shape, ())

    obs_spec_keys = [set(x) for x in obs_spec]
    self.assertSequenceEqual(obs_spec_keys, num_sprites * [set(factors)])

  @parameterized.parameters(
      (0.5, 0.5, 'square', 0, 0, 255, 0.5, 0),
      (0.5, 0.5, 'square', 255, 0, 0, 0.5, 0),
      (0.5, 0.8, 'octagon', 0.4, 0.8, 0.5, 0.6, 90),
      (0.5, 0.3, 'star_5', 180, 180, 0, 0.2, 240),
  )
  def testAttributesSingleton(self, x, y, shape, c0, c1, c2, scale, angle):
    sprite = sprite_lib.Sprite(
        x=x, y=y, shape=shape, c0=c0, c1=c1, c2=c2, scale=scale, angle=angle)
    renderer = handcrafted.SpriteFactors()
    outputs = renderer.render(sprites=[sprite])[0]


# Path: spriteworld/constants.py
# SHAPES = {
#     'triangle': shapes.polygon(num_sides=3, theta_0=np.pi/2),
#     'square': shapes.polygon(num_sides=4, theta_0=np.pi/4),
#     'pentagon': shapes.polygon(num_sides=5, theta_0=np.pi/2),
#     'hexagon': shapes.polygon(num_sides=6),
#     'octagon': shapes.polygon(num_sides=8),
#     'circle': shapes.polygon(num_sides=30),
#     'star_4': shapes.star(num_sides=4, theta_0=np.pi/4),
#     'star_5': shapes.star(num_sides=5, theta_0=np.pi + np.pi/10),
#     'star_6': shapes.star(num_sides=6),
#     'spoke_4': shapes.spokes(num_sides=4, theta_0=np.pi/4),
#     'spoke_5': shapes.spokes(num_sides=5, theta_0=np.pi + np.pi/10),
#     'spoke_6': shapes.spokes(num_sides=6),
# }
# class ShapeType(enum.IntEnum):
# 
# Path: spriteworld/sprite.py
# FACTOR_NAMES = (
#     'x',  # x-position of sprite center-of-mass (float)
#     'y',  # y-position of sprite center-of-mass (float)
#     'shape',  # shape (string)
#     'angle',  # angle in degrees (scalar)
#     'scale',  # size of sprite (float)
#     'c0',  # first color component (scalar)
#     'c1',  # second color component (scalar)
#     'c2',  # third color component (scalar)
#     'x_vel',  # x-component of velocity (float)
#     'y_vel',  # y-component of velocity (float)
# )
# _MAX_TRIES = int(1e6)
# class Sprite(object):
#   def __init__(self,
#                x=0.5,
#                y=0.5,
#                shape='square',
#                angle=0,
#                scale=0.1,
#                c0=0,
#                c1=0,
#                c2=0,
#                x_vel=0.0,
#                y_vel=0.0):
#   def _reset_centered_path(self):
#   def move(self, motion, keep_in_frame=False):
#   def update_position(self, keep_in_frame=False):
#   def contains_point(self, point):
#   def sample_contained_position(self):
#   def vertices(self):
#   def out_of_frame(self):
#   def x(self):
#   def y(self):
#   def shape(self):
#   def shape(self, s):
#   def angle(self):
#   def angle(self, a):
#   def scale(self):
#   def scale(self, s):
#   def c0(self):
#   def c1(self):
#   def c2(self):
#   def x_vel(self):
#   def y_vel(self):
#   def color(self):
#   def position(self):
#   def velocity(self):
#   def factors(self):
# 
# Path: spriteworld/renderers/handcrafted.py
# class SpriteFactors(abstract_renderer.AbstractRenderer):
# class SpritePassthrough(abstract_renderer.AbstractRenderer):
# class Success(abstract_renderer.AbstractRenderer):
#   def __init__(self, factors=sprite_lib.FACTOR_NAMES):
#   def render(self, sprites=(), global_state=None):
#     def _process_factor(name, value):
#     def _sprite_to_factors(sprite):
#   def observation_spec(self):
#   def __init__(self):
#   def render(self, sprites=(), global_state=None):
#   def observation_spec(self):
#   def render(self, sprites=(), global_state=None):
#   def observation_spec(self):


# Path: tests/renderers/handcrafted_test.py
from absl.testing import absltest
from absl.testing import parameterized
from six.moves import range
from spriteworld import constants as const
from spriteworld import sprite as sprite_lib
from spriteworld.renderers import handcrafted
import numpy as np


  @parameterized.parameters((1, ('x', 'y', 'scale', 'c0', 'c1', 'c2', 'shape',
                                 'angle', 'x_vel', 'y_vel')),
                            (1, ('x', 'y', 'scale', 'c0', 'c1', 'c2', 'shape')),
                            (2, ('x', 'y', 'scale', 'c0', 'c1', 'c2', 'shape')),
                            (5, ('x', 'y', 'angle', 'x_vel', 'y_vel')))
  def testObservationSpec(self, num_sprites, factors):
    sprites = [sprite_lib.Sprite() for _ in range(num_sprites)]
    renderer = handcrafted.SpriteFactors(factors=factors)
    renderer.render(sprites=sprites)
    obs_spec = renderer.observation_spec()

    for v in obs_spec[0].values():
      self.assertEqual(v.shape, ())

    obs_spec_keys = [set(x) for x in obs_spec]
    self.assertSequenceEqual(obs_spec_keys, num_sprites * [set(factors)])

  @parameterized.parameters(
      (0.5, 0.5, 'square', 0, 0, 255, 0.5, 0),
      (0.5, 0.5, 'square', 255, 0, 0, 0.5, 0),
      (0.5, 0.8, 'octagon', 0.4, 0.8, 0.5, 0.6, 90),
      (0.5, 0.3, 'star_5', 180, 180, 0, 0.2, 240),
  )
  def testAttributesSingleton(self, x, y, shape, c0, c1, c2, scale, angle):
    sprite = sprite_lib.Sprite(
        x=x, y=y, shape=shape, c0=c0, c1=c1, c2=c2, scale=scale, angle=angle)
    renderer = handcrafted.SpriteFactors()
    outputs = renderer.render(sprites=[sprite])[0]


    self.assertEqual(outputs['shape'], const.ShapeType[shape].value)