/*
 * This file implements business day functionality for NumPy datetime.
 *
 * Written by Mark Wiebe (mwwiebe@gmail.com)
 * Copyright (c) 2011 by Enthought, Inc.
 *
 * See LICENSE.txt for the license.
 */

#define PY_SSIZE_T_CLEAN
#include <Python.h>

#define NPY_NO_DEPRECATED_API NPY_API_VERSION
#define _MULTIARRAYMODULE
#include <numpy/arrayobject.h>

#include "npy_config.h"
#include "npy_pycompat.h"

#include "numpy/arrayscalars.h"
#include "lowlevel_strided_loops.h"
#include "_datetime.h"
#include "datetime_busday.h"
#include "datetime_busdaycal.h"

/* Gets the day of the week for a datetime64[D] value */
static int
get_day_of_week(npy_datetime date)
{
    int day_of_week;

    /* Get the day of the week for 'date' (1970-01-05 is Monday) */
    day_of_week = (int)((date - 4) % 7);
    if (day_of_week < 0) {
        day_of_week += 7;
    }

    return day_of_week;
}

/*
 * Returns 1 if the date is a holiday (contained in the sorted
 * list of dates), 0 otherwise.
 *
 * The holidays list should be normalized, which means any NaT (not-a-time)
 * values, duplicates, and dates already excluded by the weekmask should
 * be removed, and the list should be sorted.
 */
static int
is_holiday(npy_datetime date,
            npy_datetime *holidays_begin, npy_datetime *holidays_end)
{
    npy_datetime *trial;

    /* Simple binary search */
    while (holidays_begin < holidays_end) {
        trial = holidays_begin + (holidays_end - holidays_begin) / 2;

        if (date < *trial) {
            holidays_end = trial;
        }
        else if (date > *trial) {
            holidays_begin = trial + 1;
        }
        else {
            return 1;
        }
    }

    /* Not found */
    return 0;
}

/*
 * Finds the earliest holiday which is on or after 'date'. If 'date' does not
 * appear within the holiday range, returns 'holidays_begin' if 'date'
 * is before all holidays, or 'holidays_end' if 'date' is after all
 * holidays.
 *
 * To remove all the holidays before 'date' from a holiday range, do:
 *
 *      holidays_begin = find_holiday_earliest_on_or_after(date,
 *                                          holidays_begin, holidays_end);
 *
 * The holidays list should be normalized, which means any NaT (not-a-time)
 * values, duplicates, and dates already excluded by the weekmask should
 * be removed, and the list should be sorted.
 */
static npy_datetime *
find_earliest_holiday_on_or_after(npy_datetime date,
            npy_datetime *holidays_begin, npy_datetime *holidays_end)
{
    npy_datetime *trial;

    /* Simple binary search */
    while (holidays_begin < holidays_end) {
        trial = holidays_begin + (holidays_end - holidays_begin) / 2;

        if (date < *trial) {
            holidays_end = trial;
        }
        else if (date > *trial) {
            holidays_begin = trial + 1;
        }
        else {
            return trial;
        }
    }

    return holidays_begin;
}

/*
 * Finds the earliest holiday which is after 'date'. If 'date' does not
 * appear within the holiday range, returns 'holidays_begin' if 'date'
 * is before all holidays, or 'holidays_end' if 'date' is after all
 * holidays.
 *
 * To remove all the holidays after 'date' from a holiday range, do:
 *
 *      holidays_end = find_holiday_earliest_after(date,
 *                                          holidays_begin, holidays_end);
 *
 * The holidays list should be normalized, which means any NaT (not-a-time)
 * values, duplicates, and dates already excluded by the weekmask should
 * be removed, and the list should be sorted.
 */
static npy_datetime *
find_earliest_holiday_after(npy_datetime date,
            npy_datetime *holidays_begin, npy_datetime *holidays_end)
{
    npy_datetime *trial;

    /* Simple binary search */
    while (holidays_begin < holidays_end) {
        trial = holidays_begin + (holidays_end - holidays_begin) / 2;

        if (date < *trial) {
            holidays_end = trial;
        }
        else if (date > *trial) {
            holidays_begin = trial + 1;
        }
        else {
            return trial + 1;
        }
    }

    return holidays_begin;
}

/*
 * Applies the 'roll' strategy to 'date', placing the result in 'out'
 * and setting 'out_day_of_week' to the day of the week that results.
 *
 * Returns 0 on success, -1 on failure.
 */
static int
apply_business_day_roll(npy_datetime date, npy_datetime *out,
                    int *out_day_of_week,
                    NPY_BUSDAY_ROLL roll,
                    npy_bool *weekmask,
                    npy_datetime *holidays_begin, npy_datetime *holidays_end)
{
    int day_of_week;

    /* Deal with NaT input */
    if (date == NPY_DATETIME_NAT) {
        *out = NPY_DATETIME_NAT;
        if (roll == NPY_BUSDAY_RAISE) {
            PyErr_SetString(PyExc_ValueError,
                    "NaT input in busday_offset");
            return -1;
        }
        else {
            return 0;
        }
    }

    /* Get the day of the week for 'date' */
    day_of_week = get_day_of_week(date);

    /* Apply the 'roll' if it's not a business day */
    if (weekmask[day_of_week] == 0 ||
                        is_holiday(date, holidays_begin, holidays_end)) {
        npy_datetime start_date = date;
        int start_day_of_week = day_of_week;

        switch (roll) {
            case NPY_BUSDAY_FOLLOWING:
            case NPY_BUSDAY_MODIFIEDFOLLOWING: {
                do {
                    ++date;
                    if (++day_of_week == 7) {
                        day_of_week = 0;
                    }
                } while (weekmask[day_of_week] == 0 ||
                            is_holiday(date, holidays_begin, holidays_end));

                if (roll == NPY_BUSDAY_MODIFIEDFOLLOWING) {
                    /* If we crossed a month boundary, do preceding instead */
                    if (days_to_month_number(start_date) !=
                                days_to_month_number(date)) {
                        date = start_date;
                        day_of_week = start_day_of_week;

                        do {
                            --date;
                            if (--day_of_week == -1) {
                                day_of_week = 6;
                            }
                        } while (weekmask[day_of_week] == 0 ||
                            is_holiday(date, holidays_begin, holidays_end));
                    }
                }
                break;
            }
            case NPY_BUSDAY_PRECEDING:
            case NPY_BUSDAY_MODIFIEDPRECEDING: {
                do {
                    --date;
                    if (--day_of_week == -1) {
                        day_of_week = 6;
                    }
                } while (weekmask[day_of_week] == 0 ||
                            is_holiday(date, holidays_begin, holidays_end));

                if (roll == NPY_BUSDAY_MODIFIEDPRECEDING) {
                    /* If we crossed a month boundary, do following instead */
                    if (days_to_month_number(start_date) !=
                                days_to_month_number(date)) {
                        date = start_date;
                        day_of_week = start_day_of_week;

                        do {
                            ++date;
                            if (++day_of_week == 7) {
                                day_of_week = 0;
                            }
                        } while (weekmask[day_of_week] == 0 ||
                            is_holiday(date, holidays_begin, holidays_end));
                    }
                }
                break;
            }
            case NPY_BUSDAY_NAT: {
                date = NPY_DATETIME_NAT;
                break;
            }
            case NPY_BUSDAY_RAISE: {
                *out = NPY_DATETIME_NAT;
                PyErr_SetString(PyExc_ValueError,
                        "Non-business day date in busday_offset");
                return -1;
            }
        }
    }

    *out = date;
    *out_day_of_week = day_of_week;

    return 0;
}

/*
 * Applies a single business day offset. See the function
 * business_day_offset for the meaning of all the parameters.
 *
 * Returns 0 on success, -1 on failure.
 */
static int
apply_business_day_offset(npy_datetime date, npy_int64 offset,
                    npy_datetime *out,
                    NPY_BUSDAY_ROLL roll,
                    npy_bool *weekmask, int busdays_in_weekmask,
                    npy_datetime *holidays_begin, npy_datetime *holidays_end)
{
    int day_of_week = 0;
    npy_datetime *holidays_temp;

    /* Roll the date to a business day */
    if (apply_business_day_roll(date, &date, &day_of_week,
                                roll,
                                weekmask,
                                holidays_begin, holidays_end) < 0) {
        return -1;
    }

    /* If we get a NaT, just return it */
    if (date == NPY_DATETIME_NAT) {
        return 0;
    }

    /* Now we're on a valid business day */
    if (offset > 0) {
        /* Remove any earlier holidays */
        holidays_begin = find_earliest_holiday_on_or_after(date,
                                            holidays_begin, holidays_end);

        /* Jump by as many weeks as we can */
        date += (offset / busdays_in_weekmask) * 7;
        offset = offset % busdays_in_weekmask;

        /* Adjust based on the number of holidays we crossed */
        holidays_temp = find_earliest_holiday_after(date,
                                            holidays_begin, holidays_end);
        offset += holidays_temp - holidays_begin;
        holidays_begin = holidays_temp;

        /* Step until we use up the rest of the offset */
        while (offset > 0) {
            ++date;
            if (++day_of_week == 7) {
                day_of_week = 0;
            }
            if (weekmask[day_of_week] && !is_holiday(date,
                                            holidays_begin, holidays_end)) {
                offset--;
            }
        }
    }
    else if (offset < 0) {
        /* Remove any later holidays */
        holidays_end = find_earliest_holiday_after(date,
                                            holidays_begin, holidays_end);

        /* Jump by as many weeks as we can */
        date += (offset / busdays_in_weekmask) * 7;
        offset = offset % busdays_in_weekmask;

        /* Adjust based on the number of holidays we crossed */
        holidays_temp = find_earliest_holiday_on_or_after(date,
                                            holidays_begin, holidays_end);
        offset -= holidays_end - holidays_temp;
        holidays_end = holidays_temp;

        /* Step until we use up the rest of the offset */
        while (offset < 0) {
            --date;
            if (--day_of_week == -1) {
                day_of_week = 6;
            }
            if (weekmask[day_of_week] && !is_holiday(date,
                                            holidays_begin, holidays_end)) {
                offset++;
            }
        }
    }

    *out = date;
    return 0;
}

/*
 * Applies a single business day count operation. See the function
 * business_day_count for the meaning of all the parameters.
 *
 * Returns 0 on success, -1 on failure.
 */
static int
apply_business_day_count(npy_datetime date_begin, npy_datetime date_end,
                    npy_int64 *out,
                    npy_bool *weekmask, int busdays_in_weekmask,
                    npy_datetime *holidays_begin, npy_datetime *holidays_end)
{
    npy_int64 count, whole_weeks;
    int day_of_week = 0;
    int swapped = 0;

    /* If we get a NaT, raise an error */
    if (date_begin == NPY_DATETIME_NAT || date_end == NPY_DATETIME_NAT) {
        PyErr_SetString(PyExc_ValueError,
                "Cannot compute a business day count with a NaT (not-a-time) "
                "date");
        return -1;
    }

    /* Trivial empty date range */
    if (date_begin == date_end) {
        *out = 0;
        return 0;
    }
    else if (date_begin > date_end) {
        npy_datetime tmp = date_begin;
        date_begin = date_end;
        date_end = tmp;
        swapped = 1;
    }

    /* Remove any earlier holidays */
    holidays_begin = find_earliest_holiday_on_or_after(date_begin,
                                        holidays_begin, holidays_end);
    /* Remove any later holidays */
    holidays_end = find_earliest_holiday_on_or_after(date_end,
                                        holidays_begin, holidays_end);

    /* Start the count as negative the number of holidays in the range */
    count = -(holidays_end - holidays_begin);

    /* Add the whole weeks between date_begin and date_end */
    whole_weeks = (date_end - date_begin) / 7;
    count += whole_weeks * busdays_in_weekmask;
    date_begin += whole_weeks * 7;

    if (date_begin < date_end) {
        /* Get the day of the week for 'date_begin' */
        day_of_week = get_day_of_week(date_begin);

        /* Count the remaining days one by one */
        while (date_begin < date_end) {
            if (weekmask[day_of_week]) {
                count++;
            }
            ++date_begin;
            if (++day_of_week == 7) {
                day_of_week = 0;
            }
        }
    }

    if (swapped) {
        count = -count;
    }

    *out = count;
    return 0;
}

/*
 * Applies the given offsets in business days to the dates provided.
 * This is the low-level function which requires already cleaned input
 * data.
 *
 * dates:    An array of dates with 'datetime64[D]' data type.
 * offsets:  An array safely convertible into type int64.
 * out:      Either NULL, or an array with 'datetime64[D]' data type
 *              in which to place the resulting dates.
 * roll:     A rule for how to treat non-business day dates.
 * weekmask: A 7-element boolean mask, 1 for possible business days and 0
 *              for non-business days.
 * busdays_in_weekmask: A count of how many 1's there are in weekmask.
 * holidays_begin/holidays_end: A sorted list of dates matching '[D]'
 *           unit metadata, with any dates falling on a day of the
 *           week without weekmask[i] == 1 already filtered out.
 *
 * For each (date, offset) in the broadcasted pair of (dates, offsets),
 * does the following:
 *  + Applies the 'roll' rule to the date to either produce NaT, raise
 *    an exception, or land on a valid business day.
 *  + Adds 'offset' business days to the valid business day found.
 *  + Sets the value in 'out' if provided, or the allocated output array
 *    otherwise.
 */
NPY_NO_EXPORT PyArrayObject *
business_day_offset(PyArrayObject *dates, PyArrayObject *offsets,
                    PyArrayObject *out,
                    NPY_BUSDAY_ROLL roll,
                    npy_bool *weekmask, int busdays_in_weekmask,
                    npy_datetime *holidays_begin, npy_datetime *holidays_end)
{
    PyArray_DatetimeMetaData temp_meta;
    PyArray_Descr *dtypes[3] = {NULL, NULL, NULL};

    NpyIter *iter = NULL;
    PyArrayObject *op[3] = {NULL, NULL, NULL};
    npy_uint32 op_flags[3], flags;

    PyArrayObject *ret = NULL;

    if (busdays_in_weekmask == 0) {
        PyErr_SetString(PyExc_ValueError,
                "the business day weekmask must have at least one "
                "valid business day");
        return NULL;
    }

    /* First create the data types for dates and offsets */
    temp_meta.base = NPY_FR_D;
    temp_meta.num = 1;
    dtypes[0] = create_datetime_dtype(NPY_DATETIME, &temp_meta);
    if (dtypes[0] == NULL) {
        goto fail;
    }
    dtypes[1] = PyArray_DescrFromType(NPY_INT64);
    if (dtypes[1] == NULL) {
        goto fail;
    }
    dtypes[2] = dtypes[0];
    Py_INCREF(dtypes[2]);

    /* Set up the iterator parameters */
    flags = NPY_ITER_EXTERNAL_LOOP|
            NPY_ITER_BUFFERED|
            NPY_ITER_ZEROSIZE_OK;
    op[0] = dates;
    op_flags[0] = NPY_ITER_READONLY | NPY_ITER_ALIGNED;
    op[1] = offsets;
    op_flags[1] = NPY_ITER_READONLY | NPY_ITER_ALIGNED;
    op[2] = out;
    op_flags[2] = NPY_ITER_WRITEONLY | NPY_ITER_ALLOCATE | NPY_ITER_ALIGNED;

    /* Allocate the iterator */
    iter = NpyIter_MultiNew(3, op, flags, NPY_KEEPORDER, NPY_SAFE_CASTING,
                            op_flags, dtypes);
    if (iter == NULL) {
        goto fail;
    }

    /* Loop over all elements */
    if (NpyIter_GetIterSize(iter) > 0) {
        NpyIter_IterNextFunc *iternext;
        char **dataptr;
        npy_intp *strideptr, *innersizeptr;

        iternext = NpyIter_GetIterNext(iter, NULL);
        if (iternext == NULL) {
            goto fail;
        }
        dataptr = NpyIter_GetDataPtrArray(iter);
        strideptr = NpyIter_GetInnerStrideArray(iter);
        innersizeptr = NpyIter_GetInnerLoopSizePtr(iter);

        do {
            char *data_dates = dataptr[0];
            char *data_offsets = dataptr[1];
            char *data_out = dataptr[2];
            npy_intp stride_dates = strideptr[0];
            npy_intp stride_offsets = strideptr[1];
            npy_intp stride_out = strideptr[2];
            npy_intp count = *innersizeptr;

            while (count--) {
                if (apply_business_day_offset(*(npy_int64 *)data_dates,
                                       *(npy_int64 *)data_offsets,
                                       (npy_int64 *)data_out,
                                       roll,
                                       weekmask, busdays_in_weekmask,
                                       holidays_begin, holidays_end) < 0) {
                    goto fail;
                }

                data_dates += stride_dates;
                data_offsets += stride_offsets;
                data_out += stride_out;
            }
        } while (iternext(iter));
    }

    /* Get the return object from the iterator */
    ret = NpyIter_GetOperandArray(iter)[2];
    Py_INCREF(ret);

    goto finish;

fail:
    Py_XDECREF(ret);
    ret = NULL;

finish:
    Py_XDECREF(dtypes[0]);
    Py_XDECREF(dtypes[1]);
    Py_XDECREF(dtypes[2]);
    if (iter != NULL) {
        if (NpyIter_Deallocate(iter) != NPY_SUCCEED) {
            Py_XDECREF(ret);
            ret = NULL;
        }
    }
    return ret;
}

/*
 * Counts the number of business days between two dates, not including
 * the end date. This is the low-level function which requires already
 * cleaned input data.
 *
 * If dates_begin is before dates_end, the result is positive.  If
 * dates_begin is after dates_end, it is negative.
 *
 * dates_begin:  An array of dates with 'datetime64[D]' data type.
 * dates_end:    An array of dates with 'datetime64[D]' data type.
 * out:      Either NULL, or an array with 'int64' data type
 *              in which to place the resulting dates.
 * weekmask: A 7-element boolean mask, 1 for possible business days and 0
 *              for non-business days.
 * busdays_in_weekmask: A count of how many 1's there are in weekmask.
 * holidays_begin/holidays_end: A sorted list of dates matching '[D]'
 *           unit metadata, with any dates falling on a day of the
 *           week without weekmask[i] == 1 already filtered out.
 */
NPY_NO_EXPORT PyArrayObject *
business_day_count(PyArrayObject *dates_begin, PyArrayObject *dates_end,
                    PyArrayObject *out,
                    npy_bool *weekmask, int busdays_in_weekmask,
                    npy_datetime *holidays_begin, npy_datetime *holidays_end)
{
    PyArray_DatetimeMetaData temp_meta;
    PyArray_Descr *dtypes[3] = {NULL, NULL, NULL};

    NpyIter *iter = NULL;
    PyArrayObject *op[3] = {NULL, NULL, NULL};
    npy_uint32 op_flags[3], flags;

    PyArrayObject *ret = NULL;

    if (busdays_in_weekmask == 0) {
        PyErr_SetString(PyExc_ValueError,
                "the business day weekmask must have at least one "
                "valid business day");
        return NULL;
    }

    /* First create the data types for the dates and the int64 output */
    temp_meta.base = NPY_FR_D;
    temp_meta.num = 1;
    dtypes[0] = create_datetime_dtype(NPY_DATETIME, &temp_meta);
    if (dtypes[0] == NULL) {
        goto fail;
    }
    dtypes[1] = dtypes[0];
    Py_INCREF(dtypes[1]);
    dtypes[2] = PyArray_DescrFromType(NPY_INT64);
    if (dtypes[2] == NULL) {
        goto fail;
    }

    /* Set up the iterator parameters */
    flags = NPY_ITER_EXTERNAL_LOOP|
            NPY_ITER_BUFFERED|
            NPY_ITER_ZEROSIZE_OK;
    op[0] = dates_begin;
    op_flags[0] = NPY_ITER_READONLY | NPY_ITER_ALIGNED;
    op[1] = dates_end;
    op_flags[1] = NPY_ITER_READONLY | NPY_ITER_ALIGNED;
    op[2] = out;
    op_flags[2] = NPY_ITER_WRITEONLY | NPY_ITER_ALLOCATE | NPY_ITER_ALIGNED;

    /* Allocate the iterator */
    iter = NpyIter_MultiNew(3, op, flags, NPY_KEEPORDER, NPY_SAFE_CASTING,
                            op_flags, dtypes);
    if (iter == NULL) {
        goto fail;
    }

    /* Loop over all elements */
    if (NpyIter_GetIterSize(iter) > 0) {
        NpyIter_IterNextFunc *iternext;
        char **dataptr;
        npy_intp *strideptr, *innersizeptr;

        iternext = NpyIter_GetIterNext(iter, NULL);
        if (iternext == NULL) {
            goto fail;
        }
        dataptr = NpyIter_GetDataPtrArray(iter);
        strideptr = NpyIter_GetInnerStrideArray(iter);
        innersizeptr = NpyIter_GetInnerLoopSizePtr(iter);

        do {
            char *data_dates_begin = dataptr[0];
            char *data_dates_end = dataptr[1];
            char *data_out = dataptr[2];
            npy_intp stride_dates_begin = strideptr[0];
            npy_intp stride_dates_end = strideptr[1];
            npy_intp stride_out = strideptr[2];
            npy_intp count = *innersizeptr;

            while (count--) {
                if (apply_business_day_count(*(npy_int64 *)data_dates_begin,
                                       *(npy_int64 *)data_dates_end,
                                       (npy_int64 *)data_out,
                                       weekmask, busdays_in_weekmask,
                                       holidays_begin, holidays_end) < 0) {
                    goto fail;
                }

                data_dates_begin += stride_dates_begin;
                data_dates_end += stride_dates_end;
                data_out += stride_out;
            }
        } while (iternext(iter));
    }

    /* Get the return object from the iterator */
    ret = NpyIter_GetOperandArray(iter)[2];
    Py_INCREF(ret);

    goto finish;

fail:
    Py_XDECREF(ret);
    ret = NULL;

finish:
    Py_XDECREF(dtypes[0]);
    Py_XDECREF(dtypes[1]);
    Py_XDECREF(dtypes[2]);
    if (iter != NULL) {
        if (NpyIter_Deallocate(iter) != NPY_SUCCEED) {
            Py_XDECREF(ret);
            ret = NULL;
        }
    }
    return ret;
}

/*
 * Returns a boolean array with True for input dates which are valid
 * business days, and False for dates which are not. This is the
 * low-level function which requires already cleaned input data.
 *
 * dates:  An array of dates with 'datetime64[D]' data type.
 * out:      Either NULL, or an array with 'bool' data type
 *              in which to place the resulting dates.
 * weekmask: A 7-element boolean mask, 1 for possible business days and 0
 *              for non-business days.
 * busdays_in_weekmask: A count of how many 1's there are in weekmask.
 * holidays_begin/holidays_end: A sorted list of dates matching '[D]'
 *           unit metadata, with any dates falling on a day of the
 *           week without weekmask[i] == 1 already filtered out.
 */
NPY_NO_EXPORT PyArrayObject *
is_business_day(PyArrayObject *dates, PyArrayObject *out,
                    npy_bool *weekmask, int busdays_in_weekmask,
                    npy_datetime *holidays_begin, npy_datetime *holidays_end)
{
    PyArray_DatetimeMetaData temp_meta;
    PyArray_Descr *dtypes[2] = {NULL, NULL};

    NpyIter *iter = NULL;
    PyArrayObject *op[2] = {NULL, NULL};
    npy_uint32 op_flags[2], flags;

    PyArrayObject *ret = NULL;

    if (busdays_in_weekmask == 0) {
        PyErr_SetString(PyExc_ValueError,
                "the business day weekmask must have at least one "
                "valid business day");
        return NULL;
    }

    /* First create the data types for the dates and the bool output */
    temp_meta.base = NPY_FR_D;
    temp_meta.num = 1;
    dtypes[0] = create_datetime_dtype(NPY_DATETIME, &temp_meta);
    if (dtypes[0] == NULL) {
        goto fail;
    }
    dtypes[1] = PyArray_DescrFromType(NPY_BOOL);
    if (dtypes[1] == NULL) {
        goto fail;
    }

    /* Set up the iterator parameters */
    flags = NPY_ITER_EXTERNAL_LOOP|
            NPY_ITER_BUFFERED|
            NPY_ITER_ZEROSIZE_OK;
    op[0] = dates;
    op_flags[0] = NPY_ITER_READONLY | NPY_ITER_ALIGNED;
    op[1] = out;
    op_flags[1] = NPY_ITER_WRITEONLY | NPY_ITER_ALLOCATE | NPY_ITER_ALIGNED;

    /* Allocate the iterator */
    iter = NpyIter_MultiNew(2, op, flags, NPY_KEEPORDER, NPY_SAFE_CASTING,
                            op_flags, dtypes);
    if (iter == NULL) {
        goto fail;
    }

    /* Loop over all elements */
    if (NpyIter_GetIterSize(iter) > 0) {
        NpyIter_IterNextFunc *iternext;
        char **dataptr;
        npy_intp *strideptr, *innersizeptr;

        iternext = NpyIter_GetIterNext(iter, NULL);
        if (iternext == NULL) {
            goto fail;
        }
        dataptr = NpyIter_GetDataPtrArray(iter);
        strideptr = NpyIter_GetInnerStrideArray(iter);
        innersizeptr = NpyIter_GetInnerLoopSizePtr(iter);

        do {
            char *data_dates = dataptr[0];
            char *data_out = dataptr[1];
            npy_intp stride_dates = strideptr[0];
            npy_intp stride_out = strideptr[1];
            npy_intp count = *innersizeptr;

            npy_datetime date;
            int day_of_week;

            while (count--) {
                /* Check if it's a business day */
                date = *(npy_datetime *)data_dates;
                day_of_week = get_day_of_week(date);
                *(npy_bool *)data_out = weekmask[day_of_week] &&
                                        !is_holiday(date,
                                            holidays_begin, holidays_end) &&
                                        date != NPY_DATETIME_NAT;

                data_dates += stride_dates;
                data_out += stride_out;
            }
        } while (iternext(iter));
    }

    /* Get the return object from the iterator */
    ret = NpyIter_GetOperandArray(iter)[1];
    Py_INCREF(ret);

    goto finish;

fail:
    Py_XDECREF(ret);
    ret = NULL;

finish:
    Py_XDECREF(dtypes[0]);
    Py_XDECREF(dtypes[1]);
    if (iter != NULL) {
        if (NpyIter_Deallocate(iter) != NPY_SUCCEED) {
            Py_XDECREF(ret);
            ret = NULL;
        }
    }
    return ret;
}

static int
PyArray_BusDayRollConverter(PyObject *roll_in, NPY_BUSDAY_ROLL *roll)
{
    PyObject *obj = roll_in;
    char *str;
    Py_ssize_t len;

    /* Make obj into an ASCII string */
    Py_INCREF(obj);
    if (PyUnicode_Check(obj)) {
        /* accept unicode input */
        PyObject *obj_str;
        obj_str = PyUnicode_AsASCIIString(obj);
        if (obj_str == NULL) {
            Py_DECREF(obj);
            return 0;
        }
        Py_DECREF(obj);
        obj = obj_str;
    }

    if (PyBytes_AsStringAndSize(obj, &str, &len) < 0) {
        Py_DECREF(obj);
        return 0;
    }

    /* Use switch statements to quickly isolate the right enum value */
    switch (str[0]) {
        case 'b':
            if (strcmp(str, "backward") == 0) {
                *roll = NPY_BUSDAY_BACKWARD;
                goto finish;
            }
            break;
        case 'f':
            if (len > 2) switch (str[2]) {
                case 'r':
                    if (strcmp(str, "forward") == 0) {
                        *roll = NPY_BUSDAY_FORWARD;
                        goto finish;
                    }
                    break;
                case 'l':
                    if (strcmp(str, "following") == 0) {
                        *roll = NPY_BUSDAY_FOLLOWING;
                        goto finish;
                    }
                    break;
            }
            break;
        case 'm':
            if (len > 8) switch (str[8]) {
                case 'f':
                    if (strcmp(str, "modifiedfollowing") == 0) {
                        *roll = NPY_BUSDAY_MODIFIEDFOLLOWING;
                        goto finish;
                    }
                    break;
                case 'p':
                    if (strcmp(str, "modifiedpreceding") == 0) {
                        *roll = NPY_BUSDAY_MODIFIEDPRECEDING;
                        goto finish;
                    }
                    break;
            }
            break;
        case 'n':
            if (strcmp(str, "nat") == 0) {
                *roll = NPY_BUSDAY_NAT;
                goto finish;
            }
            break;
        case 'p':
            if (strcmp(str, "preceding") == 0) {
                *roll = NPY_BUSDAY_PRECEDING;
                goto finish;
            }
            break;
        case 'r':
            if (strcmp(str, "raise") == 0) {
                *roll = NPY_BUSDAY_RAISE;
                goto finish;
            }
            break;
    }

    PyErr_Format(PyExc_ValueError,
            "Invalid business day roll parameter \"%s\"",
            str);
    Py_DECREF(obj);
    return 0;

finish:
    Py_DECREF(obj);
    return 1;
}

/*
 * This is the 'busday_offset' function exposed for calling
 * from Python.
 */
NPY_NO_EXPORT PyObject *
array_busday_offset(PyObject *NPY_UNUSED(self),
                      PyObject *args, PyObject *kwds)
{
    char *kwlist[] = {"dates", "offsets", "roll",
                      "weekmask", "holidays", "busdaycal", "out", NULL};

    PyObject *dates_in = NULL, *offsets_in = NULL, *out_in = NULL;

    PyArrayObject *dates = NULL, *offsets = NULL, *out = NULL, *ret;
    NPY_BUSDAY_ROLL roll = NPY_BUSDAY_RAISE;
    npy_bool weekmask[7] = {2, 1, 1, 1, 1, 0, 0};
    NpyBusDayCalendar *busdaycal = NULL;
    int i, busdays_in_weekmask;
    npy_holidayslist holidays = {NULL, NULL};
    int allocated_holidays = 1;

    if (!PyArg_ParseTupleAndKeywords(args, kwds,
                                    "OO|O&O&O&O!O:busday_offset", kwlist,
                                    &dates_in,
                                    &offsets_in,
                                    &PyArray_BusDayRollConverter, &roll,
                                    &PyArray_WeekMaskConverter, &weekmask[0],
                                    &PyArray_HolidaysConverter, &holidays,
                                    &NpyBusDayCalendar_Type, &busdaycal,
                                    &out_in)) {
        goto fail;
    }

    /* Make sure only one of the weekmask/holidays and busdaycal is supplied */
    if (busdaycal != NULL) {
        if (weekmask[0] != 2 || holidays.begin != NULL) {
            PyErr_SetString(PyExc_ValueError,
                    "Cannot supply both the weekmask/holidays and the "
                    "busdaycal parameters to busday_offset()");
            goto fail;
        }

        /* Indicate that the holidays weren't allocated by us */
        allocated_holidays = 0;

        /* Copy the private normalized weekmask/holidays data */
        holidays = busdaycal->holidays;
        busdays_in_weekmask = busdaycal->busdays_in_weekmask;
        memcpy(weekmask, busdaycal->weekmask, 7);
    }
    else {
        /*
         * Fix up the weekmask from the uninitialized
         * signal value to a proper default.
         */
        if (weekmask[0] == 2) {
            weekmask[0] = 1;
        }

        /* Count the number of business days in a week */
        busdays_in_weekmask = 0;
        for (i = 0; i < 7; ++i) {
            busdays_in_weekmask += weekmask[i];
        }

        /* The holidays list must be normalized before using it */
        normalize_holidays_list(&holidays, weekmask);
    }

    /* Make 'dates' into an array */
    if (PyArray_Check(dates_in)) {
        dates = (PyArrayObject *)dates_in;
        Py_INCREF(dates);
    }
    else {
        PyArray_Descr *datetime_dtype;

        /* Use the datetime dtype with generic units so it fills it in */
        datetime_dtype = PyArray_DescrFromType(NPY_DATETIME);
        if (datetime_dtype == NULL) {
            goto fail;
        }

        /* This steals the datetime_dtype reference */
        dates = (PyArrayObject *)PyArray_FromAny(dates_in, datetime_dtype,
                                                0, 0, 0, dates_in);
        if (dates == NULL) {
            goto fail;
        }
    }

    /* Make 'offsets' into an array */
    offsets = (PyArrayObject *)PyArray_FromAny(offsets_in,
                            PyArray_DescrFromType(NPY_INT64),
                            0, 0, 0, offsets_in);
    if (offsets == NULL) {
        goto fail;
    }

    /* Make sure 'out' is an array if it's provided */
    if (out_in != NULL) {
        if (!PyArray_Check(out_in)) {
            PyErr_SetString(PyExc_ValueError,
                    "busday_offset: must provide a NumPy array for 'out'");
            goto fail;
        }
        out = (PyArrayObject *)out_in;
    }

    ret = business_day_offset(dates, offsets, out, roll,
                    weekmask, busdays_in_weekmask,
                    holidays.begin, holidays.end);

    Py_DECREF(dates);
    Py_DECREF(offsets);
    if (allocated_holidays && holidays.begin != NULL) {
        PyArray_free(holidays.begin);
    }

    return out == NULL ? PyArray_Return(ret) : (PyObject *)ret;

fail:
    Py_XDECREF(dates);
    Py_XDECREF(offsets);
    if (allocated_holidays && holidays.begin != NULL) {
        PyArray_free(holidays.begin);
    }

    return NULL;
}

/*
 * This is the 'busday_count' function exposed for calling
 * from Python.
 */
NPY_NO_EXPORT PyObject *
array_busday_count(PyObject *NPY_UNUSED(self),
                      PyObject *args, PyObject *kwds)
{
    char *kwlist[] = {"begindates", "enddates",
                      "weekmask", "holidays", "busdaycal", "out", NULL};

    PyObject *dates_begin_in = NULL, *dates_end_in = NULL, *out_in = NULL;

    PyArrayObject *dates_begin = NULL, *dates_end = NULL, *out = NULL, *ret;
    npy_bool weekmask[7] = {2, 1, 1, 1, 1, 0, 0};
    NpyBusDayCalendar *busdaycal = NULL;
    int i, busdays_in_weekmask;
    npy_holidayslist holidays = {NULL, NULL};
    int allocated_holidays = 1;

    if (!PyArg_ParseTupleAndKeywords(args, kwds,
                                    "OO|O&O&O!O:busday_count", kwlist,
                                    &dates_begin_in,
                                    &dates_end_in,
                                    &PyArray_WeekMaskConverter, &weekmask[0],
                                    &PyArray_HolidaysConverter, &holidays,
                                    &NpyBusDayCalendar_Type, &busdaycal,
                                    &out_in)) {
        goto fail;
    }

    /* Make sure only one of the weekmask/holidays and busdaycal is supplied */
    if (busdaycal != NULL) {
        if (weekmask[0] != 2 || holidays.begin != NULL) {
            PyErr_SetString(PyExc_ValueError,
                    "Cannot supply both the weekmask/holidays and the "
                    "busdaycal parameters to busday_count()");
            goto fail;
        }

        /* Indicate that the holidays weren't allocated by us */
        allocated_holidays = 0;

        /* Copy the private normalized weekmask/holidays data */
        holidays = busdaycal->holidays;
        busdays_in_weekmask = busdaycal->busdays_in_weekmask;
        memcpy(weekmask, busdaycal->weekmask, 7);
    }
    else {
        /*
         * Fix up the weekmask from the uninitialized
         * signal value to a proper default.
         */
        if (weekmask[0] == 2) {
            weekmask[0] = 1;
        }

        /* Count the number of business days in a week */
        busdays_in_weekmask = 0;
        for (i = 0; i < 7; ++i) {
            busdays_in_weekmask += weekmask[i];
        }

        /* The holidays list must be normalized before using it */
        normalize_holidays_list(&holidays, weekmask);
    }

    /* Make 'dates_begin' into an array */
    if (PyArray_Check(dates_begin_in)) {
        dates_begin = (PyArrayObject *)dates_begin_in;
        Py_INCREF(dates_begin);
    }
    else {
        PyArray_Descr *datetime_dtype;

        /* Use the datetime dtype with generic units so it fills it in */
        datetime_dtype = PyArray_DescrFromType(NPY_DATETIME);
        if (datetime_dtype == NULL) {
            goto fail;
        }

        /* This steals the datetime_dtype reference */
        dates_begin = (PyArrayObject *)PyArray_FromAny(dates_begin_in,
                                                datetime_dtype,
                                                0, 0, 0, dates_begin_in);
        if (dates_begin == NULL) {
            goto fail;
        }
    }

    /* Make 'dates_end' into an array */
    if (PyArray_Check(dates_end_in)) {
        dates_end = (PyArrayObject *)dates_end_in;
        Py_INCREF(dates_end);
    }
    else {
        PyArray_Descr *datetime_dtype;

        /* Use the datetime dtype with generic units so it fills it in */
        datetime_dtype = PyArray_DescrFromType(NPY_DATETIME);
        if (datetime_dtype == NULL) {
            goto fail;
        }

        /* This steals the datetime_dtype reference */
        dates_end = (PyArrayObject *)PyArray_FromAny(dates_end_in,
                                                datetime_dtype,
                                                0, 0, 0, dates_end_in);
        if (dates_end == NULL) {
            goto fail;
        }
    }

    /* Make sure 'out' is an array if it's provided */
    if (out_in != NULL) {
        if (!PyArray_Check(out_in)) {
            PyErr_SetString(PyExc_ValueError,
                    "busday_offset: must provide a NumPy array for 'out'");
            goto fail;
        }
        out = (PyArrayObject *)out_in;
    }

    ret = business_day_count(dates_begin, dates_end, out,
                    weekmask, busdays_in_weekmask,
                    holidays.begin, holidays.end);

    Py_DECREF(dates_begin);
    Py_DECREF(dates_end);
    if (allocated_holidays && holidays.begin != NULL) {
        PyArray_free(holidays.begin);
    }

    return out == NULL ? PyArray_Return(ret) : (PyObject *)ret;

fail:
    Py_XDECREF(dates_begin);
    Py_XDECREF(dates_end);
    if (allocated_holidays && holidays.begin != NULL) {
        PyArray_free(holidays.begin);
    }

    return NULL;
}

/*
 * This is the 'is_busday' function exposed for calling
 * from Python.
 */
NPY_NO_EXPORT PyObject *
array_is_busday(PyObject *NPY_UNUSED(self),
                      PyObject *args, PyObject *kwds)
{
    char *kwlist[] = {"dates",
                      "weekmask", "holidays", "busdaycal", "out", NULL};

    PyObject *dates_in = NULL, *out_in = NULL;

    PyArrayObject *dates = NULL,*out = NULL, *ret;
    npy_bool weekmask[7] = {2, 1, 1, 1, 1, 0, 0};
    NpyBusDayCalendar *busdaycal = NULL;
    int i, busdays_in_weekmask;
    npy_holidayslist holidays = {NULL, NULL};
    int allocated_holidays = 1;

    if (!PyArg_ParseTupleAndKeywords(args, kwds,
                                    "O|O&O&O!O:is_busday", kwlist,
                                    &dates_in,
                                    &PyArray_WeekMaskConverter, &weekmask[0],
                                    &PyArray_HolidaysConverter, &holidays,
                                    &NpyBusDayCalendar_Type, &busdaycal,
                                    &out_in)) {
        goto fail;
    }

    /* Make sure only one of the weekmask/holidays and busdaycal is supplied */
    if (busdaycal != NULL) {
        if (weekmask[0] != 2 || holidays.begin != NULL) {
            PyErr_SetString(PyExc_ValueError,
                    "Cannot supply both the weekmask/holidays and the "
                    "busdaycal parameters to is_busday()");
            goto fail;
        }

        /* Indicate that the holidays weren't allocated by us */
        allocated_holidays = 0;

        /* Copy the private normalized weekmask/holidays data */
        holidays = busdaycal->holidays;
        busdays_in_weekmask = busdaycal->busdays_in_weekmask;
        memcpy(weekmask, busdaycal->weekmask, 7);
    }
    else {
        /*
         * Fix up the weekmask from the uninitialized
         * signal value to a proper default.
         */
        if (weekmask[0] == 2) {
            weekmask[0] = 1;
        }

        /* Count the number of business days in a week */
        busdays_in_weekmask = 0;
        for (i = 0; i < 7; ++i) {
            busdays_in_weekmask += weekmask[i];
        }

        /* The holidays list must be normalized before using it */
        normalize_holidays_list(&holidays, weekmask);
    }

    /* Make 'dates' into an array */
    if (PyArray_Check(dates_in)) {
        dates = (PyArrayObject *)dates_in;
        Py_INCREF(dates);
    }
    else {
        PyArray_Descr *datetime_dtype;

        /* Use the datetime dtype with generic units so it fills it in */
        datetime_dtype = PyArray_DescrFromType(NPY_DATETIME);
        if (datetime_dtype == NULL) {
            goto fail;
        }

        /* This steals the datetime_dtype reference */
        dates = (PyArrayObject *)PyArray_FromAny(dates_in,
                                                datetime_dtype,
                                                0, 0, 0, dates_in);
        if (dates == NULL) {
            goto fail;
        }
    }

    /* Make sure 'out' is an array if it's provided */
    if (out_in != NULL) {
        if (!PyArray_Check(out_in)) {
            PyErr_SetString(PyExc_ValueError,
                    "busday_offset: must provide a NumPy array for 'out'");
            goto fail;
        }
        out = (PyArrayObject *)out_in;
    }

    ret = is_business_day(dates, out,
                    weekmask, busdays_in_weekmask,
                    holidays.begin, holidays.end);

    Py_DECREF(dates);
    if (allocated_holidays && holidays.begin != NULL) {
        PyArray_free(holidays.begin);
    }

    return out == NULL ? PyArray_Return(ret) : (PyObject *)ret;

fail:
    Py_XDECREF(dates);
    if (allocated_holidays && holidays.begin != NULL) {
        PyArray_free(holidays.begin);
    }

    return NULL;
}
