import numpy as np


class BitDepthConverter:
    """
    Convert bit depth of numpy arrays that represent grayscale images.
    """
    @staticmethod
    def create_LUT(levels, source, output):
        """
        Create a LUT for converting from one bit depth to another.
        :param levels:      min and max levels with which to create the LUT
        :type levels:       tuple
        :param source:      bit depth of the source. 16, 32, or 64.
        :type source:       int
        :param output:      desired output bit depth
        :type output:       int
        :return:            8 bit LUT to convert the source's bit depth to 8 bit
        :rtype:             np.ndarray
        """

        accepted_srcs = [16, 32, 64]
        if source not in accepted_srcs:
            raise TypeError('Can only convert from uint16, uint32, or uint64')

        accepted_outs = [8, 16, 32]
        if output not in accepted_outs:
            raise TypeError('Can only output uint8, uint16, or uint32')

        typecode = 'uint' + str(source)
        LUT = np.arange(2**source, dtype=typecode)
        LUT.clip(levels[0], levels[1], out=LUT)
        LUT -= levels[0]
        np.floor_divide(LUT, (levels[1] - levels[0] + 1) / (2**output), out=LUT, casting='unsafe')

        typecode = 'uint' + str(output)
        return LUT.astype(typecode)

    @staticmethod
    def apply_LUT(image, LUT):
        """
        Apply LUT
        :param image:   The image upon which to apply the LUT and change its bit depth
        :type image:    np.ndarray
        :param LUT:     The LUT to use for downscaling the bit depth. Generated by BitDepthConvert.create_LUT
        :type LUT:      np.ndarray
        :return:        Downscaled image with the LUT applied to it
        :rtype:         np.ndarray
        """
        return np.take(LUT, image).astype(np.uint8)
