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

# %% auto 0
__all__ = ['DualSourcingBaseAgent', 'SingleIndexAgent', 'DualIndexAgent', 'CappedDualIndexAgent', 'SingleIndexPolicy',
           'DualIndexPolicy', 'CappedDualIndexPolicy']

# %% ../../../nbs/agents/benchmark_agents/12_Dual_Sourcing_Agents.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

import time

# %% ../../../nbs/agents/benchmark_agents/12_Dual_Sourcing_Agents.ipynb 6
class DualSourcingBaseAgent(Agent):

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

        """ 
        xxx
        """

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

        self.policy.set_params(demand, self._preprocessors[0])


class SingleIndexAgent(DualSourcingBaseAgent):

    """
    xxx

    """

    def __init__(self,
                  mdp_info,
                  mdp,
                  l, # lead time
                  l_e, # expedited lead time
                  unit_size = 0.01,
                  preprocessors = None,
                  postprocessors = None,
                  agent_name = None,
                  precision = 5,
                  num_iterations_per_parameter = 12,
                  algorithm = "brute_0_1", 
                  full_grid=False,# "brute_0_1": brute force for demand between 0 and 1 and unit-size <1, "brute_variable": brute force for demand large 1 and unit size 1 
                  #search_space = 3 # applies only for brute_variable - grid size in multiple of average demand
        ):

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

        policy = SingleIndexPolicy(
            l = l,
            l_e = l_e,
            mdp = mdp,
            unit_size = unit_size,
            preprocessors = preprocessors,
            postprocessors = postprocessors,
            num_iterations_per_parameter = num_iterations_per_parameter,
            algorithm = algorithm,
            full_grid = full_grid,
            #search_space = search_space
        )

        self.precision=precision

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

        super().__init__(mdp_info, policy)

class DualIndexAgent(DualSourcingBaseAgent):

    """
    xxx

    """

    def __init__(self,
                  mdp_info,
                  mdp,
                  l, # lead time
                  l_e, # expedited lead time
                  unit_size = 0.01,
                  preprocessors = None,
                  postprocessors = None,
                  agent_name = None,
                  precision = 5,
                  num_iterations_per_parameter = 12,
                  algorithm = "brute_0_1",
                  full_grid=False, # "brute_0_1": brute force for demand between 0 and 1 and unit-size <1, "brute_variable": brute force for demand large 1 and unit size 1 
                  return_levels=False,
                  #search_space = 3 # applies only for brute_variable - grid size in multiple of average demand
        ):

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

        policy = DualIndexPolicy(
            l = l,
            l_e = l_e,
            mdp = mdp,
            unit_size = unit_size,
            preprocessors = preprocessors,
            postprocessors = postprocessors,
            num_iterations_per_parameter = num_iterations_per_parameter,
            algorithm = algorithm,
            full_grid = full_grid,
            return_levels = return_levels,
            #search_space = search_space
        )

        self.precision=precision

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

        super().__init__(mdp_info, policy)

class CappedDualIndexAgent(DualSourcingBaseAgent):

    """
    xxx

    """

    def __init__(self,
                  mdp_info,
                  mdp,
                  l, # lead time
                  l_e, # expedited lead time
                  unit_size = 0.01,
                  preprocessors = None,
                  postprocessors = None,
                  agent_name = None,
                  precision = 5,
                  num_iterations_per_parameter = 12,
                  algorithm = "brute_0_1",
                  full_grid=False, # "brute_0_1": brute force for demand between 0 and 1 and unit-size <1, "brute_variable": brute force for demand large 1 and unit size 1 
                  return_levels=False,
                  #search_space = 3 # applies only for brute_variable - grid size in multiple of average demand
        ):

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

        policy = CappedDualIndexPolicy(
            l = l,
            l_e = l_e,
            mdp = mdp,
            unit_size = unit_size,
            preprocessors = preprocessors,
            postprocessors = postprocessors,
            num_iterations_per_parameter = num_iterations_per_parameter,
            algorithm = algorithm,
            full_grid = full_grid,
            #search_space = search_space
        )

        self.precision=precision

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

        super().__init__(mdp_info, policy)


class SingleIndexPolicy():

    """
    xxx
    """

    def __init__(self,
                 l, 
                 l_e, 
                 mdp,
                 unit_size = 0.01,
                 preprocessors = None,
                 postprocessors = None,
                 num_iterations_per_parameter = 12,
                 algorithm = "brute_0_1", 
                 full_grid = False,# "brute_0_1": brute force for demand between 0 and 1 and unit-size <1, "brute_0_1": brute force for demand large 1 and unit size 1
                 ):
        self.l = l
        self.l_e = l_e
        self.unit_size = unit_size
        self.num_products = len(l)

        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.S = 0.5 #0.97 # initial value
        self.S_e = 0.5 #0.86 # initial value

        self.algorithm = algorithm

        self.full_grid = full_grid
    
    def set_and_calculate_grid(self, S, S_e, step_size, full_grid=False, plot=False,):
        


        if full_grid:
            S_candidates = np.arange(0, 1+step_size, step_size)
            S_e_candidates = np.arange(0, 1+step_size, step_size)
            
        else:

            S_candidates = np.arange(S-step_size*self.num_iterations_per_parameter, S+step_size*(self.num_iterations_per_parameter+1), step_size)
            S_e_candidates = np.arange(S_e-step_size*self.num_iterations_per_parameter, S_e+step_size*(self.num_iterations_per_parameter+1), step_size)

            # only use positive or zero candidates
            S_candidates = S_candidates[S_candidates >= 0]
            S_e_candidates = S_e_candidates[S_e_candidates >= 0]

        print("base S:", S)
        print("base S_e:", S_e)
        print("S candidates:", S_candidates)
        print("S_e candidates:", S_e_candidates)

        best_S, best_S_e, best_cost, _ = self.calculate_grid(S_candidates, S_e_candidates)

        return best_S, best_S_e, best_cost

    def calculate_grid(self, S_candidates, S_e_candidates):
    
        cost_matrix = np.zeros((len(S_candidates), len(S_e_candidates)))

        for i, S_candidate in enumerate(S_candidates):
            for j, S_e_candidate in enumerate(S_e_candidates):
                if S_e_candidate > S_candidate:
                    cost_value = np.inf
                else:
                    cost_value = self.run_simulation(S_candidate, S_e_candidate)

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

        best_cost = np.min(cost_matrix)
        best_S_index, best_S_e_index = np.unravel_index(np.argmin(cost_matrix), cost_matrix.shape)
        best_S = S_candidates[best_S_index]
        best_S_e = S_e_candidates[best_S_e_index]

        print("max S index:", len(S_candidates), "best S index:", best_S_index)
        print("max S_e index:", len(S_e_candidates), "best S_e index:", best_S_e_index)

        print("best S and S_e:", best_S, best_S_e)

        if best_S_index == 0 or best_S_index == len(S_candidates)-1:
            warn_edge = True
        elif best_S_e_index == 0 or best_S_e_index == len(S_e_candidates)-1:
            warn_edge = True
        else:
            warn_edge = False
        
        if warn_edge:
            print("Warning: Optimal value at the edge of search space")

        return best_S, best_S_e, best_cost, warn_edge

    def run_simulation(self, S, S_e):

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

        return total_cost

    def set_params(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.algorithm == "brute_0_1":

            if self.full_grid:
                self.S, self.S_e, cost = self.set_and_calculate_grid(self.S, self.S_e, step_size = self.unit_size, full_grid=True)
                print("result:", "S:", self.S, "S_e:", self.S_e, "cost:", cost)
            else:
                self.S, self.S_e, cost = self.set_and_calculate_grid(self.S, self.S_e, step_size = self.unit_size*10)
                print("iteration 1:", "S:", self.S, "S_e:", self.S_e, "cost:", cost)
                self.S, self.S_e, cost = self.set_and_calculate_grid(self.S, self.S_e, step_size = self.unit_size*3)
                print("iteration 2:", "S:", self.S, "S_e:", self.S_e, "cost:", cost)
                self.S, self.S_e, cost = self.set_and_calculate_grid(self.S, self.S_e, step_size = self.unit_size)
                print("iteration 3:", "S:", self.S, "S_e:", self.S_e, "cost:", cost)
        
        elif self.algorithm == "brute_variable":
            # rais not implemented

            raise NotImplementedError("only implemented for values normalized between 0 and 1")


            # self.run_brute_variable(demand)
        
        else:
            # error
            raise ValueError("Algorithm unknown")
    
    

    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)
        
        # print("input")
        total_inventory_position = np.sum(input)
        # print(total_inventory_position)
        if total_inventory_position < self.S_e:
            q_e = self.S_e - total_inventory_position
            q = self.S-self.S_e
        elif total_inventory_position < self.S:
            q_e=0
            q = self.S - total_inventory_position
        else:
            q_e=0
            q = 0

        # print("q", q, "q_e", q_e)
        action = np.array([q, q_e])

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

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

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

        total_inventory_position = np.sum(input)
        if total_inventory_position < S_e:
            q_e = S_e - total_inventory_position
            q = S-S_e
        elif total_inventory_position < S:
            q_e=0
            q = S - total_inventory_position 
        else:
            q_e=0
            q = 0 

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

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

        return action
    
    def reset(self):
        pass

class DualIndexPolicy(SingleIndexPolicy):

    """
    xxx
    """

    def __init__(self,
                 l, 
                 l_e, 
                 mdp,
                 unit_size = 0.01,
                 preprocessors = None,
                 postprocessors = None,
                 num_iterations_per_parameter = 12,
                 algorithm = "brute_0_1", # "brute_0_1": brute force for demand between 0 and 1 and unit-size <1, "brute_0_1": brute force for demand large 1 and unit size 1
                 full_grid = False,
                 return_levels = False
                 ):
        
        self.return_levels = return_levels
        
        super().__init__(l, l_e, mdp, unit_size, preprocessors, postprocessors, num_iterations_per_parameter, algorithm, full_grid)
    
    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)

        total_inventory_position = np.sum(input)
        expedited_inventory_position = input[0]+np.sum(input[-(self.l_e)]) # first position is inventory, rest is pipeline vector

        if expedited_inventory_position < self.S_e:
            q_e = self.S_e - expedited_inventory_position
            q = np.maximum(self.S-total_inventory_position-q_e,0)
        elif total_inventory_position < self.S:
            q_e=0
            q = self.S - total_inventory_position 
        else:
            q_e=0
            q = 0

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

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

        if self.return_levels:
            action = np.array([self.S, self.S_e])

        return action

    def draw_action_train(self, input, S, S_e):

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

        total_inventory_position = np.sum(input)
        expedited_inventory_position = input[0]+np.sum(input[-(self.l_e)]) # first position is inventory, rest is pipeline vector

        if expedited_inventory_position < S_e:
            q_e = S_e - expedited_inventory_position
            q = np.maximum(S-total_inventory_position-q_e,0)
        elif total_inventory_position < S:
            q_e=0
            q = S - total_inventory_position 
        else:
            q_e=0
            q = 0

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

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

        if self.return_levels:
            action = np.array([S, S_e])

        return action

class CappedDualIndexPolicy(SingleIndexPolicy):

    """
    xxx
    """

    def __init__(self,
                 l, 
                 l_e, 
                 mdp,
                 unit_size = 0.01,
                 preprocessors = None,
                 postprocessors = None,
                 num_iterations_per_parameter = 12,
                 algorithm = "brute_0_1", # "brute_0_1": brute force for demand between 0 and 1 and unit-size <1, "brute_0_1": brute force for demand large 1 and unit size 1
                 full_grid = False,
                 ):
        
        self.cap = 0.5
        
        super().__init__(l, l_e, mdp, unit_size, preprocessors, postprocessors, num_iterations_per_parameter, algorithm, full_grid)
    
    def set_and_calculate_grid(self, S, S_e, cap, step_size, full_grid=False, plot=False,):

        if full_grid:
            S_candidates = np.arange(0, 1+step_size, step_size)
            S_e_candidates = np.arange(0, 1+step_size, step_size)
            cap_candidates = np.arange(0, 1+step_size, step_size)
            
        else:
            S_candidates = np.arange(S-step_size*self.num_iterations_per_parameter, S+step_size*(self.num_iterations_per_parameter+1), step_size)
            S_e_candidates = np.arange(S_e-step_size*self.num_iterations_per_parameter, S_e+step_size*(self.num_iterations_per_parameter+1), step_size)
            cap_candidates = np.arange(cap-step_size*self.num_iterations_per_parameter, cap+step_size*(self.num_iterations_per_parameter+1), step_size)

            # only use positive or zero candidates
            S_candidates = S_candidates[S_candidates >= 0]
            S_e_candidates = S_e_candidates[S_e_candidates >= 0]
            cap_candidates = cap_candidates[cap_candidates >= 0]

        print("base S:", S)
        print("base S_e:", S_e)
        print("base cap:", cap)
        print("S candidates:", S_candidates)
        print("S_e candidates:", S_e_candidates)
        print("cap candidates:", cap_candidates)

        best_S, best_S_e, best_cap, best_cost, _ = self.calculate_grid(S_candidates, S_e_candidates, cap_candidates)

        return best_S, best_S_e, best_cap, best_cost
    
    def calculate_grid(self, S_candidates, S_e_candidates, cap_candidates):
    
        cost_matrix = np.zeros((len(S_candidates), len(S_e_candidates), len(cap_candidates)))

        for i, S_candidate in enumerate(S_candidates):
            for j, S_e_candidate in enumerate(S_e_candidates):
                if S_e_candidate > S_candidate:
                    cost_value = np.inf
                else:
                    for k, cap_candidate in enumerate(cap_candidates):
                        cost_value = self.run_simulation(S_candidate, S_e_candidate, cap_candidate)
                        cost_matrix[i, j, k] = cost_value
                        print("S:", np.round(S_candidate,2), "S_e:", np.round(S_e_candidate,2), "cap:", np.round(cap_candidate,2), "cost:", np.round(cost_value))
                cost_matrix[i, j, k] = cost_value
                print("S:", np.round(S_candidate,2), "S_e:", np.round(S_e_candidate,2), "cap:", np.nan, "cost:", np.round(cost_value))

        best_cost = np.min(cost_matrix)
        best_S_index, best_S_e_index, best_cap_index = np.unravel_index(np.argmin(cost_matrix), cost_matrix.shape)
        best_S = S_candidates[best_S_index]
        best_S_e = S_e_candidates[best_S_e_index]
        best_cap = cap_candidates[best_cap_index]

        print("max S index:", len(S_candidates), "best S index:", best_S_index)
        print("max S_e index:", len(S_e_candidates), "best S_e index:", best_S_e_index)
        print("max cap index:", len(cap_candidates), "best cap index:", best_cap_index)

        print("best S, S_e and cap:", best_S, best_S_e, best_cap)

        if best_S_index == 0 or best_S_index == len(S_candidates)-1:
            warn_edge = True
        elif best_S_e_index == 0 or best_S_e_index == len(S_e_candidates)-1:
            warn_edge = True
        elif best_cap_index == 0 or best_cap_index == len(cap_candidates)-1:
            warn_edge = True
        else:
            warn_edge = False
        
        if warn_edge:
            print("Warning: Optimal value at the edge of search space")

        return best_S, best_S_e, best_cap, best_cost, warn_edge

    def run_simulation(self, S, S_e, cap):

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

        return total_cost
    
    def set_params(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.algorithm == "brute_0_1":

            if self.full_grid:
                self.S, self.S_e, cost = self.set_and_calculate_grid(self.S, self.S_e, self.cap, step_size = self.unit_size, full_grid=True)
                print("result:", "S:", self.S, "S_e:", self.S_e, "cost:", cost)
            else:
                self.S, self.S_e, cost = self.set_and_calculate_grid(self.S, self.S_e, self.cap, step_size = self.unit_size*10)
                print("iteration 1:", "S:", self.S, "S_e:", self.S_e, "cap:", self.cap, "cost:", cost)
                self.S, self.S_e, cost = self.set_and_calculate_grid(self.S, self.S_e, self.cap, step_size = self.unit_size*3)
                print("iteration 2:", "S:", self.S, "S_e:", self.S_e, "cap:", self.cap, "cost:", cost)
                self.S, self.S_e, cost = self.set_and_calculate_grid(self.S, self.S_e, self.cap, step_size = self.unit_size)
                print("iteration 3:", "S:", self.S, "S_e:", self.S_e, "cap:", self.cap, "cost:", cost)
        
        elif self.algorithm == "brute_variable":
            # rais not implemented

            raise NotImplementedError("only implemented for values normalized between 0 and 1")


            # self.run_brute_variable(demand)
        
        else:
            # error
            raise ValueError("Algorithm unknown")
    
    

    
    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)

        total_inventory_position = np.sum(input)
        expedited_inventory_position = input[0]+np.sum(input[-(self.l_e)]) # first position is inventory, rest is pipeline vector

        if expedited_inventory_position < self.S_e:
            q_e = self.S_e - expedited_inventory_position
            q = np.maximum(np.maximum(self.S-total_inventory_position-q_e,0), self.cap)
        elif total_inventory_position < self.S:
            q_e=0
            q = np.maximum(self.S - total_inventory_position, self.cap)
        else:
            q_e=0
            q = 0

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

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

        return action

    def draw_action_train(self, input, S, S_e, cap):

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

        total_inventory_position = np.sum(input)
        expedited_inventory_position = input[0]+np.sum(input[-(self.l_e)]) # first position is inventory, rest is pipeline vector

        if expedited_inventory_position < S_e:
            q_e = S_e - expedited_inventory_position
            q = np.maximum(np.maximum(S-total_inventory_position-q_e,0), cap)
        elif total_inventory_position < S:
            q_e=0
            q = np.maximum(S - total_inventory_position, cap)
        else:
            q_e=0
            q = 0

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

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

        return action
