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

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

# %% ../../../nbs/agents/benchmark_agents/12_QR.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_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
        
        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_q_r(demand, self._preprocessors[0])

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

        self.q = 0.5
        self.r = 0.5

    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 tqdm(range(50)):

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

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

                state = self.preprocessor(state)
                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, 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


        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

# %% ../../../nbs/agents/benchmark_agents/12_QR.ipynb 7
class QRAgent_v2(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 = 10,
                  include_l = True
        ):

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

        policy = QRPolicy_v2(
            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 = 'QRAgent_v2'
        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_q_r(demand, self._preprocessors[0])

class QRPolicy_v2():

    """
    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 = 10,
                 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.q = 0.5 # initial value
        self.r = 0.5 # initial value
    
    def calculate_grid(self, R, Q, step_size, plot=False):

        R_candidates = np.arange(R-step_size*self.num_iterations_per_parameter, R+step_size*self.num_iterations_per_parameter, step_size)
        Q_candidates = np.arange(Q-step_size*self.num_iterations_per_parameter, Q+step_size*self.num_iterations_per_parameter, step_size)

        # only use positive or zero candidates
        R_candidates = R_candidates[R_candidates >= 0]
        Q_candidates = Q_candidates[Q_candidates >= 0]

        cost_matrix = np.zeros((len(R_candidates), len(Q_candidates)))

        for i, R_candidate in enumerate(R_candidates):
            for j, Q_candidate in enumerate(Q_candidates):
                cost_value = self.run_simulation(Q_candidate, R_candidate)

                cost_matrix[i, j] = cost_value
                print("R:", np.round(R_candidate,2), "Q:", np.round(Q_candidate,2), "cost:", np.round(cost_value))

        best_cost = np.min(cost_matrix)
        best_R_index, best_Q_index = np.unravel_index(np.argmin(cost_matrix), cost_matrix.shape)
        best_R = R_candidates[best_R_index]
        best_Q = Q_candidates[best_Q_index]

        if plot:
            Q_grid, R_grid = np.meshgrid(Q_candidates, R_candidates)

            import matplotlib.pyplot as plt

            fig = plt.figure()
            ax = fig.add_subplot(111, projection='3d')
            ax.plot_surface(R_grid, Q_grid, cost_matrix, cmap='viridis')
            ax.set_xlabel('R')
            ax.set_ylabel('Q')
            ax.set_zlabel('Cost')
            plt.show()

        return best_R, best_Q, best_cost
    
    def run_simulation(self, Q, R):

        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, Q, R)
                
            state, reward, _, _ = self.mdp.step(action)
            
            total_cost += -reward

        return total_cost

    def set_q_r(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.r, self.q, cost = self.calculate_grid(self.r, self.q, step_size = self.unit_size*10)
        print("iteration 1:", "q:", self.q, "r:", self.r, "cost:", cost)
        self.r, self.q, cost = self.calculate_grid(self.r, self.q, step_size = self.unit_size*3)
        print("iteration 2:", "q:", self.q, "r:", self.r, "cost:", cost)
        self.r, self.q, cost = self.calculate_grid(self.r, self.q, step_size = self.unit_size)
        print("iteration 4:", "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)

        if self.include_l:
            total_inventory_position = np.sum(input)
            if total_inventory_position <= self.r:
                r = self.r 
            else:
                r = -1
        
        else:
            pipeline_vector = input[self.num_products:]
            # pipeline = np.reshape(pipeline_vector, (self.num_products, max(self.l)))
            pipeline_sum = np.sum(pipeline_vector) # currently only works for single product case
            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)

        if self.include_l:
            total_inventory_position = np.sum(input)
            if total_inventory_position <= R:
                r = R
            else:
                r = -1
        
        else:
            pipeline_vector = input[self.num_products:]
            # pipeline = np.reshape(pipeline_vector, (self.num_products, max(self.l)))
            pipeline_sum = np.sum(pipeline_vector) # currently only works for single product case
            r = -1 if pipeline_sum != 0 else R

        q = Q

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

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

        return action
    
    def reset(self):
        pass
