# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.


import numpy
from numba import prange, njit, uint64, uint32, int64, uint16, uint8


@njit(inline="always")
def binomial_kernel(k, dtype):
    """
    Creates a 1-dimensional binomial kernel.
    The coefficients in this kernel approximate a Gaussian distribution.

    :param k: The length of the kernel.

    :param dtype: The :py:class:`numpy:numpy.dtype` of the kernel array.

    :return: A :py:class:`numpy:numpy.ndarray` of shape :code:`(k, )` and integer type.
    """
    if k < 0:
        raise ValueError("The given kernel size must not be negative!")
    kernel = numpy.zeros(k, dtype=dtype)
    kernel[0] = 1
    for i in range(k - 1):
        w = kernel[0]
        for j in range(k - 1):
            w_new = kernel[1 + j]
            kernel[1 + j] += w
            w = w_new
    return kernel


@njit([(uint64, uint8[:, :, ::1], uint8[:, :, ::1], uint8[:, :, ::1]),
       (uint64, uint16[:, :, ::1], uint16[:, :, ::1], uint16[:, :, ::1])],
      parallel=True, fastmath=True, nogil=True, cache=True)
def __filter_2d_nbinomial(k, signal, buffer, out):
    """
    Convolves a batch of 2D signals with a two-dimensional normalized binomial kernel.
    The signals are padded by zeros.

    :param k: The size of the binomial kernel to use.

    :param signal: An integer array of shape :code:`(n, h, w)`. It represents the batch of 2D signals to be filtered.

    :param buffer: An integer array of shape :code:`(n, h, w)`. This procedure may write undefined values to it.

    :return: The filtered version of the array.
    """

    h, w, d = signal.shape
    if h * w * d == 0:
        return

    # Note: This procedure is not particularly optimized, because k is so general.
    kernel = binomial_kernel(k, uint32)

    # Kernel horizontal, output horizontal:
    for y in prange(h):
        for x in prange(w):
            for c in range(d):
                r = uint32(0)
                for kidx in range(k):
                    index = int64(x) - int64(k // 2) + kidx
                    index_padded = max(0, min(w - 1, index))
                    r += kernel[kidx] * (index_padded == index) * signal[y, index_padded, c]
                buffer[y, x, c] = r // 2 ** (k - 1)

    # Kernel vertical, output horizontal:
    for y in prange(h):
        for x in prange(w):
            for c in range(d):
                r = uint32(0)
                for kidx in range(k):
                    index = int64(y) - int64(k // 2) + kidx
                    index_padded = max(0, min(h - 1, index))
                    r += kernel[kidx] * (index_padded == index) * buffer[index_padded, x, c]
                out[y, x, c] = r // 2 ** (k - 1)


def filter_2d_nbinomial(k, signal):
    """
    Convolves a batch of 2D signals with a two-dimensional normalized binomial kernel.
    The signals are padded by zeros.

    :param k: The size of the binomial kernel to use.

    :param signal: An integer array of shape :code:`(n, h, w)`. It represents the batch of 2D signals to be filtered.

    :return: The filtered version of the array.
    """

    if signal.size == 0:
        return signal.copy()
    buffer = numpy.empty_like(signal)
    out = numpy.empty_like(signal)
    __filter_2d_nbinomial(uint64(k), signal, buffer, out)
    return out

