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

# %% auto 0
__all__ = ['SSAgent', 'SSPolicy']

# %% ../../../nbs/agents/benchmark_agents/12_SS.ipynb 4
# General libraries:
import numpy as np
from scipy.stats import norm
from tqdm import tqdm

# Mushroom libraries
from mushroom_rl.core import Agent

# %% ../../../nbs/agents/benchmark_agents/12_SS.ipynb 6
class SSAgent(Agent):

    """
    Agent implementing the QR policy.

    # TODO adjust description

    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.
        l (numpy.ndarray): The lead time per product.
        preprocessors (list): List of preprocessors to be applied to the state.
        postprocessors (list): List of postprocessors to be applied to the policy.
        agent_name (str): Name of the agent. If set to None will use some default name.
        precision (int): Number of decimal places to round the demand input to.

    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,
                  mdp,
                  s, # fixed ordering cost
                  h, # holding cost per unit per period
                  l, # lead time
                  p, # penalty cost per unit
                  unit_size = 0.01,
                  preprocessors = None,
                  postprocessors = None,
                  agent_name = None,
                  precision = 5,
                  num_iterations_per_parameter = 12,
                  include_l = True
        ):

        mdp._mdp_info.horizon = mdp.demand.shape[0]
        mdp.reset(0)

        policy = SSPolicy(
            d = None,
            s = s,
            h = h,
            l = l,
            p = p,
            mdp = mdp,
            unit_size = unit_size,
            preprocessors = preprocessors,
            postprocessors = postprocessors,
            num_iterations_per_parameter = num_iterations_per_parameter,
            include_l = include_l,
        )

        self.precision=precision

        if agent_name is None:
            self.name = 'SSAgent'
        else:
            self.name = agent_name
        
        self.train_directly=True
        self.train_mode = "direct"

        super().__init__(mdp_info, policy)


    def fit(self, features = None, demand=None):

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

        # TODO adjust description

        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.set_s_S(demand, self._preprocessors[0])

class SSPolicy():

    """
    Policy implementing the QR strategy.

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

    # TODO adjust description

    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.
        l (numpy.ndarray): The lead time per 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.
        l (numpy.ndarray): The lead time per 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, 
                 l, 
                 p,
                 mdp,
                 unit_size = 0.01,
                 preprocessors = None,
                 postprocessors = None,
                 num_iterations_per_parameter = 12,
                 include_l = True
                 ):
        self.d = d
        self.s = s
        self.h = h
        self.l = l
        self.p = p
        self.unit_size = unit_size
        self.num_products = len(s)

        if preprocessors is None:
            self.preprocessors = []
        else:
            self.preprocessors = (preprocessors)
        if postprocessors is None:
            self.postprocessors = []
        else:
            self.postprocessors = (postprocessors)

        self.mdp = mdp

        self.num_iterations_per_parameter = num_iterations_per_parameter

        self.include_l = include_l

        self.s = 0.5 # initial value
        self.S = 0.5 # initial value
    
    def calculate_grid(self, s, S, step_size, plot=False):

        s_candidates = np.arange(s-step_size*self.num_iterations_per_parameter, s+step_size*self.num_iterations_per_parameter, step_size)
        S_candidates = np.arange(S-step_size*self.num_iterations_per_parameter, S+step_size*self.num_iterations_per_parameter, step_size)

        # only use positive or zero candidates
        s_candidates = s_candidates[s_candidates >= -0.1]
        S_candidates = S_candidates[S_candidates >= 0]

        cost_matrix = np.zeros((len(s_candidates), len(S_candidates)))

        for i, s_candidate in enumerate(s_candidates):
            for j, S_candidate in enumerate(S_candidates):
                cost_value = self.run_simulation(s_candidate, S_candidate)

                cost_matrix[i, j] = cost_value
                print("s:", np.round(s_candidate,2), "S:", np.round(S_candidate,2), "cost:", np.round(cost_value))

        best_cost = np.min(cost_matrix)
        best_s_index, best_S_index = np.unravel_index(np.argmin(cost_matrix), cost_matrix.shape)
        best_s = s_candidates[best_s_index]
        best_S = S_candidates[best_S_index]

        return best_s, best_S, best_cost
    
    def run_simulation(self, s, S):

        total_cost = 0

        state = self.mdp.reset(0)

        for t in range(self.mdp.info.horizon):

            state = self.preprocessor(state)
            action = self.draw_action_train(state, s, S)
                
            state, reward, _, _ = self.mdp.step(action)
            
            total_cost += -reward

        return total_cost

    def set_s_S(self, demand, preprocessor):
        
        """
        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.preprocessor = preprocessor

        self.s, self.S, cost = self.calculate_grid(self.s, self.S, step_size = self.unit_size*10)
        print("iteration 1:", "s:", self.s, "S:", self.S, "cost:", cost)
        self.s, self.S, cost = self.calculate_grid(self.s, self.S, step_size = self.unit_size*3)
        print("iteration 2:", "s:", self.s, "S:", self.S, "cost:", cost)
        self.s, self.S, cost = self.calculate_grid(self.s, self.S, step_size = self.unit_size)
        print("iteration 4:", "s:", self.s, "S:", self.S, "cost:", cost)


    def draw_action(self, input):

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

        # TODO adjust description

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

        Parameters:
            input (numpy.ndarray): The current inventory level and potentially order pipeline for each product.

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

        """

        for preprocessor in self.preprocessors:
            input = preprocessor(input)

        if self.include_l:
            total_inventory_position = np.sum(input)
            if total_inventory_position <= self.s:
                q = self.S - total_inventory_position
            else:
                q = 0 
        else:
            print("Only policy based on total inventory position implemented")
        

        action = np.array([q])

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

        return action
    
    def draw_action_train(self, input, s, S):

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

        # TODO adjust description

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

        Parameters:
            input (numpy.ndarray): The current inventory level and potentially order pipeline for each product.

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

        """

        #print("R:", R, "Q:", Q)

        for preprocessor in self.preprocessors:
            input = preprocessor(input)

        if self.include_l:
            total_inventory_position = np.sum(input)
            if total_inventory_position <= s:
                q = S - total_inventory_position
            else:
                q = 0 
                
        action = np.array([q])

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

        return action
    
    def reset(self):
        pass
