from ..core.exceptions import BadInput
from contracts.interface import describe_value


# TODO: make naming uniform
def check_2d_array(value, name=None):
    import numpy as np

    """ Checks that we have 2D numpy array """
    if not isinstance(value, np.ndarray):
        raise BadInput("Expected 2d array, got %s." % value.__class__.__name__, None, name)

    if len(value.shape) != 2:
        raise BadInput("Expected 2D array, got %s." % str(value.shape), None, name)


def assert_rgb_image(image, name=None):
    import numpy as np

    if not isinstance(image, np.ndarray):
        raise BadInput("Expected RGB image for %r, got %s." % (name, image.__class__.__name__), None, name)

    if image.dtype != "uint8":
        raise BadInput(
            "Expected RGB image for %r, got an array %s %s." % (name, str(image.shape), image.dtype),
            None,
            name,
        )

    if len(image.shape) != 3 or image.shape[2] != 3:
        raise BadInput("Bad shape for %s, expected RGB, got %s." % (name, str(image.shape)), None, name)


def assert_gray_image(image, name=None):
    import numpy as np

    if not isinstance(image, np.ndarray):
        raise BadInput(
            "Expected a grayscale image for %r, got %s." % (name, image.__class__.__name__), None, name
        )

    if image.dtype != "uint8":
        raise BadInput(
            "Expected a grayscale image for %r, got an array %s %s." % (name, str(image.shape), image.dtype),
            None,
            name,
        )

    if len(image.shape) != 2:
        raise BadInput("Bad shape for %r, expected grayscale, got %s." % (name, str(image.shape)), None, name)


def check_numpy_array(value):
    import numpy as np

    if not isinstance(value, np.ndarray):
        msg = "Expected array: %s" % describe_value(value)
        raise ValueError(msg)


def check_convertible_to_rgb(image):
    check_numpy_array(image)

    if image.ndim == 3:
        good = image.shape[2] in [3, 4]
        if not good:
            msg = "Wrong dimensions: %s" % image.shape
            raise ValueError(msg)

    if image.ndim > 3 or image.ndim <= 1:
        msg = "Bad shape: %s" % image.shape
        raise ValueError(msg)

    if image.dtype == "uint8":
        pass
    elif image.dtype == "float32":
        imin = image.min()
        imax = image.max()
        ok = imin >= 0 and imax <= 1
        if not ok:
            msg = "Invalid values range: %s %s" % (imin, imax)
            raise ValueError(msg)
    else:
        msg = "invalid dtype %r" % image.dtype
        raise ValueError(msg)


def input_check_numpy_array(block, signal):
    try:
        check_numpy_array(block.get_input(signal))
    except ValueError as e:
        raise BadInput(str(e), block, signal) from e


def input_check_convertible_to_rgb(block, signal):
    try:
        check_convertible_to_rgb(block.get_input(signal))
    except ValueError as e:
        raise BadInput(str(e), block, signal) from e


def input_check_rgb_or_grayscale(block, input):  # @ReservedAssignment
    """ Checks that the selected input is either a grayscale or RGB image.
        That is, a numpy array of uint8 either H x W,  H x W x 3,
        or HxWx4. 
        Raises BadInput if it is not. 
    """
    # TODO: write this better

    image = block.get_input(input)
    input_check_numpy_array(block, input)

    if image.dtype != "uint8":
        msg = "Expected an image, got an array %s %s." % (str(image.shape), image.dtype)
        raise BadInput(msg, block, input)

    shape = image.shape
    if len(shape) == 3:
        depth = shape[2]
        if not depth in [3, 4]:
            msg = "Bad shape for image: %s" % str(shape)
            raise BadInput(msg, block, input)
    elif len(shape) == 2:
        pass
    else:
        msg = "Bad shape for image: %s" % str(shape)
        raise BadInput(msg, block, input)

    if shape[0] == 0 or shape[1] == 0:
        msg = "Image with zero width or height: %s" % (str(shape))
        raise BadInput(msg, block, input)


def check_rgb(block, input):  # @ReservedAssignment
    """ Checks that the selected input is either  a RGB image.
        That is, a numpy array of uint8 of shape H x W x 3. 
        Raises BadInput if it is not. 
    """
    image = block.get_input(input)
    import numpy as np

    if not isinstance(image, np.ndarray):
        raise BadInput(
            "Expected RGB, this is not even a " "numpy array: %s" % image.__class__.__name__, block, input
        )
    if image.dtype != "uint8":
        raise BadInput(
            "Expected an image, got an array %s %s." % (str(image.shape), image.dtype), block, input
        )
    shape = image.shape
    if len(shape) == 3:
        if shape[2] != 3:
            raise BadInput("Bad shape for image: %s" % str(shape), block, input)
    else:
        raise BadInput("Bad shape for image: %s" % str(shape), block, input)
