# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/environments/10_calculation_functions.ipynb.

# %% auto 0
__all__ = ['normalize_reward_FC', 'normalize_reward_NV', 'get_orders_arriving_stochastic', 'eoq_qr_calculations_single_action',
           'eoq_qr_calculations_qr_action', 'mpfc_calculations_ss_action', 'nv_calculations', 'get_fixed_ordering_cost']

# %% ../../nbs/environments/10_calculation_functions.ipynb 4
# General libraries:
import numpy as np
import time

# %% ../../nbs/environments/10_calculation_functions.ipynb 8
# helper functions:
def normalize_reward_FC(self, reward):  

    # print(self.normalize_reward)
    # print("normalize_reward")
    # print(self.normalize_reward)
    if self.normalize_reward:
        # print("doning normalization")
        # print(np.sum(self.holding_cost*self.inventory_cap))
        # print(np.sum(self.fixed_ordering_cost))
        reward_norm = reward / np.maximum(np.sum(self.holding_cost*self.inventory_cap), np.sum(self.fixed_ordering_cost))
        reward_norm = np.clip(reward_norm, -10, 0)
        if reward_norm <= self.min_reward:
            self.min_reward = reward_norm
            print("before:", reward, "after:", reward_norm, "with:", (np.maximum(np.sum(self.holding_cost*self.inventory_cap), np.sum(self.fixed_ordering_cost))))
    else:
        # print("not normalizing")
        reward_norm = reward

    return reward_norm

# %% ../../nbs/environments/10_calculation_functions.ipynb 9
# helper functions:
def normalize_reward_NV(self, reward):

    if self.normalize_reward:

        reward_norm = reward / np.maximum(np.sum(self.underage_cost*self.avg_demand), np.sum(self.overage_cost*self.avg_demand))

        reward_norm = np.clip(reward_norm, -10, 0)

        if reward_norm <= self.min_reward:
            self.min_reward = reward_norm
            print("before:", reward, "after:", reward_norm, "with:", (np.maximum(np.sum(self.underage_cost*self.avg_demand), np.sum(self.overage_cost*self.avg_demand))))
    else:
        # print("not normalizing")
        reward_norm = reward

    return reward_norm

# %% ../../nbs/environments/10_calculation_functions.ipynb 10
def get_orders_arriving_stochastic(self, action):


    # print("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx")
    # print("initial state")

    # print("lead times:", self.lead_time_realizations)
    # print("full order pipeline:", self.full_order_pipeline)
    # print("order pipeline:", self.order_pipeline)

    # print("order placed:", action)
    
    all_orders_arriving = np.zeros(self.num_products) 
            
    for i in range(self.num_products):
        arriving_mask = self.lead_time_realizations[i] == 0
        orders_arriving = np.sum(self.full_order_pipeline[i][arriving_mask])
        self.full_order_pipeline[i][arriving_mask] = 0
        all_orders_arriving[i] = orders_arriving
    
    # print("arriving orders:", all_orders_arriving)
    
    self.full_order_pipeline[:,:-1] = self.full_order_pipeline[:,1:].copy()
    self.full_order_pipeline[:,-1] = action

    # print("full order pipeline after update1:", self.full_order_pipeline)

    self.lead_time_realizations[:,:-1] = self.lead_time_realizations[:,1:].copy()

    if self.lead_time_stochastic == "gamma":
        lead_time_drawn_raw = np.random.gamma(self.lead_time-1, 1, self.num_products)+1
    elif self.lead_time_stochastic == "normal":
        lead_time_drawn_raw = np.random.normal(self.lead_time, self.lead_time_variance, self.num_products)
        lead_time_drawn_raw = np.maximum(lead_time_drawn_raw, 1)
    else:
        raise ValueError("lead time distribution unknown")
    
    self.lead_time_realizations[:,-1] = int(np.round(np.minimum(lead_time_drawn_raw, self.max_lead_time), 0))
    # print("lead time of order:", self.lead_time_realizations[:,-1])
    self.lead_time_realizations = np.maximum(self.lead_time_realizations-1, 0)

    # below also accounts for the c  ase that lag window is larger than max lead time

    # print("full order pipeline after update2:", self.full_order_pipeline)

    if self.max_lead_time > self.max_observation_lead_time:
        self.order_pipeline = self.full_order_pipeline[:,-self.max_observation_lead_time:].copy()
    else:
        self.order_pipeline[:,-self.max_lead_time:] = self.full_order_pipeline.copy()


    # print("full order pipeline after update3:", self.full_order_pipeline)

    self.order_pipeline[:,0] = self.order_pipeline[:,0] + np.sum(self.full_order_pipeline[:,:-self.max_observation_lead_time], axis = 1) 


    # print("final state")

    # print("lead times:", self.lead_time_realizations)  
    # print("full order pipeline:", self.full_order_pipeline)
    # print("order pipeline:", self.order_pipeline)


    # time.sleep(10)

    return all_orders_arriving

# %% ../../nbs/environments/10_calculation_functions.ipynb 11
def eoq_qr_calculations_single_action(self, action):

    # TODO Think about naming of the function

    # print(f"action: {action}")
    # print(f"self.inventory: {self.inventory}")
    # # print pipeline
    # print(f"self.order_pipeline: {self.order_pipeline}")
    # # print demand
    # print(f"self.demand: {self.demand}")

    fixed_ordering_cost_step = get_fixed_ordering_cost(self, action, positive_only = True)

    # todo: add variable ordering cost:
        # variable cost positive order ...
        # variable cost negative order would build on overage cost

    # print("################# new step #################")

    # print("start inventory:", self.inventory)
    # print("start pipeline:", self.order_pipeline)

    # print("pipeline last element:" self.order_pipeline[:,-1])

    # print("action:", action)

    order = action

    if self.use_order_pipeline:
        if self.lead_time_stochastic:

            orders_arriving = get_orders_arriving_stochastic(self, order)
         
        else:
            orders_arriving = self.order_pipeline[:,-1].copy()
            # print("orders_arriving: ", orders_arriving)
            #Move elements one slot to the right in order_pipeline
            self.order_pipeline[:, 1:] = self.order_pipeline[:, :-1]
            product_indices = np.arange(self.num_products)
            insert_indices = -self.lead_time
            self.order_pipeline[product_indices, insert_indices] = order

    else:
        orders_arriving = order

    # print("orders_arriving:", orders_arriving)
    # print("new pipeline:", self.order_pipeline)

    self.inventory += orders_arriving # first, the order will be added to inventory (meaning lead-time is 0)

    # print("inventory after order:", self.inventory)

    self.inventory = np.maximum(self.inventory, 0) # inventory cannot be negative

    self.inventory = np.minimum(self.inventory, self.inventory_cap) # inventory cannot be higher than inventory_cap

    # print("demand:", self.demand[self.period, :])

    self.inventory -= self.demand[self.period, :] # Then demand is subtracted from inventory

    outages = np.where(self.inventory < 0, np.abs(self.inventory), 0)
    penalty_cost_step = self.underage_cost * outages.sum()

    self.inventory = np.maximum(self.inventory, 0) # inventory cannot be negative
    holding_cost_step = self.holding_cost * self.inventory

    total_cost_step = fixed_ordering_cost_step + penalty_cost_step + holding_cost_step
    reward = - total_cost_step.sum() # maximize negative cost (sum over all products)

    # print(f"fixed_ordering_cost_step: {fixed_ordering_cost_step}")
    # print(f"penalty_cost_step: {penalty_cost_step}")
    # print(f"holding_cost_step: {holding_cost_step}")


    # print(f"fixed_ordering_cost_step: {fixed_ordering_cost_step}")
    # print(f"penalty_cost_step: {penalty_cost_step}")
    # print(f"holding_cost_step: {holding_cost_step}")
    # print(f"total_cost_step: {total_cost_step}")

    reward_norm = normalize_reward_FC(self, reward)

    # print("new inventory:", self.inventory)

    # # print(f"reward_norm: {reward_norm}")
    # if self.normalize_reward:
    # time.sleep(2)

    info = dict(
            inventory = self.inventory,
            demand = self.demand[self.period, :],
            action = action,
            order = action,
            penalty = penalty_cost_step,
            holding = holding_cost_step,
            fixed_ordering = fixed_ordering_cost_step,
    )

    self.period += 1

    return reward_norm, info

# %% ../../nbs/environments/10_calculation_functions.ipynb 12
def eoq_qr_calculations_qr_action(self, action):

    # TODO Think about naming of the function
    
    # print("inventory: ", self.inventory)
    # print("pipeline: ", self.order_pipeline)
    # print("action: ", action.shape)
    # print("demand: ", self.demand[self.period, :])

    if len(action.shape) == 1:
        r = action[:self.num_products]
        q = action[self.num_products:]

    # print("r: ", r.shape)
    # print("q: ", q.shape)

    order = np.where(self.inventory <= r, q, 0)

    fixed_ordering_cost_step = get_fixed_ordering_cost(self, order, positive_only = True)

    # print(fixed_ordering_cost_step.shape)

    # print("pipeline becfore: ", self.order_pipeline)

    if self.use_order_pipeline:

        if self.lead_time_stochastic:
            orders_arriving = get_orders_arriving_stochastic(self, order)

        else:
            orders_arriving = self.order_pipeline[:,-1].copy()
            # print("orders_arriving: ", orders_arriving)
            #Move elements one slot to the right in order_pipeline
            self.order_pipeline[:, 1:] = self.order_pipeline[:, :-1]
            product_indices = np.arange(self.num_products)
            insert_indices = -self.lead_time
            self.order_pipeline[product_indices, insert_indices] = order
    else:
        orders_arriving = order

    # print("pipeline after: ", self.order_pipeline)

    # print(np.round(self.order_pipeline,0))




    # todo: add variable ordering cost:
        # variable cost positive order ...
        # variable cost negative order would build on overage cost

    inventory_start = self.inventory.copy()

    self.inventory += orders_arriving # first, the order will be added to inventory (meaning lead-time is 0)

    np.set_printoptions(suppress=True)

    self.inventory = np.maximum(self.inventory, 0) # inventory cannot be negative

    self.inventory = np.minimum(self.inventory, self.inventory_cap) # inventory cannot be higher than inventory_cap

    self.inventory -= self.demand[self.period, :] # Then demand is subtracted from inventory

    outages = np.where(self.inventory < 0, np.abs(self.inventory), 0)
    penalty_cost_step = self.underage_cost * outages


    self.inventory = np.maximum(self.inventory, 0) # inventory cannot be negative
    holding_cost_step = self.holding_cost * self.inventory

    total_cost_step = fixed_ordering_cost_step + penalty_cost_step + holding_cost_step
    reward = - total_cost_step.sum() # maximize negative cost (sum over all products)

    # print("end inventory: ", self.inventory)
    # print("end pipeline: ", self.order_pipeline)

    # print(f"fixed_ordering_cost_step: {fixed_ordering_cost_step}")
    # print(f"penalty_cost_step: {penalty_cost_step}")
    # print(f"holding_cost_step: {holding_cost_step}")
    # print(f"total_cost_step: {total_cost_step}")

    reward_norm = normalize_reward_FC(self, reward)


    # print(f"reward_norm: {reward_norm}")

    # if self.normalize_reward:


    info = dict(
            inventory = inventory_start,
            demand = self.demand[self.period, :],
            action = action,
            order = order,
            penalty = penalty_cost_step,
            holding = holding_cost_step,
            fixed_ordering = fixed_ordering_cost_step,
    )

    self.period += 1

    return reward_norm, info

# %% ../../nbs/environments/10_calculation_functions.ipynb 13
def mpfc_calculations_ss_action(self, action):

    # TODO Think about naming of the function

    if len(action.shape) == 1:
        s = action[0]
        S = action[1]
    else:
        s = action[0][0]
        S = action[0][1]

    total_inventory_position= self.inventory + np.sum(self.order_pipeline, axis = 1)
    q = S - total_inventory_position

    # print("inventory:",  self.inventory)
    # print("order pipeline:", self.order_pipeline)
    # print("total_inventory_position:", total_inventory_position)
    # print("S:", S)
    # print("q:", q)

    order = np.where(self.inventory <= s, q, 0)

    # print("order:", order)

    fixed_ordering_cost_step = get_fixed_ordering_cost(self, order, positive_only = True)

    if self.use_order_pipeline:
        orders_arriving = self.order_pipeline[:,-1].copy()
            # Move elements one slot to the right in order_pipeline
        self.order_pipeline[:, 1:] = self.order_pipeline[:, :-1]
        self.order_pipeline[:, 0] = order
    else:
        orders_arriving = order

    # todo: add variable ordering cost:
        # variable cost positive order ...
        # variable cost negative order would build on overage cost

    inventory_start = self.inventory.copy()

    self.inventory += orders_arriving # first, the order will be added to inventory (meaning lead-time is 0)

    self.inventory = np.maximum(self.inventory, 0) # inventory cannot be negative

    self.inventory = np.minimum(self.inventory, self.inventory_cap) # inventory cannot be higher than inventory_cap

    self.inventory -= self.demand[self.period, :] # Then demand is subtracted from inventory

    outages = np.where(self.inventory < 0, np.abs(self.inventory), 0)
    penalty_cost_step = self.underage_cost * outages.sum()

    self.inventory = np.maximum(self.inventory, 0) # inventory cannot be negative
    holding_cost_step = self.holding_cost * self.inventory

    total_cost_step = fixed_ordering_cost_step + penalty_cost_step + holding_cost_step
    reward = - total_cost_step.sum() # maximize negative cost (sum over all products)

    reward_norm = normalize_reward_FC(self, reward)

    info = dict(
            inventory = inventory_start,
            demand = self.demand[self.period, :],
            action = action,
            order = order,
            penalty = penalty_cost_step,
            holding = holding_cost_step,
            fixed_ordering = fixed_ordering_cost_step,
    )

    self.period += 1

    return reward_norm, info

# %% ../../nbs/environments/10_calculation_functions.ipynb 14
def nv_calculations(self, action):
    
    if self.full_demand:
        demand = self.demand[self.period+self.full_additional_context, :]
    else:
        demand = self.demand[self.period, :]

    # assert that action is not nan

    if self.mask is not None:
        mask = self.mask.copy()
        if len(self.mask.shape)==1:
            mask = mask.reshape(-1,1)
        mask = mask[self.period, :]
        action = action*mask

    # action is minimum of 0 and action
    action = np.maximum(action, 0)

    remaining_demand = demand.copy() - action

    underage_cost = self.underage_cost.reshape(1, -1)  # Reshape to (1, product)
    overage_cost = self.overage_cost.reshape(1, -1) 

    cost = np.where(remaining_demand > 0, underage_cost * remaining_demand, 0)
    cost += np.where(remaining_demand < 0, overage_cost * np.abs(remaining_demand), 0)
    # cost[remaining_demand > 0] = self.underage_cost * remaining_demand[remaining_demand > 0]
    # cost[remaining_demand < 0] = self.overage_cost * np.abs(remaining_demand[remaining_demand < 0])

    reward = -cost.sum()

    # print("demand:", demand, "action:", action, "reward: ", reward)

    reward_norm = normalize_reward_NV(self, reward)

    # print(demand[32], action[32])
    
    info = dict(
        demand = demand,
        cost = cost
    )

    self.period += 1

    return reward_norm, info

# %% ../../nbs/environments/10_calculation_functions.ipynb 16
def get_fixed_ordering_cost(self, action, positive_only = True):
    """ 
    Returns the fixed ordering cost for the period.

    Parameters
    ----------
    action : array
        The action taken in the period.
    positive_only : bool
        If True, only positive orders will be charged the fixed ordering cost. Otherwise, all non-zero orders will be charged the fixed ordering cost.

    Returns
    -------
    fixed_ordering_cost : array
        The fixed ordering cost for the period per product.
    
    """

    if positive_only:
        positive_order = np.where(action > 0, 1, 0)
    else:
        # nonzero = 1, zero = 0
        positive_order = np.where(action != 0, 1, 0)
    return self.fixed_ordering_cost * positive_order
