# cython: language_level=3
# cython: boundscheck=False
# cython: wraparound=False
# cython: initializedcheck=False
# cython: nonecheck=False
# cython: cdivision=True

from libc.stdlib cimport malloc, free
from libc.string cimport memset, memcpy
from cython.parallel cimport prange
cimport cython
import numpy as np
cimport numpy as cnp

cnp.import_array()

@cython.boundscheck(False)
@cython.wraparound(False)
@cython.cdivision(True)
cdef void _rosa_row(const cnp.int64_t* x, int n, int alphabet, cnp.int64_t* y_out) noexcept nogil:
    cdef int s = 2 * n + 1
    cdef int *trans = <int*> malloc(<size_t>s * alphabet * sizeof(int))
    cdef int *c = <int*> malloc(<size_t>s * sizeof(int))
    cdef int *d = <int*> malloc(<size_t>s * sizeof(int))
    cdef cnp.int64_t *e = <cnp.int64_t*> malloc(<size_t>s * sizeof(cnp.int64_t))

    memset(trans, 0xFF, <size_t>s * alphabet * sizeof(int))  # all -1

    cdef int g = 0, z = 1, i, t, r, p, q, u, v
    cdef cnp.int64_t a

    d[0] = 0
    c[0] = -1
    e[0] = -1

    for i in range(n):
        t = <int> x[i]
        r = z
        z += 1
        d[r] = d[g] + 1
        e[r] = -1
        p = g
        while p != -1 and trans[p * alphabet + t] == -1:
            trans[p * alphabet + t] = r
            p = c[p]
        if p == -1:
            c[r] = 0
        else:
            q = trans[p * alphabet + t]
            if d[p] + 1 == d[q]:
                c[r] = q
            else:
                u = z
                z += 1
                memcpy(trans + <size_t>u * alphabet, trans + <size_t>q * alphabet,
                       <size_t>alphabet * sizeof(int))
                d[u] = d[p] + 1
                c[u] = c[q]
                e[u] = e[q]
                while p != -1 and trans[p * alphabet + t] == q:
                    trans[p * alphabet + t] = u
                    p = c[p]
                c[q] = u
                c[r] = u
        v = g = r
        a = -1
        while v != -1:
            if d[v] > 0 and e[v] >= 0:
                a = x[e[v] + 1]
                break
            v = c[v]
        y_out[i] = a
        v = g
        while v != -1 and e[v] < i:
            e[v] = i
            v = c[v]

    free(trans)
    free(c)
    free(d)
    free(e)


def rosa_batch(cnp.int64_t[:, :] x not None, int alphabet):
    """
    x: (num_rows, n) int64, values in [0, alphabet)
    returns: (num_rows, n) int64 numpy array, -1 where no next-distinct-symbol exists
    """
    cdef Py_ssize_t num_rows = x.shape[0]
    cdef Py_ssize_t n = x.shape[1]
    y_np = np.empty((num_rows, n), dtype=np.int64)
    cdef cnp.int64_t[:, :] y = y_np
    cdef Py_ssize_t i

    for i in prange(num_rows, nogil=True, schedule='static'):
        _rosa_row(&x[i, 0], <int>n, alphabet, &y[i, 0])

    return y_np


cpdef list rosa(list x):
    cdef int n = len(x)
    cdef list y = [-1] * n
    cdef int s = 2 * n + 1
    cdef list b = [None] * s
    cdef list c = [-1] * s
    cdef list d = [0] * s
    cdef list e = [-1] * s
    b[0] = {}
    cdef int g = 0
    cdef int z = 1

    cdef int i, t
    cdef int r, p, q, u, v, a
    for i in range(n):
        t = x[i]
        r = z
        z += 1
        b[r] = {}
        d[r] = d[g] + 1
        p = g
        while p != -1 and t not in b[p]:
            b[p][t] = r
            p = c[p]
        if p == -1:
            c[r] = 0
        else:
            q = b[p][t]
            if d[p] + 1 == d[q]:
                c[r] = q
            else:
                u = z
                z += 1
                b[u] = b[q].copy()
                d[u] = d[p] + 1
                c[u] = c[q]
                e[u] = e[q]
                while p != -1 and b[p][t] == q:
                    b[p][t] = u
                    p = c[p]
                c[q] = c[r] = u
        v = g = r
        a = -1
        while v != -1:
            if d[v] > 0 and e[v] >= 0:
                a = x[e[v] + 1]
                break
            v = c[v]
        y[i] = a
        v = g
        while v != -1 and e[v] < i:
            e[v] = i
            v = c[v]
    return y