fields.pyx

"""
Functions for accessing attributes of Timestamp/datetime64/datetime-like
objects and arrays
"""
from locale import LC_TIME

from _strptime import LocaleTime

cimport cython
from cython cimport Py_ssize_t

import numpy as np

cimport numpy as cnp
from numpy cimport (
    int8_t,
    int32_t,
    int64_t,
    ndarray,
    uint32_t,
)

cnp.import_array()

from pandas._config.localization import set_locale

from pandas._libs.tslibs.ccalendar import (
    DAYS_FULL,
    MONTHS_FULL,
)

from pandas._libs.tslibs.ccalendar cimport (
    dayofweek,
    get_day_of_year,
    get_days_in_month,
    get_firstbday,
    get_iso_calendar,
    get_lastbday,
    get_week_of_year,
    iso_calendar_t,
)
from pandas._libs.tslibs.nattype cimport NPY_NAT
from pandas._libs.tslibs.np_datetime cimport (
    NPY_DATETIMEUNIT,
    NPY_FR_ns,
    npy_datetimestruct,
    pandas_datetime_to_datetimestruct,
    pandas_timedelta_to_timedeltastruct,
    pandas_timedeltastruct,
)


@cython.wraparound(False)
@cython.boundscheck(False)
def build_field_sarray(const int64_t[:] dtindex, NPY_DATETIMEUNIT reso):
    """
    Datetime as int64 representation to a structured array of fields
    """
    cdef:
        Py_ssize_t i, count = len(dtindex)
        npy_datetimestruct dts
        ndarray[int32_t] years, months, days, hours, minutes, seconds, mus

    sa_dtype = [
        ("Y", "i4"),  # year
        ("M", "i4"),  # month
        ("D", "i4"),  # day
        ("h", "i4"),  # hour
        ("m", "i4"),  # min
        ("s", "i4"),  # second
        ("u", "i4"),  # microsecond
    ]

    out = np.empty(count, dtype=sa_dtype)

    years = out["Y"]
    months = out["M"]
    days = out["D"]
    hours = out["h"]
    minutes = out["m"]
    seconds = out["s"]
    mus = out["u"]

    for i in range(count):
        pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
        years[i] = dts.year
        months[i] = dts.month
        days[i] = dts.day
        hours[i] = dts.hour
        minutes[i] = dts.min
        seconds[i] = dts.sec
        mus[i] = dts.us

    return out


def month_position_check(fields, weekdays) -> str | None:
    cdef:
        int32_t daysinmonth, y, m, d
        bint calendar_end = True
        bint business_end = True
        bint calendar_start = True
        bint business_start = True
        bint cal
        int32_t[:] years = fields["Y"]
        int32_t[:] months = fields["M"]
        int32_t[:] days = fields["D"]

    for y, m, d, wd in zip(years, months, days, weekdays):
        if calendar_start:
            calendar_start &= d == 1
        if business_start:
            business_start &= d == 1 or (d <= 3 and wd == 0)

        if calendar_end or business_end:
            daysinmonth = get_days_in_month(y, m)
            cal = d == daysinmonth
            if calendar_end:
                calendar_end &= cal
            if business_end:
                business_end &= cal or (daysinmonth - d < 3 and wd == 4)
        elif not calendar_start and not business_start:
            break

    if calendar_end:
        return "ce"
    elif business_end:
        return "be"
    elif calendar_start:
        return "cs"
    elif business_start:
        return "bs"
    else:
        return None


@cython.wraparound(False)
@cython.boundscheck(False)
def get_date_name_field(
    const int64_t[:] dtindex,
    str field,
    object locale=None,
    NPY_DATETIMEUNIT reso=NPY_FR_ns,
):
    """
    Given a int64-based datetime index, return array of strings of date
    name based on requested field (e.g. day_name)
    """
    cdef:
        Py_ssize_t i, count = dtindex.shape[0]
        ndarray[object] out, names
        npy_datetimestruct dts
        int dow

    out = np.empty(count, dtype=object)

    if field == "day_name":
        if locale is None:
            names = np.array(DAYS_FULL, dtype=np.object_)
        else:
            names = np.array(_get_locale_names("f_weekday", locale),
                             dtype=np.object_)
        for i in range(count):
            if dtindex[i] == NPY_NAT:
                out[i] = np.nan
                continue

            pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
            dow = dayofweek(dts.year, dts.month, dts.day)
            out[i] = names[dow].capitalize()

    elif field == "month_name":
        if locale is None:
            names = np.array(MONTHS_FULL, dtype=np.object_)
        else:
            names = np.array(_get_locale_names("f_month", locale),
                             dtype=np.object_)
        for i in range(count):
            if dtindex[i] == NPY_NAT:
                out[i] = np.nan
                continue

            pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
            out[i] = names[dts.month].capitalize()

    else:
        raise ValueError(f"Field {field} not supported")

    return out


cdef bint _is_on_month(int month, int compare_month, int modby) nogil:
    """
    Analogous to DateOffset.is_on_offset checking for the month part of a date.
    """
    if modby == 1:
        return True
    elif modby == 3:
        return (month - compare_month) % 3 == 0
    else:
        return month == compare_month


@cython.wraparound(False)
@cython.boundscheck(False)
def get_start_end_field(
    const int64_t[:] dtindex,
    str field,
    str freqstr=None,
    int month_kw=12,
    NPY_DATETIMEUNIT reso=NPY_FR_ns,
):
    """
    Given an int64-based datetime index return array of indicators
    of whether timestamps are at the start/end of the month/quarter/year
    (defined by frequency).

    Parameters
    ----------
    dtindex : ndarray[int64]
    field : str
    frestr : str or None, default None
    month_kw : int, default 12
    reso : NPY_DATETIMEUNIT, default NPY_FR_ns

    Returns
    -------
    ndarray[bool]
    """
    cdef:
        Py_ssize_t i
        int count = dtindex.shape[0]
        bint is_business = 0
        int end_month = 12
        int start_month = 1
        ndarray[int8_t] out
        npy_datetimestruct dts
        int compare_month, modby

    out = np.zeros(count, dtype="int8")

    if freqstr:
        if freqstr == "C":
            raise ValueError(f"Custom business days is not supported by {field}")
        is_business = freqstr[0] == "B"

        # YearBegin(), BYearBegin() use month = starting month of year.
        # QuarterBegin(), BQuarterBegin() use startingMonth = starting
        # month of year. Other offsets use month, startingMonth as ending
        # month of year.

        if (freqstr[0:2] in ["MS", "QS", "AS"]) or (
                freqstr[1:3] in ["MS", "QS", "AS"]):
            end_month = 12 if month_kw == 1 else month_kw - 1
            start_month = month_kw
        else:
            end_month = month_kw
            start_month = (end_month % 12) + 1
    else:
        end_month = 12
        start_month = 1

    compare_month = start_month if "start" in field else end_month
    if "month" in field:
        modby = 1
    elif "quarter" in field:
        modby = 3
    else:
        modby = 12

    if field in ["is_month_start", "is_quarter_start", "is_year_start"]:
        if is_business:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = 0
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)

                if _is_on_month(dts.month, compare_month, modby) and (
                        dts.day == get_firstbday(dts.year, dts.month)):
                    out[i] = 1

        else:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = 0
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)

                if _is_on_month(dts.month, compare_month, modby) and dts.day == 1:
                    out[i] = 1

    elif field in ["is_month_end", "is_quarter_end", "is_year_end"]:
        if is_business:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = 0
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)

                if _is_on_month(dts.month, compare_month, modby) and (
                        dts.day == get_lastbday(dts.year, dts.month)):
                    out[i] = 1

        else:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = 0
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)

                if _is_on_month(dts.month, compare_month, modby) and (
                        dts.day == get_days_in_month(dts.year, dts.month)):
                    out[i] = 1

    else:
        raise ValueError(f"Field {field} not supported")

    return out.view(bool)


@cython.wraparound(False)
@cython.boundscheck(False)
def get_date_field(
    const int64_t[:] dtindex,
    str field,
    NPY_DATETIMEUNIT reso=NPY_FR_ns,
):
    """
    Given a int64-based datetime index, extract the year, month, etc.,
    field and return an array of these values.
    """
    cdef:
        Py_ssize_t i, count = len(dtindex)
        ndarray[int32_t] out
        npy_datetimestruct dts

    out = np.empty(count, dtype="i4")

    if field == "Y":
        with nogil:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = -1
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
                out[i] = dts.year
        return out

    elif field == "M":
        with nogil:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = -1
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
                out[i] = dts.month
        return out

    elif field == "D":
        with nogil:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = -1
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
                out[i] = dts.day
        return out

    elif field == "h":
        with nogil:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = -1
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
                out[i] = dts.hour
                # TODO: can we de-dup with period.pyx <accessor>s?
        return out

    elif field == "m":
        with nogil:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = -1
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
                out[i] = dts.min
        return out

    elif field == "s":
        with nogil:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = -1
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
                out[i] = dts.sec
        return out

    elif field == "us":
        with nogil:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = -1
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
                out[i] = dts.us
        return out

    elif field == "ns":
        with nogil:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = -1
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
                out[i] = dts.ps // 1000
        return out
    elif field == "doy":
        with nogil:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = -1
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
                out[i] = get_day_of_year(dts.year, dts.month, dts.day)
        return out

    elif field == "dow":
        with nogil:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = -1
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
                out[i] = dayofweek(dts.year, dts.month, dts.day)
        return out

    elif field == "woy":
        with nogil:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = -1
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
                out[i] = get_week_of_year(dts.year, dts.month, dts.day)
        return out

    elif field == "q":
        with nogil:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = -1
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
                out[i] = dts.month
                out[i] = ((out[i] - 1) // 3) + 1
        return out

    elif field == "dim":
        with nogil:
            for i in range(count):
                if dtindex[i] == NPY_NAT:
                    out[i] = -1
                    continue

                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
                out[i] = get_days_in_month(dts.year, dts.month)
        return out
    elif field == "is_leap_year":
        return isleapyear_arr(get_date_field(dtindex, "Y", reso=reso))

    raise ValueError(f"Field {field} not supported")


@cython.wraparound(False)
@cython.boundscheck(False)
def get_timedelta_field(
    const int64_t[:] tdindex,
    str field,
    NPY_DATETIMEUNIT reso=NPY_FR_ns,
):
    """
    Given a int64-based timedelta index, extract the days, hrs, sec.,
    field and return an array of these values.
    """
    cdef:
        Py_ssize_t i, count = len(tdindex)
        ndarray[int32_t] out
        pandas_timedeltastruct tds

    out = np.empty(count, dtype="i4")

    if field == "seconds":
        with nogil:
            for i in range(count):
                if tdindex[i] == NPY_NAT:
                    out[i] = -1
                    continue

                pandas_timedelta_to_timedeltastruct(tdindex[i], reso, &tds)
                out[i] = tds.seconds
        return out

    elif field == "microseconds":
        with nogil:
            for i in range(count):
                if tdindex[i] == NPY_NAT:
                    out[i] = -1
                    continue

                pandas_timedelta_to_timedeltastruct(tdindex[i], reso, &tds)
                out[i] = tds.microseconds
        return out

    elif field == "nanoseconds":
        with nogil:
            for i in range(count):
                if tdindex[i] == NPY_NAT:
                    out[i] = -1
                    continue

                pandas_timedelta_to_timedeltastruct(tdindex[i], reso, &tds)
                out[i] = tds.nanoseconds
        return out

    raise ValueError(f"Field {field} not supported")


@cython.wraparound(False)
@cython.boundscheck(False)
def get_timedelta_days(
    const int64_t[:] tdindex,
    NPY_DATETIMEUNIT reso=NPY_FR_ns,
):
    """
    Given a int64-based timedelta index, extract the days,
    field and return an array of these values.
    """
    cdef:
        Py_ssize_t i, count = len(tdindex)
        ndarray[int64_t] out
        pandas_timedeltastruct tds

    out = np.empty(count, dtype="i8")

    with nogil:
        for i in range(count):
            if tdindex[i] == NPY_NAT:
                out[i] = -1
                continue

            pandas_timedelta_to_timedeltastruct(tdindex[i], reso, &tds)
            out[i] = tds.days
    return out


cpdef isleapyear_arr(ndarray years):
    """vectorized version of isleapyear; NaT evaluates as False"""
    cdef:
        ndarray[int8_t] out

    out = np.zeros(len(years), dtype="int8")
    out[np.logical_or(years % 400 == 0,
                      np.logical_and(years % 4 == 0,
                                     years % 100 > 0))] = 1
    return out.view(bool)


@cython.wraparound(False)
@cython.boundscheck(False)
def build_isocalendar_sarray(const int64_t[:] dtindex, NPY_DATETIMEUNIT reso):
    """
    Given a int64-based datetime array, return the ISO 8601 year, week, and day
    as a structured array.
    """
    cdef:
        Py_ssize_t i, count = len(dtindex)
        npy_datetimestruct dts
        ndarray[uint32_t] iso_years, iso_weeks, days
        iso_calendar_t ret_val

    sa_dtype = [
        ("year", "u4"),
        ("week", "u4"),
        ("day", "u4"),
    ]

    out = np.empty(count, dtype=sa_dtype)

    iso_years = out["year"]
    iso_weeks = out["week"]
    days = out["day"]

    with nogil:
        for i in range(count):
            if dtindex[i] == NPY_NAT:
                ret_val = 0, 0, 0
            else:
                pandas_datetime_to_datetimestruct(dtindex[i], reso, &dts)
                ret_val = get_iso_calendar(dts.year, dts.month, dts.day)

            iso_years[i] = ret_val[0]
            iso_weeks[i] = ret_val[1]
            days[i] = ret_val[2]
    return out


def _get_locale_names(name_type: str, locale: object = None):
    """
    Returns an array of localized day or month names.

    Parameters
    ----------
    name_type : str
        Attribute of LocaleTime() in which to return localized names.
    locale : str

    Returns
    -------
    list of locale names
    """
    with set_locale(locale, LC_TIME):
        return getattr(LocaleTime(), name_type)


# ---------------------------------------------------------------------
# Rounding


class RoundTo:
    """
    enumeration defining the available rounding modes

    Attributes
    ----------
    MINUS_INFTY
        round towards -∞, or floor [2]_
    PLUS_INFTY
        round towards +∞, or ceil [3]_
    NEAREST_HALF_EVEN
        round to nearest, tie-break half to even [6]_
    NEAREST_HALF_MINUS_INFTY
        round to nearest, tie-break half to -∞ [5]_
    NEAREST_HALF_PLUS_INFTY
        round to nearest, tie-break half to +∞ [4]_


    References
    ----------
    .. [1] "Rounding - Wikipedia"
           https://en.wikipedia.org/wiki/Rounding
    .. [2] "Rounding down"
           https://en.wikipedia.org/wiki/Rounding#Rounding_down
    .. [3] "Rounding up"
           https://en.wikipedia.org/wiki/Rounding#Rounding_up
    .. [4] "Round half up"
           https://en.wikipedia.org/wiki/Rounding#Round_half_up
    .. [5] "Round half down"
           https://en.wikipedia.org/wiki/Rounding#Round_half_down
    .. [6] "Round half to even"
           https://en.wikipedia.org/wiki/Rounding#Round_half_to_even
    """
    @property
    def MINUS_INFTY(self) -> int:
        return 0

    @property
    def PLUS_INFTY(self) -> int:
        return 1

    @property
    def NEAREST_HALF_EVEN(self) -> int:
        return 2

    @property
    def NEAREST_HALF_PLUS_INFTY(self) -> int:
        return 3

    @property
    def NEAREST_HALF_MINUS_INFTY(self) -> int:
        return 4


cdef ndarray[int64_t] _floor_int64(const int64_t[:] values, int64_t unit):
    cdef:
        Py_ssize_t i, n = len(values)
        ndarray[int64_t] result = np.empty(n, dtype="i8")
        int64_t res, value

    with cython.overflowcheck(True):
        for i in range(n):
            value = values[i]
            if value == NPY_NAT:
                res = NPY_NAT
            else:
                res = value - value % unit
            result[i] = res

    return result


cdef ndarray[int64_t] _ceil_int64(const int64_t[:] values, int64_t unit):
    cdef:
        Py_ssize_t i, n = len(values)
        ndarray[int64_t] result = np.empty(n, dtype="i8")
        int64_t res, value

    with cython.overflowcheck(True):
        for i in range(n):
            value = values[i]

            if value == NPY_NAT:
                res = NPY_NAT
            else:
                remainder = value % unit
                if remainder == 0:
                    res = value
                else:
                    res = value + (unit - remainder)

            result[i] = res

    return result


cdef ndarray[int64_t] _rounddown_int64(values, int64_t unit):
    return _ceil_int64(values - unit // 2, unit)


cdef ndarray[int64_t] _roundup_int64(values, int64_t unit):
    return _floor_int64(values + unit // 2, unit)


def round_nsint64(values: np.ndarray, mode: RoundTo, nanos: int) -> np.ndarray:
    """
    Applies rounding mode at given frequency

    Parameters
    ----------
    values : np.ndarray[int64_t]`
    mode : instance of `RoundTo` enumeration
    nanos : np.int64
        Freq to round to, expressed in nanoseconds

    Returns
    -------
    np.ndarray[int64_t]
    """
    cdef:
        int64_t unit = nanos

    if mode == RoundTo.MINUS_INFTY:
        return _floor_int64(values, unit)
    elif mode == RoundTo.PLUS_INFTY:
        return _ceil_int64(values, unit)
    elif mode == RoundTo.NEAREST_HALF_MINUS_INFTY:
        return _rounddown_int64(values, unit)
    elif mode == RoundTo.NEAREST_HALF_PLUS_INFTY:
        return _roundup_int64(values, unit)
    elif mode == RoundTo.NEAREST_HALF_EVEN:
        # for odd unit there is no need of a tie break
        if unit % 2:
            return _rounddown_int64(values, unit)
        quotient, remainder = np.divmod(values, unit)
        mask = np.logical_or(
            remainder > (unit // 2),
            np.logical_and(remainder == (unit // 2), quotient % 2)
        )
        quotient[mask] += 1
        return quotient * unit

    # if/elif above should catch all rounding modes defined in enum 'RoundTo':
    # if flow of control arrives here, it is a bug
    raise ValueError("round_nsint64 called with an unrecognized rounding mode")