1"""# TowerHcraft Environment
  2
  3Simple environment with tower-structured constructor rules
  4to evaluate polynomial sub-behaviors reusability.
  5
  6The goal of the environment is to get the item on top of the tower.
  7
  8The tower has 'height' layers and 'width' items per layer,
  9plus the final item on top of the last layer.
 10
 11Each item in the tower requires all the items of the previous layer to be built.
 12Items of the floor layer require nothing and can be built from the start.
 13
 14## Example
 15
 16For example here is a tower of height 2 and width 3:
 17
 18|   | 6 |   |
 19|:-:|:-:|:-:|
 20| 3 | 4 | 5 |
 21| 0 | 1 | 2 |
 22
 23The goal here is to get the item 6.
 24Item 6 requires the items {3, 4, 5}.
 25Each of the items 3, 4 and 5 requires items {0, 1, 2}.
 26Each of the items 0, 1 and 2 requires nothing and can be crafted from the start.
 27
 28Requirements graph for H2-W3:
 29<div class="graph">
 30.. include:: ../../../docs/images/requirements_graphs/TowerHcraft-H2-W3.html
 31</div>
 32
 33"""
 34
 35from typing import List
 36
 37from hcraft.elements import Item
 38from hcraft.env import HcraftEnv
 39from hcraft.transformation import Transformation, Use, Yield, PLAYER
 40from hcraft.world import world_from_transformations
 41from hcraft.task import GetItemTask
 42
 43try:
 44    import gymnasium as gym
 45
 46    gym.register(
 47        id="TowerHcraft-v1",
 48        entry_point="hcraft.examples.tower:TowerHcraftEnv",
 49    )
 50
 51except ImportError:
 52    pass
 53
 54
 55class TowerHcraftEnv(HcraftEnv):
 56    """Tower, a tower-structured hierarchical Environment.
 57
 58    Item of given layer requires all items of the previous.
 59    The goal is to obtain the last item on top of the tower.
 60
 61    """
 62
 63    def __init__(self, height: int = 2, width: int = 3, **kwargs):
 64        """
 65        Args:
 66            height (int): Number of layers of the tower (ignoring goal item).
 67            width (int): Number of items per layer.
 68        """
 69        self.height = height
 70        self.width = width
 71        n_items = self.height * self.width + 1
 72        self.items = [Item(str(i)) for i in range(n_items)]
 73        name = f"TowerHcraft-H{self.height}-W{self.width}"
 74        if not isinstance(kwargs.get("max_step"), int):
 75            if self.width == 1:
 76                kwargs["max_step"] = 1 + int(self.width * (self.height + 1))
 77            else:
 78                # 1 + w + w**2 + ... + w**h
 79                kwargs["max_step"] = 1 + int(
 80                    (1 - self.width ** (self.height + 1)) / (1 - self.width)
 81                )
 82        transformations = self.build_transformations(self.items)
 83        world = world_from_transformations(transformations)
 84        if "purpose" not in kwargs:
 85            kwargs["purpose"] = GetItemTask(self.items[-1])
 86        super().__init__(world, name=name, **kwargs)
 87
 88    def build_transformations(self, items: List[Item]) -> List[Transformation]:
 89        """Build transformations to make every item accessible.
 90
 91        Args:
 92            items: List of items.
 93
 94        Returns:
 95            List of craft recipes as transformations.
 96
 97        """
 98        transformations = []
 99
100        # First layer recipes
101        for first_layer_id in range(self.width):
102            item = items[first_layer_id]
103            new_recipe = Transformation(inventory_changes=[Yield(PLAYER, item)])
104            transformations.append(new_recipe)
105
106        # Tower recipes
107        for layer in range(1, self.height):
108            for item_layer_id in range(self.width):
109                item_id = layer * self.width + item_layer_id
110                item = items[item_id]
111
112                inventory_changes = [Yield(PLAYER, item)]
113
114                prev_layer_id = (layer - 1) * self.width
115                for prev_item_id in range(self.width):
116                    required_item = items[prev_layer_id + prev_item_id]
117                    inventory_changes.append(Use(PLAYER, required_item, consume=1))
118
119                new_recipe = Transformation(inventory_changes=inventory_changes)
120                transformations.append(new_recipe)
121
122        # Last item recipe
123        last_item = items[-1]
124        inventory_changes = [Yield(PLAYER, last_item)]
125        last_layer_id = (self.height - 1) * self.width
126        for prev_item_id in range(self.width):
127            required_item = items[last_layer_id + prev_item_id]
128            inventory_changes.append(Use(PLAYER, required_item, consume=1))
129
130        new_recipe = Transformation(inventory_changes=inventory_changes)
131        transformations.append(new_recipe)
132
133        return transformations