# Copyright (c) 2020, salesforce.com, inc.
# All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# For full license text, see the LICENSE file in the repo root
# or https://opensource.org/licenses/BSD-3-Clause

import random
from abc import ABC, abstractmethod
from copy import deepcopy

import numpy as np

from ai_economist.foundation.agents import agent_registry
from ai_economist.foundation.base.registrar import Registry
from ai_economist.foundation.base.world import World
from ai_economist.foundation.components import component_registry
from ai_economist.foundation.entities import (
    endogenous_registry,
    landmark_registry,
    resource_registry,
)


class BaseEnvironment(ABC):
    """
    Base Environment class. Should be used as the parent class for Scenario classes.
    Instantiates world, agent, and component objects.

    Provides Gym-style API for resetting and stepping:
        obs                  <-- env.reset()
        obs, rew, done, info <-- env.step(actions)

    Also provides Gym-style API for controlling random behavior:
        env.seed(seed) # Sets numpy and built-in RNG seeds to seed

    Reference: OpenAI Gym [https://github.com/openai/gym]

    Environments in this framework are instances of Scenario classes (which are built
    as extensions of BaseEnvironment). A Scenario must implement the following
    abstract methods (method docstrings provide detailed explanations):
        reset_starting_layout
        reset_agent_states
        scenario_step
        generate_observations
        compute_reward

    Scenario classes define their own passive dynamics--that is, dynamics that do not
    depend on agent actions--and supply observations. It is up to the Scenario class
    to handle reward.

    Interactions with the environment are handled through components, which define
    actions that agents can perform. Components are defined through distinct
    Component classes (which extend BaseComponent [see base_component.py]) and must
    be included in the components_registry in order to be used (see below).
    Components influence the environment dynamics through effects they have on
    agent/world states. They also (optionally) supply observations.

    The actions available to the agents, observations they receive, the dynamics of
    the environment, and the rewards depend of the choice of which Scenario class and
    Component class(es) to use.

    In multi_action_mode, an agent may choose an action for each of the action
    subspaces defined by the included Component classes. A Component can define 0, 1,
    or several action subspaces for a given agent type. If not using
    multi_action_mode, these action subspaces are combined into a single action space
    and the agent may select one action within this aggregated space.

    For additional detail regarding actions and action subspaces, see the
    BaseComponent class in base_component.py.

    There are 2 types of agents: mobile agents and the planner agent. There can be
    two or more mobile agents and a single planner agent. Conceptually, mobile agents
    represent the individual actors in the economic simulation while the planner
    agent represents a social planner that sets macroeconomic policy.

    This environment framework makes extensive use of Python classes. Scenarios,
    Components, Agents, and environment entities such as Resources, Landmarks,
    and Endogenous variables are all implemented as classes. These classes are
    accessed via registries. See top example.

    Example:
        from ai_economist import foundation
        # foundation.scenarios  <-- Scenario class registry
        # foundation.components <-- Component class registry
        # foundation.agents     <-- Agent class registry
        # foundation.resources  <-- Resource class registry
        # foundation.landmarks  <-- Landmark class registry
        # foundation.endogenous <-- Endogenous class registry

        # see ../scenarios/simple_wood_and_stone/dynamic_layout.py
        UniScenarioClass = foundation.scenarios.get("uniform/simple_wood_and_stone")

        # see ../components/build.py and ../components/move.py
        BuildComponentClass  = foundation.components.get("Build")
        GatherComponentClass = foundation.components.get("Gather")

    Example:
        from ai_economist import foundation
        from ai_economist.foundation.base.base_env import BaseEnvironment

        ScenarioClass = foundation.scenarios.get(...)
        assert issubclass(ScenarioClass, BaseEnvironment)

        env = ScenarioClass(
            components=[
                ("Build", {"payment": 20}),
                ("Gather", {"move_labor": 1.0, "collect_labor": 2.0}),
            ],
            n_agents=20,
            world_size=[25, 25],
        )

        obs = env.reset()

        actions = {agent.idx: ... for agent in env.all_agents}
        obs, rew, done, info = env.step(actions)

    Args:
        components (list): A list of tuples ("Component Name", {Component kwargs}) or
            list of dicts {"Component Name": {Component kwargs}} specifying the
            components that the instantiated environment will include.
            "Component Name" must be a string matching the name of a registered
            Component class.
            {Component kwargs} must be a dictionary of kwargs that can be passed as
            arguments to the Component class with name "Component Name".
            Resetting, stepping, and observation generation will be carried out in
            the order in which components are listed. This should be considered,
            as re-ordering the components list may impact the dynamics of the
            environment.
        agent_composition (dict): Agent Class name in string paired with amount of agents of that class
        n_agents (int): The number of mobile agents (does not include planner).
            Number of agents must be > 1.
        world_size (list): A length-2 list specifying the dimensions of the 2D world.
            Interpreted as [height, width].
        episode_length (int): Number of timesteps in a single episode.
        multi_action_mode_agents (bool): Whether mobile agents use multi_action_mode.
        multi_action_mode_planner (bool): Whether the planner uses multi_action_mode.
        flatten_observations (bool): Whether to preprocess observations by
            concatenating all scalar/vector observation subfields into a single
            "flat" observation field. If not, return observations as minimally
            processed dictionaries.
        flatten_masks (bool): Whether to flatten action masks into a single array or
            to keep as a {"action_subspace_name": action_subspace_mask} dictionary.
            For integration with deep RL, it is helpful to set this to True, for the
            purpose of action masking: flattened masks have the same semantics as
            policy logits.
        allow_observation_scaling (bool): Whether to enable certain observation
            fields to be scaled to a range better suited for deep RL.
        dense_log_frequency (int): [optional] How often (in completed episodes) to
            create a dense log while playing an episode. By default, dense logging is
            turned off (dense_log_frequency=None). If dense_log_frequency=20,
            a dense log will be created when the total episode count is a multiple of
            20.
            Dense logs provide a log of agent states, actions, and rewards at each
            timestep of an episode. They also log world states at a coarser timescale
            (see below). Component classes optionally contribute additional
            information to the dense log.
            Note: dense logging is time consuming (especially with many agents).
        world_dense_log_frequency (int): When dense logging, how often (in timesteps) to
            log a snapshot of the world state. If world_dense_log_frequency=50
            (the default), the world state will be included in the dense log for
            timesteps where t is a multiple of 50.
            Note: More frequent world snapshots increase the dense log memory footprint.
        seed (int, optional): If provided, sets the numpy and built-in random number
            generator seeds to seed. You can control the seed after env construction
            using the 'seed' method.
    """

    # The name associated with this Scenario class (must be unique)
    # Note: This is what will identify the Scenario class in the scenario registry.
    name = ""

    # The (sub)classes of agents that this scenario applies to
    agent_subclasses = []

    # The (non-agent) game entities that are expected to be in play
    required_entities = None  # Replace with list or tuple (can be empty)

    def __init__(
        self,
        components=None,
        agent_composition=None,
        world_size=None,
        episode_length=1000,
        multi_action_mode_agents=False,
        multi_action_mode_planner=True,
        flatten_observations=True,
        flatten_masks=True,
        allow_observation_scaling=True,
        dense_log_frequency=None,
        world_dense_log_frequency=50,
        collate_agent_step_and_reset_data=False,
        seed=None,
    ):

        # Make sure a name was declared by child class
        assert self.name
        # Make sure the agent_subclasses was declared by child class
        # and does not create potential conflicts
        assert isinstance(self.agent_subclasses, (tuple, list))
        assert len(self.agent_subclasses) > 0
        if len(self.agent_subclasses) > 1:
            for i in range(len(self.agent_subclasses)):
                for j in range(len(self.agent_subclasses)):
                    if i == j:
                        continue
                    a_i = agent_registry.get(self.agent_subclasses[i])
                    a_j = agent_registry.get(self.agent_subclasses[j])
                    assert not issubclass(a_i, a_j)

        # Make sure the required_entities was declared by child class
        # (will typecheck later)
        assert isinstance(self.required_entities, (tuple, list))

        # World size must be a tuple or list of length 2,
        # specifying [Height, Width] of the game map
        assert isinstance(world_size, (tuple, list))
        assert len(world_size) == 2
        self.world_size = world_size

        # Set n_agents
        self.agent_composition=agent_composition
        n_agents=0
        for k,v in agent_composition.items():
            n_agents+=v

        # Number of agents must be an integer and there must be at least 2 agents
        assert isinstance(n_agents, int)
        assert n_agents >= 2
        self.n_agents = n_agents

        # Foundation assumes there's only a single planner
        n_planners = 1
        self.num_agents = (
            n_agents + n_planners
        )  # used in the warp_drive env wrapper (+ 1 for the planner)
        
        # Components must be a tuple/list where each element is either a...
        #   tuple: ('Component Name', {Component kwargs})
        #   dict : {'Component Name': {Component kwargs}}
        assert isinstance(components, (tuple, list))

        def spec_is_valid(spec):
            """Return True if component specification is validly configured."""
            if isinstance(spec, (tuple, list)):
                if len(spec) != 2:
                    return False
                return isinstance(spec[0], str) and isinstance(spec[1], dict)
            if isinstance(spec, dict):
                if len(spec) != 1:
                    return False
                key_is_str = isinstance(list(spec.keys())[0], str)
                val_is_dict = isinstance(list(spec.values())[0], dict)
                return key_is_str and val_is_dict
            return False

        assert all(spec_is_valid(component) for component in components)

        self._episode_length = int(episode_length)
        assert self._episode_length >= 1

        # Can an agent/planner execute multiple actions (1 per action subspace) per
        # timestep (=True) or just one action (=False)
        self.multi_action_mode_agents = bool(multi_action_mode_agents)
        self.multi_action_mode_planner = bool(multi_action_mode_planner)

        # Whether to allow the world to scale observations
        self._allow_observation_scaling = bool(allow_observation_scaling)

        # Whether to flatten the observation dictionaries before returning them
        # Note: flattened observations are still returned as dictionaries, but with
        # all scalar/vector observation fields concatenated into a single "flat" field.
        self._flatten_observations = bool(flatten_observations)

        # Whether to flatten the mask dictionaries before putting them in the obs
        self._flatten_masks = bool(flatten_masks)

        # How often (in episode completions) to create a dense log
        self._dense_log_this_episode = False
        if dense_log_frequency is None:  # Only create a dense log
            # if manually specified during reset
            self._create_dense_log_every = None
        else:  # Create a dense log every dense_log_frequency episodes
            self._create_dense_log_every = int(dense_log_frequency)
            assert self._create_dense_log_every >= 1

        # How often (in timesteps) to snapshot the world map when creating the denselog
        self._world_dense_log_frequency = int(world_dense_log_frequency)
        assert self._world_dense_log_frequency >= 1

        # Seed control
        if seed is not None:
            self.seed(seed)

        # Initialize the set of entities used in the game that's being created.
        # Coin and Labor are always included.
        self._entities = {
            "resources": ["Coin"],
            "landmarks": [],
            "endogenous": ["Labor"],
        }
        self._register_entities(self.required_entities)

        # Register all the components to get the entities they rely on.
        self._components = []
        self._components_dict = {}
        self._shorthand_lookup = {}
        component_classes = []
        for component_spec in components:
            if isinstance(component_spec, (tuple, list)):
                component_name, component_config = component_spec
            elif isinstance(component_spec, dict):
                assert len(component_spec) == 1
                component_name = list(component_spec.keys())[0]
                component_config = list(component_spec.values())[0]
            else:
                raise TypeError
            component_cls = component_registry.get(component_name)
            self._register_entities(component_cls.required_entities)
            component_classes.append([component_cls, component_config])

        # Initialize the world object (contains agents and world map),
        # now that we know all the entities we'll use.
        self.world = World(
            self.world_size,
            self.agent_composition,
            self.resources,
            self.landmarks,
            self.multi_action_mode_agents,
            self.multi_action_mode_planner,
        )

        # Initialize the component objects.
        for component_cls, component_kwargs in component_classes:
            component_object = component_cls(
                self.world,
                self._episode_length,
                inventory_scale=self.inv_scale,
                **component_kwargs
            )
            self._components.append(component_object)
            self._components_dict[component_object.name] = component_object
            self._shorthand_lookup[component_object.shorthand] = component_object

       
        self.world.planner.register_inventory(self.resources)
        self.world.planner.register_components(self._components)
        self.reapply_scenario_config_to_agents()
       

        self._completions = 0
        self._finish_episode=False
        self._last_ep_metrics = None

        # For dense logging
        self._dense_log = {"world": [], "states": [], "actions": [], "rewards": []}
        self._last_ep_dense_log = self.dense_log.copy()

        # For episode replay
        self._replay_log = {"reset": dict(seed_state=None), "step": []}
        self._last_ep_replay_log = self.replay_log.copy()

        self._packagers = {}

        # To collate all the agents ('0', '1', ...) data during reset and step
        # into a single agent with index 'a'
        self.collate_agent_step_and_reset_data = collate_agent_step_and_reset_data
    
    def reapply_scenario_config_to_agents(self):
        # Register the components with the agents
        # to finish setting up their state/action spaces.
        for agent in self.world.agents:
            agent.register_inventory(self.resources)
            agent.register_endogenous(self.endogenous)
            agent.register_components(self._components)
        self._agent_lookup = {str(agent.idx): agent for agent in self.all_agents}
        self.world.apply_agent_db_to_world()

    def _register_entities(self, entities):
        for entity in entities:
            if resource_registry.has(entity):
                if entity not in self._entities["resources"]:
                    self._entities["resources"].append(entity)
            elif landmark_registry.has(entity):
                if entity not in self._entities["landmarks"]:
                    self._entities["landmarks"].append(entity)
            elif endogenous_registry.has(entity):
                if entity not in self._entities["endogenous"]:
                    self._entities["endogenous"].append(entity)
            else:
                raise KeyError("Unknown entity: {}".format(entity))

    # Properties
    # ----------

    @property
    def episode_length(self):
        """Length of an episode, in timesteps."""
        return int(self._episode_length)

    @property
    def inv_scale(self):
        """Scale value to be used for inventory scaling. 1 if no scaling enabled."""
        return 0.01 if self._allow_observation_scaling else 1

    @property
    def resources(self):
        """List of resources managed by this environment instance."""
        return sorted(list(self._entities["resources"]))

    @property
    def landmarks(self):
        """List of landmarks managed by this environment instance."""
        return sorted(list(self._entities["landmarks"]))

    @property
    def endogenous(self):
        """List of endogenous quantities managed by this environment instance."""
        return sorted(list(self._entities["endogenous"]))

    @property
    def all_agents(self):
        """List of mobile agents and the planner agent."""
        return self.world.agents + [self.world.planner]

    @property
    def previous_episode_metrics(self):
        """Metrics from the end of the last completed episode."""
        return self._last_ep_metrics

    @property
    def metrics(self):
        """The combined metrics yielded by the scenario and the components."""
        metrics = self.scenario_metrics() or {}

        for component in self._components:
            m_metrics = component.get_metrics()
            if not m_metrics:
                continue
            for k, v in m_metrics.items():
                metrics["{}/{}".format(component.shorthand, k)] = v

        return metrics

    @property
    def components(self):
        """The list of components associated with this scenario."""
        return self._components

    @property
    def dense_log(self):
        """The contents of the current (potentially incomplete) dense log."""
        return self._dense_log

    @property
    def replay_log(self):
        """The contents of the current (potentially incomplete) replay log."""
        return self._replay_log

    @property
    def previous_episode_dense_log(self):
        """Dense log from the last completed episode that was being logged."""
        return self._last_ep_dense_log

    @property
    def previous_episode_replay_log(self):
        """
        Replay log from the last completed episode. Serves as a compact encoding of
        an episode by allowing the episode to be perfectly reproduced.

        Examples:
            # replay log of the episode to be reproduced
            replay_log = env.previous_episode_replay_log

            # recover episode metrics and dense log via replay
            _ = env.reset(force_dense_logging=True, **replay_log['reset'])
            for replay_step in replay_log['step']:
                _ = env.step(**replay_step)
            dense_log = env.previous_episode_dense_log
            metrics = env.previous_episode_metrics
        """
        return self._last_ep_replay_log

    @property
    def generate_rewards(self):
        """Compute the rewards for each agent."""
        return self._generate_rewards

    # Seed control
    # -----------------

    @staticmethod
    def seed(seed):
        """Sets the numpy and built-in random number generator seed.

        Args:
            seed (int, float): Seed value to use. Must be > 0. Converted to int
                internally if provided value is a float.
        """
        assert isinstance(seed, (int, float))
        seed = int(seed)
        assert seed > 0

        np.random.seed(seed)
        random.seed(seed)

    # Getters & Setters
    # -----------------
    def set_finish_episode(self,done):
        self._finish_episode=done

    def get_component(self, component_name):
        """
        Get the component object instance wrapped in the environment.

        Args:
            component_name (str): Name or shorthand name of the Component class to get.
                Must correspond to a name or shorthand of one of the components that
                is included in this environment instance.

        Returns:
            component (BaseComponent object)
        """
        if component_name not in self._components_dict:
            if component_name not in self._shorthand_lookup:
                raise KeyError(
                    "No component with name or shorthand name {} found; "
                    "registered components are:\n".format(component_name)
                    + "\n\t".join(list(self._components_dict.keys()))
                )
            return self._shorthand_lookup[component_name]
        return self._components_dict[component_name]

    def get_agent(self, agent_idx):
        """
        Get the agent object instance with idx agent_idx.

        Args:
            agent_idx (int or str): Identifier of the agent to return. Must match the
                idx property of one of the agent objects in self.all_agents.

        Returns:
            agent (BaseAgent object)
        """
        agent = self._agent_lookup.get(str(agent_idx), None)
        if agent is None:
            raise ValueError("No agent with associated index {}".format(agent_idx))
        return agent

    def set_agent_component_action(self, agent_idx, component_name, action):
        """
        Set agent with idx <agent_idx> to take action <action> for the action
        subspace with name <component_name>

        Args:
            agent_idx (int or str): Identifier of the agent taking the action. Must
                match the idx property of one of the agent objects in self.all_agents.
            component_name (str): Name of the action subspace to set the action value
                of.
            action (int): Index of the chosen action.
        """
        agent = self.get_agent(agent_idx)
        agent.set_component_action(component_name, action)

    def parse_actions(self, action_dictionary):
        """Put actions into the appropriate agent's action buffer"""
        for agent_idx, agent_actions in action_dictionary.items():
            agent = self.get_agent(agent_idx)
            agent.parse_actions(agent_actions)

    # Core control of environment execution
    # -------------------------------------

    @staticmethod
    def _build_packager(sub_obs, put_in_both=None):
        """
        Decides which keys-vals should be flattened or not.
        put_in_both: include in both (e.g., 'time')
        """
        if put_in_both is None:
            put_in_both = []
        keep_as_is = []
        flatten = []
        wrap_as_list = {}
        for k, v in sub_obs.items():
            if isinstance(v, np.ndarray):
                multi_d_array = len(v.shape) > 1
            else:
                multi_d_array = False

            if k == "action_mask" or multi_d_array:
                keep_as_is.append(k)
            else:
                flatten.append(k)
                if k in put_in_both:
                    keep_as_is.append(k)

            wrap_as_list[k] = np.isscalar(v)

        flatten = sorted(flatten)

        return keep_as_is, flatten, wrap_as_list

    @staticmethod
    def _package(obs_dict, keep_as_is, flatten, wrap_as_list):
        new_obs = {k: obs_dict[k] for k in keep_as_is}
        if len(flatten) == 1:
            k = flatten[0]
            o = obs_dict[k]
            if wrap_as_list[k]:
                o = [o]
            new_obs["flat"] = np.array(o, dtype=np.float32)
        else:
            to_flatten = [
                [obs_dict[k]] if wrap_as_list[k] else obs_dict[k] for k in flatten
            ]
            try:
                new_obs["flat"] = np.concatenate(to_flatten).astype(np.float32)
            except ValueError:
                for k, v in zip(flatten, to_flatten):
                    print(k, np.array(v).shape)
                    print(v)
                    print("")
                raise
        return new_obs

    def _generate_observations(self, flatten_observations=False, flatten_masks=False):
        def recursive_listify(d):
            assert isinstance(d, dict)
            for k, v in d.items():
                if isinstance(v, dict):
                    d[k] = recursive_listify(v)
                elif isinstance(v, (int, float)):
                    d[k] = v
                elif isinstance(v, (list, tuple, set)):
                    d[k] = list(v)
                elif isinstance(v, (np.ndarray, np.integer, np.floating)):
                    d[k] = v.tolist()
                else:
                    raise NotImplementedError(
                        "Not clear how to handle {} with type {}".format(k, type(v))
                    )
                if isinstance(d[k], list) and len(d[k]) == 1:
                    d[k] = d[k][0]
            return d

        # Initialize empty observations
        if self.collate_agent_step_and_reset_data:
            obs = {"a": {}, "p": {}}
        else:
            obs = {str(agent.idx): {} for agent in self.all_agents}
        agent_wise_planner_obs = {
            "p" + str(agent.idx): {} for agent in self.world.agents
        }

        # Get/process observations generated by the scenario
        world_obs = {str(k): v for k, v in self.generate_observations().items()}
        time_scale = self.episode_length if self._allow_observation_scaling else 1.0
        for idx, o in world_obs.items():
            if idx in obs:
                obs[idx].update({"world-" + k: v for k, v in o.items()})
                if self.collate_agent_step_and_reset_data and idx == "a":
                    obs[idx]["time"] = np.array(
                        [
                            self.world.timestep / time_scale
                            for _ in range(self.world.n_agents)
                        ]
                    )
                else:
                    obs[idx]["time"] = [self.world.timestep / time_scale]
            elif idx in agent_wise_planner_obs:
                agent_wise_planner_obs[idx].update(
                    {"world-" + k: v for k, v in o.items()}
                )
            else:
                raise KeyError

        # Get/process observations generated by the components
        for component in self._components:
            for idx, o in component.obs().items():
                if idx in obs:
                    obs[idx].update({component.name + "-" + k: v for k, v in o.items()})
                elif idx in agent_wise_planner_obs:
                    agent_wise_planner_obs[idx].update(
                        {component.name + "-" + k: v for k, v in o.items()}
                    )
                else:
                    raise KeyError

        # Process the observations
        if flatten_observations:
            for o_dict in [obs, agent_wise_planner_obs]:
                for aidx, aobs in o_dict.items():
                    if not aobs:
                        continue
                    if aidx not in self._packagers:
                        self._packagers[aidx] = self._build_packager(
                            aobs, put_in_both=["time"]
                        )
                    try:
                        o_dict[aidx] = self._package(aobs, *self._packagers[aidx])
                    except ValueError:
                        print("Error when packaging obs.")
                        print("Agent index: {}\nRaw obs: {}\n".format(aidx, aobs))
                        raise

        for k, v in agent_wise_planner_obs.items():
            if len(v) > 0:
                obs[self.world.planner.idx][k] = (
                    v["flat"] if flatten_observations else v
                )

        # Get each agent's action masks and incorporate them into the observations
        for aidx, amask in self._generate_masks(flatten_masks=flatten_masks).items():
            obs[aidx]["action_mask"] = amask

        return obs

    def _generate_masks(self, flatten_masks=True):
        if self.collate_agent_step_and_reset_data:
            masks = {"a": {}, "p": {}}
        else:
            masks = {agent.idx: {} for agent in self.all_agents}
        for component in self._components:
            # Use the component's generate_masks method to get action masks
            component_masks = component.generate_masks(completions=self._completions)

            for idx, mask in component_masks.items():
                if isinstance(mask, dict):
                    for sub_action, sub_mask in mask.items():
                        masks[idx][
                            "{}.{}".format(component.name, sub_action)
                        ] = sub_mask
                else:
                    masks[idx][component.name] = mask

        if flatten_masks:
            if self.collate_agent_step_and_reset_data:
                flattened_masks = {}
                for agent_id in masks.keys():
                    if agent_id == "a":
                        multi_action_mode = self.multi_action_mode_agents
                        no_op_mask = np.ones((1, self.n_agents))
                    elif agent_id == "p":
                        multi_action_mode = self.multi_action_mode_planner
                        no_op_mask = [1]
                    mask_dict = masks[agent_id]
                    list_of_masks = []
                    if not multi_action_mode:
                        list_of_masks.append(no_op_mask)
                    for m in mask_dict.keys():
                        if multi_action_mode:
                            list_of_masks.append(no_op_mask)
                        list_of_masks.append(mask_dict[m])
                    flattened_masks[agent_id] = np.concatenate(
                        list_of_masks, axis=0
                    ).astype(np.float32)
                return flattened_masks
            return {
                str(agent.idx): agent.flatten_masks(masks[agent.idx])
                for agent in self.all_agents
            }
        return {
            str(agent_idx): {
                k: np.array(v, dtype=np.uint8).tolist()
                for k, v in masks[agent_idx].items()
            }
            for agent_idx in list(masks.keys())
        }

    def _generate_rewards(self):
        rew = self.compute_reward()
        assert isinstance(rew, dict)
        return {str(k): v for k, v in rew.items()}

    def _finalize_logs(self):
        self._last_ep_replay_log = self._replay_log
        self._last_ep_metrics = self.metrics

        if not self._dense_log_this_episode:
            return

        def recursive_cast(d):
            if isinstance(d, (list, tuple, set)):
                new_d = [recursive_cast(v_) for v_ in d]
                return new_d
            if isinstance(d, dict):
                for k, v in d.items():
                    if isinstance(v, (list, tuple, set, dict)):
                        d[k] = recursive_cast(v)
                    elif isinstance(v, (int, float, str)):
                        d[k] = v
                    elif isinstance(v, (np.ndarray, np.integer, np.floating)):
                        d[k] = v.tolist()
                    else:
                        raise NotImplementedError(
                            "Not clear how to handle {} with type {}".format(k, type(v))
                        )
                return d
            if isinstance(d, (int, float, str)):
                return d
            if isinstance(d, (np.ndarray, np.integer, np.floating)):
                return d.tolist()
            raise NotImplementedError(
                "Not clear how to handle {} with type {}".format(d, type(d))
            )

        self._dense_log["world"].append(deepcopy(self.world.maps.state_dict))
        self._dense_log["states"].append(
            {str(agent.idx): deepcopy(agent.state) for agent in self.all_agents}
        )

        # Back-fill the log with each component's dense log to complete the aggregate
        # dense log
        for component in self._components:
            component_log = component.get_dense_log()
            if component_log is None:
                continue
            if isinstance(component_log, dict):
                for k, v in component_log.items():
                    self._dense_log[component.shorthand + "-" + k] = v
            elif isinstance(component_log, (tuple, list)):
                self._dense_log[component.shorthand] = list(component_log)
            else:
                raise TypeError

        self._last_ep_dense_log = recursive_cast(self._dense_log)

    def collate_agent_obs(self, obs):
        # Collating observations from all agents
        if "a" in obs:  # already collated!
            return obs
        num_agents = len(obs.keys()) - 1
        obs["a"] = {}
        for key in obs["0"].keys():
            obs["a"][key] = np.stack(
                [obs[str(agent_idx)][key] for agent_idx in range(num_agents)], axis=-1
            )
        for agent_idx in range(num_agents):
            del obs[str(agent_idx)]
        return obs

    def collate_agent_rew(self, rew):
        # Collating rewards from all agents
        if "a" in rew:  # already collated!
            return rew
        num_agents = len(rew.keys()) - 1
        rew["a"] = []
        for agent_idx in range(num_agents):
            rew["a"] += [rew[str(agent_idx)]]
            del rew[str(agent_idx)]
        return rew

    def collate_agent_info(self, info):
        # Collating infos from all agents
        if "a" in info:  # already collated!
            return info
        num_agents = len(info.keys()) - 1
        info["a"] = {}
        for agent_idx in range(num_agents):
            info["a"][str(agent_idx)] = info[str(agent_idx)]
            del info[str(agent_idx)]
        return info

    def reset(self, seed_state=None, force_dense_logging=False):
        """
        Reset the state of the environment to initialize a new episode.

        Arguments:
            seed_state (tuple or list): Optional state that the numpy RNG should be set
                to prior to the reset cycle must be length 5, following the format
                expected by np.random.set_state()
            force_dense_logging (bool): Optional whether to force dense logging to take
                place this episode; default behavior is to do dense logging every
                create_dense_log_every episodes

        Returns:
            obs (dict): A dictionary of {"agent_idx": agent_obs} with an entry for
                each agent receiving observations. The "agent_idx" key identifies the
                agent receiving the observations in the associated agent_obs value,
                which itself is a dictionary. The "agent_idx" key matches the
                agent.idx property for the given agent.
        """
        if seed_state is not None:
            assert isinstance(seed_state, (tuple, list))
            assert len(seed_state) == 5
            seed_state = (
                str(seed_state[0]),
                np.array(seed_state[1], dtype=np.uint32),
                int(seed_state[2]),
                int(seed_state[3]),
                float(seed_state[4]),
            )
            np.random.set_state(seed_state)

        if force_dense_logging:
            self._dense_log_this_episode = True
        elif self._create_dense_log_every is None:
            self._dense_log_this_episode = False
        else:
            self._dense_log_this_episode = (
                self._completions % self._create_dense_log_every
            ) == 0

        # For dense logging
        self._dense_log = {"world": [], "states": [], "actions": [], "rewards": []}

        # For episode replay
        self._replay_log = {"reset": dict(seed_state=np.random.get_state()), "step": []}

        # Reset the timestep counter
        self.world.timestep = 0

        # Reset done flag
        self._finish_episode=False

        # Perform the scenario reset,
        # which includes resetting the world and agent states
        self.reset_starting_layout()
        self.reset_agent_states()

        # Perform the component resets for each registered component
        for component in self._components:
            component.reset()

        # Take any customized reset actions
        self.additional_reset_steps()

        # By default, agents take the NO-OP action for each action space.
        # Reset actions to that default.
        for agent in self.all_agents:
            agent.reset_actions()
            agent.set_setup(True)

        # Produce observations
        obs = self._generate_observations(
            flatten_observations=self._flatten_observations,
            flatten_masks=self._flatten_masks,
        )

        if self.collate_agent_step_and_reset_data:
            obs = self.collate_agent_obs(obs)

        return obs

    def step(self, actions=None, seed_state=None):
        """
        Execute the components, perform the scenario step, collect observations and
        return observations, rewards, dones, and infos.

        Arguments:
            actions (dict): dictionary of {agent_idx: action} with an entry for each
                agent (which may include the planner) that is supplying an action.
                The key identifies which agent the action is associated with. It
                should match that agent's agent.idx property.
                The value indicates which action the agent will take. The environment
                supports two formats for specifying an action, with slightly
                different expectations for multi_action_mode.
                If agent.multi_action_mode, action must be a list of integers
                specifying the chosen action for each action subspace.
                Otherwise, action must be a single integer specifying the chosen
                action (where the action space is the concatenation of the subspaces).
            seed_state (tuple or list): Optional state that the numpy RNG should be set
                to prior to the reset cycle must be length 5, following the format
                expected by np.random.set_state().

        Returns:
            obs (dict): A dictionary of {"agent_idx": agent_obs} with an entry for
                each agent receiving observations. The "agent_idx" key identifies the
                agent receiving the observations in the associated agent_obs value,
                which itself is a dictionary. The "agent_idx" key matches the
                agent.idx property for the given agent.
            rew (dict): A dictionary of {"agent_idx": reward} with an entry for each
                agent that also receives an observation. Each reward value is a scalar.
            done (dict): A dictionary with a single key "__all__". The associated
                value is False when self.world.timestep < self.episode_length and True
                otherwise.
            info (dict): Placeholder dictionary with structure {"agent_idx": {}},
                with the same keys as obs and rew.
        """
        if actions is not None:
            assert isinstance(actions, dict)
            self.parse_actions(actions)

        if seed_state is not None:
            assert isinstance(seed_state, (tuple, list))
            assert len(seed_state) == 5
            seed_state = (
                str(seed_state[0]),
                np.array(seed_state[1], dtype=np.uint32),
                int(seed_state[2]),
                int(seed_state[3]),
                float(seed_state[4]),
            )
            np.random.set_state(seed_state)

        self._replay_log["step"].append(
            dict(actions=actions, seed_state=np.random.get_state())
        )

        if self._dense_log_this_episode:
            self._dense_log["world"].append(
                deepcopy(self.world.maps.state_dict)
                if (self.world.timestep % self._world_dense_log_frequency) == 0
                else {}
            )
            self._dense_log["states"].append(
                {str(agent.idx): deepcopy(agent.state) for agent in self.all_agents}
            )
            self._dense_log["actions"].append(
                {
                    str(agent.idx): {k: v for k, v in agent.action.items() if v > 0}
                    for agent in self.all_agents
                }
            )

        self.world.timestep += 1

        for component in self._components:
            component.component_step()

        self.scenario_step()

        obs = self._generate_observations(
            flatten_observations=self._flatten_observations,
            flatten_masks=self._flatten_masks,
        )
        rew = self._generate_rewards()
        done = {"__all__": self.world.timestep >= self._episode_length | self._finish_episode}
        info = {k: {} for k in obs.keys()}

        if self._dense_log_this_episode:
            self._dense_log["rewards"].append(rew)

        for agent in self.all_agents:
            agent.reset_actions()

        if done[
            "__all__"
        ]:  # Complete the dense log and stash it as well as the metrics
            self._finalize_logs()
            self._completions += 1

        if self.collate_agent_step_and_reset_data:
            obs = self.collate_agent_obs(obs)
            rew = self.collate_agent_rew(rew)
            info = self.collate_agent_info(info)

        return obs, rew, done, info

    # The following methods must be implemented for each scenario
    # -----------------------------------------------------------

    @abstractmethod
    def reset_starting_layout(self):
        """
        Part 1/2 of scenario reset. This method handles resetting the state of the
        environment managed by the scenario (i.e. resource & landmark layout).
        """

    @abstractmethod
    def reset_agent_states(self):
        """
        Part 2/2 of scenario reset. This method handles resetting the state of the
        agents themselves (i.e. inventory, locations, etc.).
        """

    @abstractmethod
    def scenario_step(self):
        """
        Update the state of the world according to whatever rules this scenario
        implements.

        This gets called in the 'step' method (of base_env) after going through each
        component step and before generating observations, rewards, etc.

        This is where things like resource regeneration, income redistribution, etc.,
        can be implemented.
        """

    @abstractmethod
    def generate_observations(self):
        """
        Generate observations associated with this scenario.

        A scenario does not need to produce observations and can provide observations
        for only some agent types; however, for a given agent type, it should either
        always or never yield an observation. If it does yield an observation,
        that observation should always have the same structure/sizes!

        Returns:
            obs (dict): A dictionary of {agent.idx: agent_obs_dict}. In words,
                return a dictionary with an entry for each agent (which can including
                the planner) for which this scenario provides an observation. For each
                entry, the key specifies the index of the agent and the value contains
                its associated observation dictionary.
        """

    @abstractmethod
    def compute_reward(self):
        """
        Apply the reward function(s) associated with this scenario to get the rewards
        from this step.

        Returns:
            rew (dict): A dictionary of {agent.idx: agent_obs_dict}. In words,
                return a  dictionary with an entry for each agent in the environment
                (including the planner). For each entry, the key specifies the index of
                the agent and the value contains the scalar reward earned this timestep.
        """

    # Optional methods for customization
    # ----------------------------------

    def additional_reset_steps(self):
        """
        Extra scenario-specific steps that should be performed at the end of the reset
        cycle.

        For each reset cycle...
            First, reset_starting_layout() and reset_agent_states() will be called.

            Second, <component>.reset() will be called for each registered component.

            Lastly, this method will be called to allow for any final customization of
            the reset cycle.
        """

    def scenario_metrics(self):
        """
        Allows the scenario to generate metrics (collected along with component metrics
        in the 'metrics' property).

        To have the scenario add metrics, this function needs to return a dictionary of
        {metric_key: value} where 'value' is a scalar (no nesting or lists!)
        """
        return


scenario_registry = Registry(BaseEnvironment)
"""The registry for Scenario classes.

This creates a registry object for Scenario classes. This registry requires that all
added classes are subclasses of BaseEnvironment. To make a Scenario class available
through the registry, decorate the class definition with @scenario_registry.add.

Example:
    from ai_economist.foundation.base.base_env
    import BaseEnvironment, scenario_registry

    @scenario_registry.add
    class ExampleScenario(BaseEnvironment):
        name = "Example"
        pass

    assert scenario_registry.has("Example")

    ScenarioClass = scenario_registry.get("Example")
    scenario = ScenarioClass(...)
    assert isinstance(scenario, ExampleScenario)

Notes:
    The foundation package exposes the scenario registry as: foundation.scenarios

    A Scenario class that is defined and registered following the above example will
    only be visible in foundation.scenarios if defined/registered in a file that is
    imported in ../scenarios/__init__.py.
"""
