import numpy as np
cimport numpy as np
cimport cython

@cython.boundscheck(False)  # Deactivate bounds checking
@cython.wraparound(False)   # Deactivate negative indexing.
cpdef np.ndarray[np.double_t, ndim=1] \
    shift(
        np.ndarray[np.double_t, ndim=1] G, # Sorted groupby array
        np.ndarray[np.double_t, ndim=1] D, # Data array
        int N # Periods to shift
    ):

    # Definitions
    cdef int i
    cdef int i_prev
    cdef int l = D.shape[0]
    cdef np.ndarray[np.double_t, ndim=1] R = np.empty(l, dtype=np.double)

    for i in range(l):
        i_prev = i - N
        if (i_prev < 0) or (G[i] != G[i_prev]):
            R[i] = np.nan
        else:
            R[i] = D[i_prev]

    return R