# 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 apfio.vision.binomial import filter_2d_nbinomial

def filter_dark(image, k=16, t=10, d0=5, d1=20):
    """
    Smoothes and tone-maps dark regions of an image.

    :param image: An integer :py:class:`numpy:numpy.ndarray` of shape :code:`(h, w, d)` that is to be filtered.

    :param k: The size of the smoothing kernel to use. :code:`k = 1` disables smoothing. Larger `k` means stronger smoothing.

    :param t: The pixel brightness (average of all color channels) under which smoothing is to occur. Pixels brighter
              than this will not be modified by smoothing. The darker a pixel, the stronger the effect of smoothing.

    :param d0: Pixels with a brightness lower than this after smoothing this will be set to zero.

    :param d1: Pixels with a brightness lower than this but higher than d0 after smoothing will be adjusted by a linear
               transformation such that the darkest pixels with be set to zero and the brightest pixels will be left
               unmodified. Pixels brighter than d1 will not be modified at all.

    :return: The filtered image.
    """

    dtype = image.dtype

    if dtype == numpy.uint8:
        num_bits = 8
        double = numpy.uint16
    elif dtype == numpy.uint16:
        num_bits = 16
        double = numpy.uint32
    else:
        raise NotImplementedError(f"filter_dark cannot handle dtype {dtype}!")

    m = 2 ** num_bits - 1

    h, w, d = image.shape

    smooth = filter_2d_nbinomial(k, image)
    smooth = smooth.astype(double)

    alpha = smooth.sum(axis=-1, keepdims=True) // d
    alpha = numpy.where(alpha < t, (m * alpha) // t, m)

    interpol = alpha * image.astype(double) + (m - alpha) * smooth
    interpol = (interpol // m).astype(dtype)

    toned = numpy.where(interpol < d1, ((d1 - 0) * (interpol.clip(d0) - d0)) // (d1 - d0), interpol)

    return toned

