# 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 numpy as np

from ai_economist.foundation.base.base_component import (
    BaseComponent,
    component_registry,
)


@component_registry.add
class ExternalMarket(BaseComponent):
    """
    Allows mobile agents to build house landmarks in the world using stone and wood,
    earning income.

    Can be configured to include heterogeneous building skill where agents earn
    different levels of income when building.

    Args:
        payment (int): Default amount of coin agents earn from building.
            Must be >= 0. Default is 10.
        market_demand (dict): Resource name -> amout of money
        skill_dist (str): Distribution type for sampling skills. Default ("none")
            gives all agents identical skill equal to a multiplier of 1. "pareto" and
            "lognormal" sample skills from the associated distributions.
        build_labor (float): Labor cost associated with building a house.
            Must be >= 0. Default is 10.
    """

    name = "ExternalMarket"
    component_type = "Trade"
    required_entities = ["Coin", "Labor"]
    agent_subclasses = ["TradingAgent"]

    def __init__(
        self,
        *base_component_args,
        market_demand={},
        trade_labor=1.0,
        **base_component_kwargs
    ):
        super().__init__(*base_component_args, **base_component_kwargs)

        self.market_demand = market_demand
        self.action_res_map={}
        for k in market_demand.keys():
            self.action_res_map[len(self.action_res_map)+1]=k

        self.trade_labor = float(trade_labor)
        assert self.trade_labor >= 0

        self.builds = []

    def agent_can_sell(self, agent,res):
        """Return True if agent can sell a res."""
        # See if the agent has the resources necessary to complete the action
       
        if agent.state["inventory"][res]>= 1:
            return True
        return False

    # Required methods for implementing components
    # --------------------------------------------

    def get_n_actions(self, agent_cls_name):
        """
        See base_component.py for detailed description.

        Add a single action (build) for mobile agents.
        """
        # This component adds 1 action that mobile agents can take: build a house
        if agent_cls_name in self.agent_subclasses:
            return len(self.action_res_map)

        return None

    def get_additional_state_fields(self, agent_cls_name):
        """
        See base_component.py for detailed description.

        For mobile agents, add state fields for building skill.
        """
      
        return {}

    def component_step(self):
        """
        See base_component.py for detailed description.

        Convert stone+wood to house+coin for agents that choose to build and can.
        """
        world = self.world
        build = []
        # Apply any building actions taken by the mobile agents
        for agent in world.get_random_order_agents():

            action = agent.get_component_action(self.name)

            # This component doesn't apply to this agent!
            if action is None:
                continue

            # NO-OP!
            if action == 0:
                continue
           
            res_name=self.action_res_map[action]
            # Build! (If you can.)
       
            if self.agent_can_sell(agent,res_name):
                    # Remove the resources
                    agent.state["inventory"][res_name] -= 1

                    # Receive payment for the house
                    agent.state["inventory"]["Coin"] += self.market_demand[res_name]

                    # Incur the labor cost for building
                    agent.state["endogenous"]["Labor"] += self.trade_labor

                    build.append(
                        {
                            "seller": agent.idx,
                            "commodity": res_name,
                            "income": self.market_demand[res_name],
                        }
                    )
            else:
                raise ValueError

        self.builds.append(build)

    def generate_observations(self):
        """
        See base_component.py for detailed description.

        Here, agents observe their build skill. The planner does not observe anything
        from this component.
        """

        obs_dict = dict()
        for agent in self.world.agents:
            if agent.name in self.agent_subclasses:

                obs_dict[agent.idx] = {}
                for res_name,coin in self.market_demand.items():            
                    obs_dict[agent.idx]["external_{}_price".format(res_name)]: self.inv_scale*coin            

        return obs_dict

    def generate_masks(self, completions=0):
        """
        See base_component.py for detailed description.

        Prevent building only if a landmark already occupies the agent's location.
        """

        masks = {}
        # Mobile agents' build action is masked if they cannot build with their
        # current location and/or endowment
        for agent in self.world.agents:
            if agent.name in self.agent_subclasses:
                mask=[]
                for res in self.market_demand:
                    mask.append(self.agent_can_sell(agent,res))
                masks[agent.idx] = mask

        return masks

    # For non-required customization
    # ------------------------------

    def get_metrics(self):
        """
        Metrics that capture what happened through this component.

        Returns:
            metrics (dict): A dictionary of {"metric_name": metric_value},
                where metric_value is a scalar.
        """
        world = self.world
        """
        build_stats = {a.idx: {"n_builds": 0} for a in world.agents}
        for builds in self.builds:
            for build in builds:
                idx = build["builder"]
                build_stats[idx]["n_builds"] += 1

        out_dict = {}
        for a in world.agents:
            for k, v in build_stats[a.idx].items():
                out_dict["{}/{}".format(a.idx, k)] = v

        num_houses = np.sum(world.maps.get("House") > 0)
        out_dict["total_builds"] = num_houses
        """ 
        return {}

    def additional_reset_steps(self):
        """
        See base_component.py for detailed description.

        Re-sample agents' building skills.
        """
        self.builds = []

    def get_dense_log(self):
        """
        Log builds.

        Returns:
            builds (list): A list of build events. Each entry corresponds to a single
                timestep and contains a description of any builds that occurred on
                that timestep.

        """
        return self.builds