# Path: difflogic/envs/blocksworld/block.py
# def randomly_generate_world(nr_blocks, random_order=False, one_stack=False):
#   """Randomly generate a blocks world case.
# 
#   Similar to classical random tree generation, incrementally add new blocks.
#   for each new block, randomly sample a valid father and stack on its father.
# 
#   Args:
#     nr_blocks: The number of blocks in the world.
#     random_order: Randomly permute the indexes of the blocks if set True.
#         Or set as a provided order. Leave the raw order unchanged in default.
#     one_stack: A special case where only one stack of blocks. If True, for each
#         new node, set its father as the last node.
# 
#   Returns:
#     A BlocksWorld instance which is randomly generated.
#   """
#   blocks = [Block(0, None)]
#   leafs = [blocks[0]]
# 
#   for i in range(1, nr_blocks + 1):
#     other = random.choice_list(leafs)
#     this = Block(i)
#     this.add_to(other)
#     if not other.placeable or one_stack:
#       leafs.remove(other)
#     blocks.append(this)
#     leafs.append(this)
# 
#   order = None
#   if random_order:
#     order = random.permutation(len(blocks))
# 
#   return BlocksWorld(blocks, random_order=order)
# 
# Path: difflogic/envs/blocksworld/represent.py
# def get_coordinates(world, absolute=False):
#   """Get the coordinates of each block in the blocks world."""
#   coordinates = [None for _ in range(world.size)]
#   raw_blocks = world.blocks.raw
# 
#   def dfs(block):
#     """Use depth-first-search to get the coordinate of each block."""
#     if block.is_ground:
#       coordinates[block.index] = (0, 0)
#       for j, c in enumerate(block.children):
#         # When using absolute coordinate, the block x directly placed on the
#         # ground gets coordinate (x, 1).
#         x = world.blocks.inv_index(c.index) if absolute else j
#         coordinates[c.index] = (x, 1)
#         dfs(c)
#     else:
#       coor = coordinates[block.index]
#       assert coor is not None
#       x, y = coor
#       for c in block.children:
#         coordinates[c.index] = (x, y + 1)
#         dfs(c)
# 
#   dfs(raw_blocks[0])
#   coordinates = world.blocks.permute(coordinates)
#   return np.array(coordinates)
# 
# Path: difflogic/envs/blocksworld/represent.py
# def decorate(state, nr_objects, world_id=None):
#   """Append world index and object index information to state."""
#   info = []
#   if world_id is not None:
#     info.append(np.ones((nr_objects, 1)) * world_id)
#   info.extend([np.array(range(nr_objects))[:, np.newaxis], state])
#   return np.hstack(info)


import numpy as np
import jacinle.random as random
from jaclearn.rl.env import SimpleRLEnvBase
from .block import randomly_generate_world
from .represent import get_coordinates
from .represent import decorate
        decorate: if True, the coordinates in the states will also include the
            world index (default: 0) and the block index (starting from 0).
        prob_unchange: The probability that an action is not effective.
        prob_fall: The probability that an action will make the object currently
            moving fall on the ground.
  """

  def __init__(self,
               nr_blocks,
               random_order=False,
               decorate=False,
               prob_unchange=0.0,
               prob_fall=0.0):
    super().__init__()
    self.nr_blocks = nr_blocks
    self.nr_objects = nr_blocks + 1
    self.random_order = random_order
    self.decorate = decorate
    self.prob_unchange = prob_unchange
    self.prob_fall = prob_fall

  def _restart(self):
    self.world = randomly_generate_world(
        self.nr_blocks, random_order=self.random_order)
    self._set_current_state(self._get_decorated_states())
    self.is_over = False
    self.cached_result = self._get_result()

  def _get_decorated_states(self, world_id=0):
    state = get_coordinates(self.world)

# Path: difflogic/envs/blocksworld/block.py
# def randomly_generate_world(nr_blocks, random_order=False, one_stack=False):
#   """Randomly generate a blocks world case.
# 
#   Similar to classical random tree generation, incrementally add new blocks.
#   for each new block, randomly sample a valid father and stack on its father.
# 
#   Args:
#     nr_blocks: The number of blocks in the world.
#     random_order: Randomly permute the indexes of the blocks if set True.
#         Or set as a provided order. Leave the raw order unchanged in default.
#     one_stack: A special case where only one stack of blocks. If True, for each
#         new node, set its father as the last node.
# 
#   Returns:
#     A BlocksWorld instance which is randomly generated.
#   """
#   blocks = [Block(0, None)]
#   leafs = [blocks[0]]
# 
#   for i in range(1, nr_blocks + 1):
#     other = random.choice_list(leafs)
#     this = Block(i)
#     this.add_to(other)
#     if not other.placeable or one_stack:
#       leafs.remove(other)
#     blocks.append(this)
#     leafs.append(this)
# 
#   order = None
#   if random_order:
#     order = random.permutation(len(blocks))
# 
#   return BlocksWorld(blocks, random_order=order)
# 
# Path: difflogic/envs/blocksworld/represent.py
# def get_coordinates(world, absolute=False):
#   """Get the coordinates of each block in the blocks world."""
#   coordinates = [None for _ in range(world.size)]
#   raw_blocks = world.blocks.raw
# 
#   def dfs(block):
#     """Use depth-first-search to get the coordinate of each block."""
#     if block.is_ground:
#       coordinates[block.index] = (0, 0)
#       for j, c in enumerate(block.children):
#         # When using absolute coordinate, the block x directly placed on the
#         # ground gets coordinate (x, 1).
#         x = world.blocks.inv_index(c.index) if absolute else j
#         coordinates[c.index] = (x, 1)
#         dfs(c)
#     else:
#       coor = coordinates[block.index]
#       assert coor is not None
#       x, y = coor
#       for c in block.children:
#         coordinates[c.index] = (x, y + 1)
#         dfs(c)
# 
#   dfs(raw_blocks[0])
#   coordinates = world.blocks.permute(coordinates)
#   return np.array(coordinates)
# 
# Path: difflogic/envs/blocksworld/represent.py
# def decorate(state, nr_objects, world_id=None):
#   """Append world index and object index information to state."""
#   info = []
#   if world_id is not None:
#     info.append(np.ones((nr_objects, 1)) * world_id)
#   info.extend([np.array(range(nr_objects))[:, np.newaxis], state])
#   return np.hstack(info)


# Path: difflogic/envs/blocksworld/envs.py
import numpy as np
import jacinle.random as random
from jaclearn.rl.env import SimpleRLEnvBase
from .block import randomly_generate_world
from .represent import get_coordinates
from .represent import decorate

        decorate: if True, the coordinates in the states will also include the
            world index (default: 0) and the block index (starting from 0).
        prob_unchange: The probability that an action is not effective.
        prob_fall: The probability that an action will make the object currently
            moving fall on the ground.
  """

  def __init__(self,
               nr_blocks,
               random_order=False,
               decorate=False,
               prob_unchange=0.0,
               prob_fall=0.0):
    super().__init__()
    self.nr_blocks = nr_blocks
    self.nr_objects = nr_blocks + 1
    self.random_order = random_order
    self.decorate = decorate
    self.prob_unchange = prob_unchange
    self.prob_fall = prob_fall

  def _restart(self):
    self.world = randomly_generate_world(
        self.nr_blocks, random_order=self.random_order)
    self._set_current_state(self._get_decorated_states())
    self.is_over = False
    self.cached_result = self._get_result()

  def _get_decorated_states(self, world_id=0):
    state = get_coordinates(self.world)

    if self.decorate: