1from typing import TYPE_CHECKING, Dict, Optional
  2
  3import numpy as np
  4
  5from hcraft.transformation import InventoryOwner
  6
  7if TYPE_CHECKING:
  8    from hcraft.world import World
  9    from hcraft.elements import Zone, Item
 10
 11
 12class HcraftState:
 13    """State manager of HierarchyCraft environments.
 14
 15    The state of every HierarchyCraft environment is composed of three parts:
 16    * The player's inventory: `state.player_inventory`
 17    * The one-hot encoded player's position: `state.position`
 18    * All zones inventories: `state.zones_inventories`
 19
 20    The mapping of items, zones, and zones items to their respective indexes is done through
 21    the given World. (See `hcraft.world`)
 22
 23    ![hcraft state](../../docs/images/hcraft_state.png)
 24
 25    """
 26
 27    def __init__(self, world: "World") -> None:
 28        """
 29        Args:
 30            world: World to build the state for.
 31        """
 32        self.player_inventory = np.array([], dtype=np.int32)
 33        self.position = np.array([], dtype=np.int32)
 34        self.zones_inventories = np.array([], dtype=np.int32)
 35
 36        self.discovered_items = np.array([], dtype=np.ubyte)
 37        self.discovered_zones = np.array([], dtype=np.ubyte)
 38        self.discovered_zones_items = np.array([], dtype=np.ubyte)
 39        self.discovered_transformations = np.array([], dtype=np.ubyte)
 40
 41        self.world = world
 42        self.reset()
 43
 44    @property
 45    def current_zone_inventory(self) -> np.ndarray:
 46        """Inventory of the zone where the player is."""
 47        if self.position.shape[0] == 0:
 48            return np.array([])  # No Zone
 49        return self.zones_inventories[self._current_zone_slot, :][0]
 50
 51    @property
 52    def observation(self) -> np.ndarray:
 53        """The player's observation is a subset of the state.
 54
 55        Only the inventory of the current zone is shown.
 56
 57        ![hcraft state](../../docs/images/hcraft_observation.png)
 58
 59        """
 60        return np.concatenate(
 61            (
 62                self.player_inventory,
 63                self.position,
 64                self.current_zone_inventory,
 65            )
 66        )
 67
 68    def amount_of(self, item: "Item", owner: Optional["Zone"] = "player") -> int:
 69        """Current amount of the given item owned by owner.
 70
 71        Args:
 72            item: Item to get the amount of.
 73            owner: Owner of the inventory to check. Defaults to player.
 74
 75        Returns:
 76            int: Amount of the item in the owner's inventory.
 77        """
 78
 79        if owner in self.world.zones:
 80            zone_index = self.world.zones.index(owner)
 81            zone_item_index = self.world.zones_items.index(item)
 82            return int(self.zones_inventories[zone_index, zone_item_index])
 83
 84        item_index = self.world.items.index(item)
 85        return int(self.player_inventory[item_index])
 86
 87    def has_discovered(self, zone: "Zone") -> bool:
 88        """Whether the given zone was discovered.
 89
 90        Args:
 91            zone (Zone): Zone to check.
 92
 93        Returns:
 94            bool: True if the zone was discovered.
 95        """
 96        zone_index = self.world.zones.index(zone)
 97        return bool(self.discovered_zones[zone_index])
 98
 99    @property
100    def current_zone(self) -> Optional["Zone"]:
101        """Current position of the player."""
102        if self.world.n_zones == 0:
103            return None
104        return self.world.zones[self._current_zone_slot[0]]
105
106    @property
107    def _current_zone_slot(self) -> int:
108        return self.position.nonzero()[0]
109
110    @property
111    def player_inventory_dict(self) -> Dict["Item", int]:
112        """Current inventory of the player."""
113        return self._inv_as_dict(self.player_inventory, self.world.items)
114
115    @property
116    def zones_inventories_dict(self) -> Dict["Zone", Dict["Item", int]]:
117        """Current inventories of the current zone and each zone containing item."""
118        zones_invs = {}
119        for zone_slot, zone_inv in enumerate(self.zones_inventories):
120            zone = self.world.zones[zone_slot]
121            zone_inv = self._inv_as_dict(zone_inv, self.world.zones_items)
122            if zone_slot == self._current_zone_slot or zone_inv:
123                zones_invs[zone] = zone_inv
124        return zones_invs
125
126    def apply(self, action: int) -> bool:
127        """Apply the given action to update the state.
128
129        Args:
130            action (int): Index of the transformation to apply.
131
132        Returns:
133            bool: True if the transformation was applied succesfuly. False otherwise.
134        """
135        choosen_transformation = self.world.transformations[action]
136        if not choosen_transformation.is_valid(self):
137            return False
138        choosen_transformation.apply(
139            self.player_inventory,
140            self.position,
141            self.zones_inventories,
142        )
143        self._update_discoveries(action)
144        return True
145
146    def reset(self) -> None:
147        """Reset the state to it's initial value."""
148        self.player_inventory = np.zeros(self.world.n_items, dtype=np.int32)
149        for stack in self.world.start_items:
150            item_slot = self.world.items.index(stack.item)
151            self.player_inventory[item_slot] = stack.quantity
152
153        self.position = np.zeros(self.world.n_zones, dtype=np.int32)
154        start_slot = 0  # Start in first Zone by default
155        if self.world.start_zone is not None:
156            start_slot = self.world.slot_from_zone(self.world.start_zone)
157        if self.position.shape[0] > 0:
158            self.position[start_slot] = 1
159
160        self.zones_inventories = np.zeros(
161            (self.world.n_zones, self.world.n_zones_items), dtype=np.int32
162        )
163        for zone, zone_stacks in self.world.start_zones_items.items():
164            zone_slot = self.world.slot_from_zone(zone)
165            for stack in zone_stacks:
166                item_slot = self.world.zones_items.index(stack.item)
167                self.zones_inventories[zone_slot, item_slot] = stack.quantity
168
169        self.discovered_items = np.zeros(self.world.n_items, dtype=np.ubyte)
170        self.discovered_zones_items = np.zeros(self.world.n_zones_items, dtype=np.ubyte)
171        self.discovered_zones = np.zeros(self.world.n_zones, dtype=np.ubyte)
172        self.discovered_transformations = np.zeros(
173            len(self.world.transformations), dtype=np.ubyte
174        )
175        self._update_discoveries()
176
177    def _update_discoveries(self, action: Optional[int] = None) -> None:
178        self.discovered_items = np.bitwise_or(
179            self.discovered_items, self.player_inventory > 0
180        )
181        self.discovered_zones_items = np.bitwise_or(
182            self.discovered_zones_items, self.current_zone_inventory > 0
183        )
184        self.discovered_zones = np.bitwise_or(self.discovered_zones, self.position > 0)
185        if action is not None:
186            self.discovered_transformations[action] = 1
187
188    @staticmethod
189    def _inv_as_dict(inventory_array: np.ndarray, obj_registry: list):
190        return {
191            obj_registry[index]: value
192            for index, value in enumerate(inventory_array)
193            if value > 0
194        }
195
196    def as_dict(self) -> dict:
197        state_dict = {
198            "pos": self.current_zone,
199            InventoryOwner.PLAYER.value: self.player_inventory_dict,
200        }
201        state_dict.update(self.zones_inventories_dict)
202        return state_dict