import random as _random
import os as _os
import logging as _logging

import tensorflow as _tf
import matplotlib.pyplot as _plt
import so_ml_tools as _soml


def load_image_as_tensor(filename: str, img_shape=(224, 224), scale=True) -> _tf.Tensor:
    """
    Reads in an image from filename, turns it into a tensor and reshapes into
    specified shape (img_shape, img_shape, channels)

    Args:
        filename: path to target image
        img_shape: A `tuple` with the dimensions of the image.
        scale: A `bool` indicating if scaling needs to be performed (0-255 -> 0-1)

    Returns:
        A 'tf.Tensor' of shape (img_shape, 3) with of dtype 'tf.int32'.
    """

    img = _tf.io.read_file(filename=filename)
    return image_as_tensor(image=img, img_shape=img_shape, scale=scale)


def image_as_tensor(image, img_shape: tuple = (224, 224), scale: bool = True) -> _tf.Tensor:
    """
    Reads in an image from filename, turns it into a tensor and reshapes into
    specified shape (img_shape, img_shape, channels)
            
    Args:
        image: A `Tensor` of type `string`. 0-D. The encoded image bytes.
        img_shape: A `tuple` with the dimensions of the image.
        scale: A `bool` indicating if scaling needs to be performed (0-255 -> 0-1)

    Returns:
        A 'tf.Tensor' of shape (img_shape, 3) with of dtype 'tf.int32'.
    """
    # Decode image into tensor
    img = _tf.io.decode_image(contents=image, channels=3)

    # # Resize the image (height / width)
    img = _tf.image.convert_image_dtype(img, _tf.float32)
    img = _tf.image.resize(images=img, size=[img_shape[0], img_shape[1]])

    if not scale:
        return _tf.cast(x=img * 255, dtype=_tf.int32)
    else:
        return img


def is_image_float32_and_not_normalized(x) -> bool:
    """
    Convenience method to check if an image in `tf.Tensor` format is of type `float32` but not normalized (values between
    0..255 instead of 0..1)

    Args:
        x: A `tf.Tensor`

    Returns:
        A `bool` indicating if the tensor is of type `float32` but not normalized.
    """
    return x.dtype == _tf.float32 and _tf.math.reduce_max(x).numpy() > 1.0


def show_images_from_nparray_or_tensor(x, y, class_names: list[str] = None, indices=None, shape: tuple = (4, 6),
                                       cmap: str = 'gray') -> None:
    """
    Shows images stored in a tensor / numpy array. The array should be a vector of images.
    
    Args:
        x: a `tf.Tensor` containg the images
        y: a 'tf.Tensor' containing the labels associated with `x`
        class_names: a `list` of class names 
        indices: a `list` of indices for the images to display, `None` to pick random images
        shape: A `tuple` specifying the number of images to display (width, height)
        cmap: A 'str' with the color map to use for displaying the images.

    Returns:

    """
    y = _soml.util.label.to_prediction(y_prob=y)

    if is_image_float32_and_not_normalized(x):
        x = _tf.cast(x=x, dtype=_tf.uint8)

    if indices:
        assert shape[0] * shape[1] <= len(
            indices), f"Size of shape ({shape[0]}, {shape[1]}), with a total of " \
                      f"{shape[0] * shape[1]} images, is larger then number of indices supplied ({len(indices)})."
        for i in indices:
            if i > len(x):
                assert False, f"Values of indices point to an index ({i}) which is out of bounds of X (length: {len(x)})"

    fig = _plt.figure(figsize=(shape[1] * 3, shape[0] * 3))
    fig.patch.set_facecolor('gray')
    for i in range(shape[0] * shape[1]):
        ax = _plt.subplot(shape[0], shape[1], i + 1)
        ax.axis('off')

        if indices is None:
            rand_index = _random.choice(range(len(x)))
        else:
            rand_index = indices[i]

        _plt.imshow(x[rand_index], cmap=cmap)

        if _soml.util.label.is_multiclass_classification(y):
            class_index = _soml.util.label.probability_to_class(y)
        else:
            # Integer encoded labels
            class_index = y[rand_index]

        if class_names is None:
            _plt.title(class_index, color='white')
        else:
            _plt.title(f", {class_index}, {x[rand_index].shape}", color='white')
    _plt.show()


def show_single_image_from_nparray_or_tensor(image, title="", figsize=(10, 8), cmap='gray') -> None:
    """
    Shows images stored in a tensor / numpy array. The array should be a vector of images.

    :param image: the image to plot
    :param title: the title to display above the image
    :param figsize: Size of output figure (default=(10, 8)).
    :param cmap: is the collor map to use, use "gray" for gray scale images, use None for default.
    """
    fig = _plt.figure(figsize=figsize)
    fig.patch.set_facecolor('gray')
    _plt.imshow(image, cmap=cmap)
    _plt.axis('off')
    _plt.title(f"{title} {image.shape}", color='white')
    _plt.show()


def show_random_image_from_disk(target_dir, target_class, shape=(4, 6), cmap='gray') -> None:
    """
    Shows a random image from the file system.
    :param target_dir: The target directory, for example /food101/train
    :param target_class: The target class name, for example /steak
    :param shape: is the number of images in a grid to display
    :param cmap: is the color map to use, use "gray" for gray scale images, use None for default.
    """
    try:
        target_folder = _os.path.join(target_dir, target_class)

        # Get a random image path
        fig = _plt.figure(figsize=(shape[1] * 3, shape[0] * 3))
        fig.patch.set_facecolor('gray')
        for i in range(0, shape[0] * shape[1]):
            ax = _plt.subplot(shape[0], shape[1], i + 1)
            ax.axis('off')
            image_filename = _random.sample(_os.listdir(target_folder), 1)
            image_fullpath = _os.path.join(target_folder, image_filename[0])
            img = _plt.imread(image_fullpath)
            _plt.imshow(img, cmap=cmap)
            _plt.title(f"{image_filename[0]}: {img.shape}", color='white')

        _plt.show()
    except Exception as e:
        _logging.error(e)


def show_images_from_dataset(dataset: _tf.data.Dataset, shape=(4, 8)) -> None:
    assert isinstance(dataset, _tf.data.Dataset), f"The dataset supplied is not a tensorflow.data.Dataset."

    # Retrieve first batch, depending on the initalization of the dataset the batch size is default 32
    # so when performing a take of (1) we retreive the first batch
    batches = dataset.take(1)

    # Use an iterator to get the first batch of images and labels
    batch_iter = iter(batches)
    x, y_true = batch_iter.next()

    assert shape[0] * shape[1] == len(x), f"Size of shape ({shape[0]}, {shape[1]}), with a total of " \
                                          f"{shape[0] * shape[1]} images, is not equal to the batch size of the dataset ({len(x)})."

    if len(x) != shape[0] * shape[1]:
        raise TypeError('dataset is not a Dataset')

    show_images_from_nparray_or_tensor(x=x, y=y_true, class_names=dataset.class_names, shape=shape)
    _plt.show()
