# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_utils/00_utils.ipynb.

# %% auto 0
__all__ = ['set_env_parameter']

# %% ../nbs/00_utils/00_utils.ipynb 3
import numpy as np

# %% ../nbs/00_utils/00_utils.ipynb 5
def set_env_parameter(
                    input,
                    num_SKUs):

    # check if input is scalar
    if isinstance(input, Parameter):
        assert input.shape == (num_SKUs,), "Parameter shape must be equal to the number of SKUs"
        output = input
        
    elif isinstance(input, (int, float)):
        if num_SKUs == 1:
            output = np.array([input])
        else:
            output = np.array([input]*num_SKUs)
    # check in numpy array and dimension 0 is equal to num_SKUs
    elif isinstance(input, np.ndarray):
        if input.shape[0] == num_SKUs:
            output = input
        elif input.shape[0] == 1:
            output = np.array([input[0]]*num_SKUs)
        else:
            raise ValueError("Input array must have the same length as the number of SKUs")
    else:
        raise TypeError("Input must be a scalar or numpy array")
    
    return output
