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

# %% auto 0
__all__ = ['QRAgent', 'QRPolicy']

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

# Mushroom libraries
from mushroom_rl.core import Agent

# %% ../../../nbs/agents/benchmark_agents/12_QR.ipynb 6
class QRAgent(Agent):

    train_directly=True

    """
    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,
                  manually = False,
                  num_iterations = 50,
                  include_l = True
        ):

        policy = QRPolicy(
            d = None,
            s = s,
            h = h,
            l = l,
            p = p,
            mdp = mdp,
            unit_size = unit_size,
            preprocessors = preprocessors,
            postprocessors = postprocessors,
            manually = manually,
            num_iterations=num_iterations,
            include_l = include_l
        )

        self.precision=precision

        if agent_name is None:
            self.name = 'QRAgent'
        else:
            self.name = agent_name

        super().__init__(mdp_info, policy)

    def fit(self, demand):

        """ 
        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_q_r(demand)

class QRPolicy():

    """
    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,
                 manually = False,
                 num_iterations = 50,
                 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)
        self.q_star = None
        if preprocessors is None:
            self.preprocessors = []
        else:
            self.preprocessors = (preprocessors)
        if postprocessors is None:
            self.postprocessors = []
        else:
            self.postprocessors = (postprocessors)

        self.mdp = mdp
        self.manually = manually

        self.num_iterations = num_iterations

        self.include_l = include_l

    def calculate_initial_R(self, demand, lam, Q):

        try:
            percentile = 1 - self.h * (Q / (self.p * lam))
            R = np.percentile(demand, percentile * 100)
        except:
            R = self.l * lam
        return R
    
    def expected_stockouts(self, demand, R):
        total_stockouts = 0
        for i in range(demand.shape[0]-(self.l.max())):
            demand_during_lead_time = demand[i:i+self.l.max(),] #! Does not work for different lead times
            stockouts = np.maximum(np.sum(demand_during_lead_time, axis=0) - R, 0)
            total_stockouts += stockouts
        
        exp_stockouts = total_stockouts/(demand.shape[0]-self.l) 

        return exp_stockouts
    
    def calculate_incremental_R(self, R):

        best_cost = self.run_simulation(self.q, R)
        best_R = R

        R_candidate = R + self.unit_size
        cost = self.run_simulation(self.q, R_candidate)

        if cost <= best_cost:
            best_cost = cost
            best_R = R_candidate
            improve = True

            while improve:
                R_candidate += self.unit_size
                cost = self.run_simulation(self.q, R_candidate)

                if cost > best_cost:
                    improve = False
                else:
                    best_cost = cost
                    best_R = R_candidate
        
        elif cost > best_cost:
            R_candidate = R - self.unit_size
            cost = self.run_simulation(self.q, R_candidate)

            if cost <= best_cost:
                best_cost = cost
                best_R = R_candidate
                improve = True

                while improve:
                    R_candidate -= self.unit_size
                    cost = self.run_simulation(self.q, R_candidate)

                    if cost > best_cost:
                        improve = False
                    else:
                        best_cost = cost
                        best_R = R_candidate
        
        return best_R, best_cost

    def calculate_incremental_Q(self, Q):

        best_cost = self.run_simulation(Q, self.r)
        best_Q = Q

        Q_candidate = Q + self.unit_size
        cost = self.run_simulation(Q_candidate, self.r)

        if cost <= best_cost:
            best_cost = cost
            best_Q = Q_candidate
            improve = True

            while improve:
                Q_candidate += self.unit_size
                cost = self.run_simulation(Q_candidate, self.r)

                if cost > best_cost:
                    improve = False
                else:
                    best_cost = cost
                    best_Q = Q_candidate

        elif cost > best_cost:
            Q_candidate = Q - self.unit_size
            cost = self.run_simulation(Q_candidate, self.r)

            if cost <= best_cost:
                best_cost = cost
                best_Q = Q_candidate
                improve = True

                while improve:
                    Q_candidate -= self.unit_size
                    cost = self.run_simulation(Q_candidate, self.r)

                    if cost > best_cost:
                        improve = False
                    else:
                        best_cost = cost
                        best_Q = Q_candidate
            
        return best_Q, best_cost
    
    
    def run_simulation(self, Q, R):

        total_cost = 0

        for _ in range(50):

            state = self.mdp.reset()
        
            cost = 0

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

                action = self.draw_action_train(state, Q, R)
                state, reward, _, _ = self.mdp.step(action)
                cost += -reward
            
            total_cost += cost
        
        cost = total_cost/50

        return cost

    def set_q_r(self, demand):
        
        """
        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

        """

        if self.manually:

            self.r = 0.01
            self.q = 0.5

        else:
            lam = np.mean(demand, axis=0)

            # Note: s in Cornell presentation denoted by K
            self.q = (np.sqrt(2*self.s*lam/self.h))[0]
            self.r = (self.calculate_initial_R(demand, lam, self.q))[0]

            print("initial r:", self.r, "initial q:", self.q)

            for i in range(self.num_iterations):
                
                self.r, _ = self.calculate_incremental_R(self.r)
                self.q, cost = self.calculate_incremental_Q(self.q)

                print("iteration:", i, "q:", self.q, "r:", self.r, "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.

        """

        assert self.q is not None, "q is not set"
        assert self.r is not None, "r is not set"

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

        pipeline_vector = input[self.num_products:]
        pipeline = np.reshape(pipeline_vector, (self.num_products, max(self.l)))
        pipeline_sum = np.sum(pipeline, axis=1)

        if self.include_l:
            r = (self.r - pipeline_sum)[0]
        else:
            r = -1 if pipeline_sum != 0 else self.r
        q = self.q

        action = np.array([r, q])

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

        return action
    
    def draw_action_train(self, input, Q, R):

        """
        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)

        pipeline_vector = input[self.num_products:]

        pipeline = np.reshape(pipeline_vector, (self.num_products, max(self.l)))
        pipeline_sum = np.sum(pipeline, axis=1)
        input = input[:self.num_products]

        if self.include_l:
            r = (R - pipeline_sum)[0]
        else:
            r = -1 if pipeline_sum != 0 else R
        q = Q

        action = np.array([r, q]) 

        #print("action:", action)

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

        return action
    
    def reset(self):
        pass
