# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/agents/benchmark_agents/01_EOQ.ipynb.

# %% auto 0
__all__ = ['EOQAgent', 'EOQPolicy']

# %% ../../../nbs/agents/benchmark_agents/01_EOQ.ipynb 4
# General libraries:
import numpy as np

# Mushroom libraries
from mushroom_rl.core import Agent

# %% ../../../nbs/agents/benchmark_agents/01_EOQ.ipynb 6
class EOQAgent(Agent):

    """
    Agent implementing the Economic Order Quantity (EOQ) policy.

    Args:
        mdp_info (MDPInfo): Information about the Markov Decision Process (MDP).
        s (numpy.ndarray): The fixed ordering cost.
        h (numpy.ndarray): The holding cost per unit per period.
        postprocessors (list): List of postprocessors to be applied to the policy.

    Attributes:
        mdp_info (MDPInfo): Information about the Markov Decision Process (MDP).
        policy (EOQPolicy): The EOQ policy implemented by the agent.

    """


    def __init__(self,
                  mdp_info,
                  s, # fixed ordering cost
                  h, # holding cost per unit per period
                  postprocessors = None
        ):

        policy = EOQPolicy(
            d = None,
            s = s,
            h = h,
            postprocessors = postprocessors
        )

        super().__init__(mdp_info, policy)

    def fit(self, demand):

        """ 
        Fit the EOQ policy to the given demand.

        This method allows the EOQ agent to adapt its policy to historic demand data, assuming a fixed demand rate without uncertainty.

        Parameters:
            demand (numpy.ndarray): The demand for each period and product. The array should have the shape (num_products, num_periods).

        Returns:
            None

        """


        assert isinstance(demand, np.ndarray)
        assert demand.ndim == 2

        self.policy.d = np.mean(demand, axis=0)

        self.policy.set_q_star()

class EOQPolicy():

    """
    Policy implementing the Economic Order Quantity (EOQ) strategy.

    Notethat d, s, and h must all have the shape (num_products,)

    Args:
        d (numpy.ndarray): The (average) demand per period for each product.
        s (numpy.ndarray): The fixed ordering cost for each product.
        h (numpy.ndarray): The holding cost per unit per period for each product.
        postprocessors (list): List of postprocessors to be applied to the action.

    Attributes:
        d (numpy.ndarray): The (average) demand per period for each product.
        s (numpy.ndarray): The fixed ordering cost for each product.
        h (numpy.ndarray): The holding cost per unit per period for each product.
        num_products (int): The number of products.
        q_star (numpy.ndarray): The optimal order quantity per product.
        postprocessors (list): List of postprocessors to be applied to the action.

    """

    def __init__(self,
                 d, 
                 s, 
                 h, 
                 postprocessors = None
                 ):
        self.d = d
        self.s = s
        self.h = h
        self.num_products = len(s)
        self.q_star = None
        if postprocessors is None:
            self.postprocessors = []
        else:
            self.postprocessors = (postprocessors)

    def set_q_star(self):
        
        """
        Set the optimal order quantity (q_star) for each product.

        This method calculates and assigns the optimal order quantity based on the EOQ formula.

        Returns:
            None

        """

        self.q_star = np.sqrt((2*self.s*self.d)/self.h)

    def draw_action(self, inventory_level):

        """
        Generate an action based on the current state.

        Returns zero for products which have still sufficient inventory, and the optimal order quantity for products which are running out of stock.

        Parameters:
            inventory_level (numpy.ndarray): The current inventory level for each product. Shape must be (num_products,).

        Returns:
            numpy.ndarray: The action to be taken, indicating the quantity to order for each product.

        """

        assert self.q_star is not None, "q_star is not set"

        action = np.zeros(self.num_products)

        # TODO account for lead time
        action = np.where(inventory_level >= self.d, 0, self.q_star)

        for postprocessor in self.postprocessors:
            action = postprocessor(action)

        return action
    
    def reset(self):
        pass
