import numpy as np
import cv2
cimport numpy as np
cimport cython
cimport libcpp
cimport libcpp.pair
cimport libcpp.queue
from libcpp.pair cimport *
from libcpp.queue  cimport *

@cython.boundscheck(False)
@cython.wraparound(False)
cdef np.ndarray[np.int32_t, ndim=2] _pse(np.ndarray[np.uint8_t, ndim=3] kernels,
                                         np.ndarray[np.int32_t, ndim=2] label,
                                         int kernel_num,
                                         int label_num,
                                         float min_area=0):
    cdef np.ndarray[np.int32_t, ndim=2] pred
    pred = np.zeros((label.shape[0], label.shape[1]), dtype=np.int32)

    for label_idx in range(1, label_num):
        if np.sum(label == label_idx) < min_area:
            label[label == label_idx] = 0

    cdef libcpp.queue.queue[libcpp.pair.pair[np.int16_t,np.int16_t]] que = \
        queue[libcpp.pair.pair[np.int16_t,np.int16_t]]()
    cdef libcpp.queue.queue[libcpp.pair.pair[np.int16_t,np.int16_t]] nxt_que = \
        queue[libcpp.pair.pair[np.int16_t,np.int16_t]]()
    cdef np.int16_t* dx = [-1, 1, 0, 0]
    cdef np.int16_t* dy = [0, 0, -1, 1]
    cdef np.int16_t tmpx, tmpy

    points = np.array(np.where(label > 0)).transpose((1, 0))
    for point_idx in range(points.shape[0]):
        tmpx, tmpy = points[point_idx, 0], points[point_idx, 1]
        que.push(pair[np.int16_t,np.int16_t](tmpx, tmpy))
        pred[tmpx, tmpy] = label[tmpx, tmpy]

    cdef libcpp.pair.pair[np.int16_t,np.int16_t] cur
    cdef int cur_label
    for kernel_idx in range(kernel_num - 1, -1, -1):
        while not que.empty():
            cur = que.front()
            que.pop()
            cur_label = pred[cur.first, cur.second]

            is_edge = True
            for j in range(4):
                tmpx = cur.first + dx[j]
                tmpy = cur.second + dy[j]
                if tmpx < 0 or tmpx >= label.shape[0] or tmpy < 0 or tmpy >= label.shape[1]:
                    continue
                if kernels[kernel_idx, tmpx, tmpy] == 0 or pred[tmpx, tmpy] > 0:
                    continue

                que.push(pair[np.int16_t,np.int16_t](tmpx, tmpy))
                pred[tmpx, tmpy] = cur_label
                is_edge = False
            if is_edge:
                nxt_que.push(cur)

        que, nxt_que = nxt_que, que

    return pred

def pse(kernels, min_area):
    kernel_num = kernels.shape[0]
    label_num, label = cv2.connectedComponents(kernels[-1], connectivity=4)
    return _pse(kernels[:-1], label, kernel_num, label_num, min_area)