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

# %% auto 0
__all__ = ['CBSAgent', 'CBSPolicy']

# %% ../../../nbs/agents/benchmark_agents/21_Capped_base_stock.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/21_Capped_base_stock.ipynb 6
class CBSAgent(Agent):

    train_directly=True

    """_summary_
    """

    def __init__(self,
                    mdp_info,
                    mdp,
                    h, # holding cost per unit per period
                    p, # penalty cost per unit
                    l, # lead time
                    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 = CBSPolicy(
                        d=None, # Demand
                        h=h, # holding cost per unit per period
                        p=p, # penalty cost per unit
                        l=l, # lead time
                        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 = 'CBSAgent'
        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):
        assert isinstance(demand, np.ndarray)
        assert demand.ndim == 2

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

class CBSPolicy():
    """
    

    """

    def __init__(self,
                    d, 
                    h,
                    p,
                    l,
                    mdp,
                    unit_size,
                    preprocessors,
                    postprocessors,
                    num_iterations_per_parameter,
                    include_l
                    ):

        self.d = d
        self.h = h
        self.p = p
        self.l = l
        self.unit_size = unit_size

        # num products len(h) if h is not scalar, otherwise 1
        self.num_products = len(h) if isinstance(h, np.ndarray) else 1

        if preprocessors is None:
            self.preprocessors = list()
        else:
            self.preoprocessors = (preprocessors)
        if postprocessors is None:
            self.postprocessors = list()
        else:
            self.postprocessors = (postprocessors)

        self.mdp=mdp
        self.num_iterations_per_parameter=num_iterations_per_parameter
        self.include_l=include_l

        self.base = 0.5
        self.cap = 0.5
    

    def calculate_grid(self, base, cap, step_size, plot=False):

        base_candidates = np.arange(base-step_size*self.num_iterations_per_parameter, base+step_size*self.num_iterations_per_parameter, step_size)
        cap_candidates = np.arange(cap-step_size*self.num_iterations_per_parameter, cap+step_size*self.num_iterations_per_parameter, step_size)

        # only use positive or zero candidates
        base_candidates = base_candidates[base_candidates >= -0.1]
        cap_candidates = cap_candidates[cap_candidates >= -0.1]

        cost_matrix = np.zeros((len(base_candidates), len(cap_candidates)))

        for i, base_candidate in enumerate(base_candidates):
            for j, cap_candidate in enumerate(cap_candidates):
                cost_value = self.run_simulation(base_candidate, cap_candidate)

                cost_matrix[i, j] = cost_value
                print("base:", np.round(base_candidate,2), "cap:", np.round(cap_candidate,2), "cost:", np.round(cost_value))

        best_cost = np.min(cost_matrix)
        best_base_index, best_cap_index = np.unravel_index(np.argmin(cost_matrix), cost_matrix.shape)

        best_base = base_candidates[best_base_index]
        best_cap = cap_candidates[best_cap_index]

        return best_base, best_cap, best_cost

    def run_simulation(self, base, 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, base, cap)

            state, reward, _, _ = self.mdp.step(action)

            total_cost += -reward

        return total_cost
    
    def set_params(self, demand, preprocessor):
        
        self.preprocessor = preprocessor

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

    def draw_action(self, input):

        """
        Note: only designed for single product case
        """

        for preprocessor in self.preprocessors:
            # input must be a vector containing the inventory and the pipeline vector
            input = preprocessor(input)
        
        # print(input)
        # breakpoint()
        
        total_inventory_position = np.sum(input)
        q = np.minimum(np.maximum(self.base - total_inventory_position, 0), self.cap)

        print("action:", q)

        return np.array([q])

    def draw_action_train(self, input, base, cap):

        """
        Note: only designed for single product case
        """

        for preprocessor in self.preprocessors:
            # input must be a vector containing the inventory and the pipeline vector
            input = preprocessor(input)
        
        total_inventory_position = np.sum(input)
        q = np.minimum(np.maximum(base - total_inventory_position, 0), cap)

        return np.array([q])

    def reset(self):
        pass
