# Path: spriteworld/action_spaces.py
# class SelectMove(object):
# class DragAndDrop(SelectMove):
# class Embodied(object):
#   def __init__(self, scale=1.0, motion_cost=0.0, noise_scale=None):
#   def get_motion(self, action):
#   def apply_noise_to_action(self, action):
#   def get_sprite_from_position(self, position, sprites):
#   def step(self, action, sprites, keep_in_frame):
#   def sample(self):
#   def action_spec(self):
#   def get_motion(self, action):
#   def __init__(self, step_size=0.05, motion_cost=0.):
#   def get_body_sprite(self, sprites):
#   def get_non_body_sprites(self, sprites):
#   def get_carried_sprite(self, sprites):
#   def step(self, action, sprites, keep_in_frame):
#   def sample(self):
#   def action_spec(self):
# 
# 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):


from absl.testing import absltest
from absl.testing import parameterized
from spriteworld import action_spaces
from spriteworld import sprite
import numpy as np



class SelectMoveTest(parameterized.TestCase):

  @parameterized.named_parameters(
      ('Motion', 1, np.array([0.5, 0.5, 0.2, 0.75]), (-0.3, 0.25)),
      ('SameMotion', 1, np.array([0.2, 0.5, 0.2, 0.75]), (-0.3, 0.25)),
      ('SmallerScale', 0.5, np.array([0.2, 0.5, 0.2, 0.75]), (-0.15, 0.125)),
  )
  def testGetMotion(self, scale, action, true_motion):
    action_space = action_spaces.SelectMove(scale=scale)
    motion = action_space.get_motion(action)
    self.assertTrue(np.allclose(motion, true_motion, atol=1e-4))

  @parameterized.named_parameters(
      ('NoCost', 1, np.array([0.5, 0.5, 0.2, 0.75]), 0., 0.),
      ('Cost', 1, np.array([0.5, 0.5, 0.2, 0.75]), 1., -0.39),
      ('SameCost', 1, np.array([0.2, 0.3, 0.2, 0.75]), 1., -0.39),
      ('LowerCost', 0.5, np.array([0.5, 0.5, 0.2, 0.75]), 1., -0.195),
  )
  def testMotionCost(self, scale, action, motion_cost, true_cost):
    action_space = action_spaces.SelectMove(
        scale=scale, motion_cost=motion_cost)
    cost = action_space.step(action, sprites=[], keep_in_frame=False)
    self.assertAlmostEqual(cost, true_cost, delta=0.01)

  def testMoveSprites(self):
    """Take a series of actions and repeatedly check sprite motions."""
    action_space = action_spaces.SelectMove(scale=0.5)

# Path: spriteworld/action_spaces.py
# class SelectMove(object):
# class DragAndDrop(SelectMove):
# class Embodied(object):
#   def __init__(self, scale=1.0, motion_cost=0.0, noise_scale=None):
#   def get_motion(self, action):
#   def apply_noise_to_action(self, action):
#   def get_sprite_from_position(self, position, sprites):
#   def step(self, action, sprites, keep_in_frame):
#   def sample(self):
#   def action_spec(self):
#   def get_motion(self, action):
#   def __init__(self, step_size=0.05, motion_cost=0.):
#   def get_body_sprite(self, sprites):
#   def get_non_body_sprites(self, sprites):
#   def get_carried_sprite(self, sprites):
#   def step(self, action, sprites, keep_in_frame):
#   def sample(self):
#   def action_spec(self):
# 
# 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: tests/action_spaces_test.py
from absl.testing import absltest
from absl.testing import parameterized
from spriteworld import action_spaces
from spriteworld import sprite
import numpy as np




class SelectMoveTest(parameterized.TestCase):

  @parameterized.named_parameters(
      ('Motion', 1, np.array([0.5, 0.5, 0.2, 0.75]), (-0.3, 0.25)),
      ('SameMotion', 1, np.array([0.2, 0.5, 0.2, 0.75]), (-0.3, 0.25)),
      ('SmallerScale', 0.5, np.array([0.2, 0.5, 0.2, 0.75]), (-0.15, 0.125)),
  )
  def testGetMotion(self, scale, action, true_motion):
    action_space = action_spaces.SelectMove(scale=scale)
    motion = action_space.get_motion(action)
    self.assertTrue(np.allclose(motion, true_motion, atol=1e-4))

  @parameterized.named_parameters(
      ('NoCost', 1, np.array([0.5, 0.5, 0.2, 0.75]), 0., 0.),
      ('Cost', 1, np.array([0.5, 0.5, 0.2, 0.75]), 1., -0.39),
      ('SameCost', 1, np.array([0.2, 0.3, 0.2, 0.75]), 1., -0.39),
      ('LowerCost', 0.5, np.array([0.5, 0.5, 0.2, 0.75]), 1., -0.195),
  )
  def testMotionCost(self, scale, action, motion_cost, true_cost):
    action_space = action_spaces.SelectMove(
        scale=scale, motion_cost=motion_cost)
    cost = action_space.step(action, sprites=[], keep_in_frame=False)
    self.assertAlmostEqual(cost, true_cost, delta=0.01)

  def testMoveSprites(self):
    """Take a series of actions and repeatedly check sprite motions."""
    action_space = action_spaces.SelectMove(scale=0.5)

    sprites = [sprite.Sprite(x=0.55, y=0.5), sprite.Sprite(x=0.5, y=0.5)]