# coding: utf8
"""Useful functions for images."""
import numpy
from skimage.io import imread
from skimage.transform import resize
from .paths import imfiles_in_folder
from .types import NDBoolMask, PathLike, NDImage, NDByteImage, Coord
from ..patches.compat import convert_coords
import itertools
from typing import Iterator, List, Tuple, Sequence, Optional, Union, Any
from nptyping import NDArray, Shape, Float


def regular_grid(shape: Coord, interval: Coord, psize: Coord) -> Iterator[Coord]:
    """
    Get a regular grid of position on a slide given its dimensions.

    Arguments:
        shape: (x, y) shape of the window to tile.
        interval: (x, y) steps between patch samples.
        psize: (w, h) size of the patches (in pixels).

    Yields:
        (x, y) positions on a regular grid.

    """
    psize = convert_coords(psize)
    interval = convert_coords(interval)
    shape = convert_coords(shape)
    step = interval + psize
    maxj, maxi = step * ((shape - psize) / step + 1)
    col = numpy.arange(start=0, stop=maxj, step=step[0], dtype=int)
    line = numpy.arange(start=0, stop=maxi, step=step[1], dtype=int)
    for i, j in itertools.product(line, col):
        yield Coord(x=j, y=i)


def get_coords_from_mask(
    mask: NDBoolMask, shape: Coord, interval: Coord, psize: Coord
) -> Iterator[Coord]:
    """
    Get tissue coordinates given a tissue binary mask and slide dimensions.

    Arguments:
        mask: binary mask where tissue is marked as True.
        shape: (x, y) shape of the window to tile.
        interval: (x, y) steps between patch samples.
        psize: (w, h) size of the patches (in pixels).

    Yields:
        (x, y) positions on a regular grid.
    """

    psize = convert_coords(psize)
    interval = convert_coords(interval)
    shape = convert_coords(shape)
    step = interval + psize
    mask_w, mask_h = (shape - psize) / step + 1
    mask = resize(mask, (mask_h, mask_w))
    for i, j in numpy.argwhere(mask):
        yield step * (j, i)


def unlabeled_regular_grid_list(shape: Coord, step: int, psize: int) -> List[Coord]:
    """
    Get a regular grid of position on a slide given its dimensions.

    Args:
        shape: shape (i, j) of the window to tile.
        step: steps in pixels between patch samples.
        psize: size of the side of the patch (in pixels).

    Returns:
        Positions (i, j) on the regular grid.

    """
    maxi = step * int((shape[0] - (psize - step)) / step) + 1
    maxj = step * int((shape[1] - (psize - step)) / step) + 1
    col = numpy.arange(start=0, stop=maxj, step=step, dtype=int)
    line = numpy.arange(start=0, stop=maxi, step=step, dtype=int)
    return list(itertools.product(line, col))


def images_in_folder(
    folder: PathLike,
    authorized: Sequence[str] = (".png", ".jpg", ".jpeg", ".tif", ".tiff"),
    forbiden: Sequence[str] = ("thumbnail",),
    randomize: bool = False,
    datalim: Optional[int] = None,
    paths: bool = False,
) -> Iterator[Union[NDByteImage, Tuple[str, NDByteImage]]]:
    """
    Get images in a given folder.

    Get all images as numpy arrays (selected by file extension).
    You can remove terms from the research.

    Args:
        folder: absolute path to an image directory.
        authorized: authorized image file extensions.
        forbiden: non-authorized words in file names.
        randomize: whether to randomize output list of files.
        datalim: maximum number of file to extract in folder.
        paths: whether to return absolute path with image data.

    Yields:
        Images as numpy arrays, optionally with path.

    """
    for imfile in imfiles_in_folder(folder, authorized, forbiden, randomize, datalim):
        if paths:
            yield imfile, imread(imfile)
        else:
            yield imread(imfile)


def sample_img(
    image: NDImage, psize: int, spl_per_image: int, mask: NDBoolMask = None
) -> List[NDArray[Shape["N"], Float]]:
    """
    Split image in patches.

    Args:
        image: numpy image to fit on.
        psize: size in pixels of the side of a patch.
        spl_per_image: maximum number of patches to extract in image.
        mask: optional boolean array, we sample in true pixels if provided.

    Returns:
        Patches in the image.

    """
    img = image.astype(float)
    spaceshape = (image.shape[0], image.shape[1])
    di, dj = spaceshape
    if mask is None:
        positions = unlabeled_regular_grid_list(spaceshape, psize)
    else:
        half_size = int(0.5 * psize)
        cropped_mask = numpy.zeros_like(mask)
        cropped_mask[mask > 0] = 1
        cropped_mask[0 : half_size + 1, :] = 0
        cropped_mask[di - half_size - 1 : :, :] = 0
        cropped_mask[:, 0 : half_size + 1] = 0
        cropped_mask[:, dj - half_size - 1 : :] = 0
        y, x = numpy.where(cropped_mask > 0)
        y -= half_size
        x -= half_size
        positions = [(i, j) for i, j in zip(y, x)]

    numpy.random.shuffle(positions)
    positions = positions[0:spl_per_image]
    patches = [img[i : i + psize, j : j + psize].reshape(-1) for i, j in positions]
    return patches


def sample_img_sep_channels(
    image: NDByteImage, psize: int, spl_per_image: int, mask: NDBoolMask = None
) -> Tuple[List[NDArray[Shape["N"], Float]], ...]:
    """Fit vocabulary on a single image.

    Split image in patches and fit on them.

    Args:
        image: numpy image to fit on.
        psize: size in pixels of the side of a patch.
        spl_per_image: maximum number of patches to extract in image.
        mask: optional boolean array, we sample in true pixels if provided.

    Returns:
        Patches in the image in separated channels.

    """
    img = image.astype(float)
    n_channels = image.shape[-1]
    spaceshape = (image.shape[0], image.shape[1])
    di, dj = spaceshape
    if mask is None:
        positions = unlabeled_regular_grid_list(spaceshape, psize)
    else:
        half_size = int(0.5 * psize)
        cropped_mask = numpy.zeros_like(mask)
        cropped_mask[mask > 0] = 1
        cropped_mask[0 : half_size + 1, :] = 0
        cropped_mask[di - half_size - 1 : :, :] = 0
        cropped_mask[:, 0 : half_size + 1] = 0
        cropped_mask[:, dj - half_size - 1 : :] = 0
        y, x = numpy.where(cropped_mask > 0)
        y -= half_size
        x -= half_size
        positions = [(i, j) for i, j in zip(y, x)]
    numpy.random.shuffle(positions)
    if len(positions) > spl_per_image:
        positions = positions[0:spl_per_image]

    patches = []
    for c in range(n_channels):
        patches.append(
            [
                img[:, :, c][i : i + psize, j : j + psize].reshape(-1)
                for i, j in positions
            ]
        )
    return tuple(patches)