"""A module to handle data generation for deep neural networks.

It uses the tf.data.Dataset object to enable parallel computing of batches.
"""
import numpy as np
import openslide
import tensorflow as tf
from typing import Sequence, Callable, Iterator, Any, Tuple, Optional, Dict, Union
from ..util.types import Patch, NDByteImage


def slide_query(patch: Patch, patch_size: int) -> NDByteImage:
    """
    Query patch image in slide.

    Get patch image given position, level and dimensions.

    Args:
        patch: the patch to query.
        patch_size: size of side of the patch in pixels.

    Returns:
        Numpy array rgb image of the patch.

    """
    slide = openslide.OpenSlide(patch["slide"])
    pil_img = slide.read_region(
        (patch["x"], patch["y"]), patch["level"], (patch_size, patch_size)
    )
    return np.array(pil_img)[:, :, 0:3]


def fast_slide_query(
    slides: Dict[str, openslide.OpenSlide],
    patch: Patch,
    patch_size: int
) -> NDByteImage:
    """
    Query patch image in slide.

    Get patch image given the slide obj, the position, level and dimensions.

    Args:
        slide: the slide to request the patch.
        patch: the patch to query.
        patch_size: size of side of the patch in pixels.

    Returns:
        Numpy array rgb image of the patch.

    """
    slide = slides[patch["slide"]]
    pil_img = slide.read_region(
        (patch["x"], patch["y"]), patch["level"], (patch_size, patch_size)
    )
    return np.array(pil_img)[:, :, 0:3]


def generator_fn(
    patch_list: Sequence[Patch],
    label_list: Sequence[Any],
    patch_size: int,
    preproc: Callable
) -> Iterator[Tuple[Patch, Any]]:
    """
    Provide a generator for tf.data.Dataset.

    Create a scope with required arguments, but produce a arg-less iterator.

    Args:
        patch_list: patch list to query.
        label_list: label of patches.
        patch_size: size of the side of the patches in pixels.
        preproc: a preprocessing function for images.
    Returns:
        A generator of tuples (patch, label).

    """
    def generator():
        for patch, y in zip(patch_list, label_list):
            x = slide_query(patch, patch_size)
            yield preproc(x), y

    return generator


def get_tf_dataset(
    patch_list: Sequence[Patch],
    label_list: Any,
    preproc: Callable,
    batch_size: int,
    patch_size: int,
    prefetch: Optional[int] = None,
    training: Optional[bool] = True,
) -> tf.data.Dataset:
    """
    Create tensorflow dataset.

    Create tf.data.Dataset object able to prefetch and batch samples from generator.

    Args:
        patch_list: patch list to query.
        label_list: label of patches.
        preproc: a preprocessing function for images.
        batch_size: number of samples per batch.
        patch_size: size (pixel) of the side of a square patch.

    Returns:
        tf.data.Dataset: a proper tensorflow dataset to fit on.

    """
    gen = generator_fn(patch_list, label_list, patch_size, preproc)
    try:
        shape_label = label_list[0].shape
    except AttributeError:
        shape_label = None
    dataset = tf.data.Dataset.from_generator(
        generator=gen,
        output_types=(np.float32, np.int32),
        output_shapes=((patch_size, patch_size, 3), shape_label),
    )
    if training:
        dataset = dataset.batch(batch_size, drop_remainder=True)
        dataset = dataset.repeat()
    else:
        dataset = dataset.batch(batch_size, drop_remainder=False)
    # prefetch
    # <=> while fitting batch b, prepare b+k in parallel
    if prefetch is None:
        dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
    else:
        dataset = dataset.prefetch(prefetch)
    return dataset