# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/dt/workday.ipynb.

# %% auto 0
__all__ = ['get_holidays_filter_for_datetime_index', 'get_holidays_calendar_from_ds', 'get_busdays_filter_for_datetime_index',
           'get_weekends_filter_for_datetime_index', 'get_busdays_filter_with_np_is_busday']

# %% ../../nbs/dt/workday.ipynb 3
from typing import Iterable
import holidays

import numpy as np
import pandas as pd

import enilm.etypes.arr
import enilm.datasets
from ..checkers import check_vaild_arr_type, ArrTypes

# %% ../../nbs/dt/workday.ipynb 5
def get_holidays_filter_for_datetime_index(
    dates: pd.DatetimeIndex,
    holidays_calendar: holidays.HolidayBase,
) -> enilm.etypes.arr.NPBool:
    res = np.array(dates.map(lambda date: date in holidays_calendar), dtype=bool)
    assert check_vaild_arr_type(res, ArrTypes.NPBool)
    return res

# %% ../../nbs/dt/workday.ipynb 7
def get_holidays_calendar_from_ds(ds: enilm.etypes.DatasetID) -> holidays.HolidayBase:
    country_code = enilm.datasets.get_country_code_of_dataset(ds)
    return holidays.country_holidays(country_code)


# %% ../../nbs/dt/workday.ipynb 17
def get_busdays_filter_for_datetime_index(
    dates: pd.DatetimeIndex,
) -> enilm.etypes.arr.NPBool:
    # https://stackoverflow.com/a/32280195/1617883
    try:
        res = dates.day_of_week // 5 != 1
    except AttributeError:
        res = dates.dayofweek // 5 != 1
    assert check_vaild_arr_type(res, ArrTypes.NPBool)
    return res


def get_weekends_filter_for_datetime_index(
    dates: pd.DatetimeIndex,
) -> enilm.etypes.arr.NPBool:
    try:
        res = dates.day_of_week >= 5
    except AttributeError:
        res = dates.dayofweek >= 5
    assert check_vaild_arr_type(res, ArrTypes.NPBool)
    return res

def get_busdays_filter_with_np_is_busday(
    dates: enilm.etypes.arr.NPTimestamps,
    np_is_busday_kwargs={},
) -> enilm.etypes.arr.NPBool:
    """
    np_is_busday_kwargs passed to np.is_busday
    """
    assert check_vaild_arr_type(dates, ArrTypes.NPTimestamps)

    busdays_filter: Iterable[bool] = np.is_busday(
        dates=dates.astype("datetime64[D]"),
        **np_is_busday_kwargs,
    )
    assert check_vaild_arr_type(busdays_filter, ArrTypes.NPBool)
    return busdays_filter

