import os

import cv2
import numpy as np
import tifffile as tiff


def load_data(data_path, img_height, img_width, img_channels, n_classes):
    """
    Custom data loader for loading of images and localization maps from given
    directory into NumPy arrays for later use in training and evaluation. First
    it tries to load pickled NumPy arrays of images and maps if they exist. If
    the arrays do not exist, images and maps of PNG and TIFF file formats are
    loaded and before beeing returned, they are pickled as NPY files in the
    directory.

    Parameters
    ----------
    data_path : str
    img_height : int
    img_width : int
    img_channels : int
    n_classes : int

    Returns
    -------
    images : numpy.ndarray
    maps : numpy.ndarray
    """

    if os.path.exists(os.path.join(data_path, "images.npy")) and os.path.exists(
        os.path.join(data_path, "maps.npy")
    ):
        images = np.load(f"{data_path}/images.npy")
        maps = np.load(f"{data_path}/maps.npy")

    else:
        images = next(os.walk(os.path.join(data_path, "images")))[-1]
        ids = [x[:-4] for x in images if x[-4:] == ".png"]

        # allocate arrays
        images = np.zeros(
            (len(ids), img_height, img_width, img_channels), dtype=np.uint8
        )
        maps = np.zeros((len(ids), img_height, img_width, n_classes), dtype=np.uint8)

        # fill arrays
        for i, id in len(ids):
            img = cv2.imread(
                os.path.join(data_path, "images", f"{id}_id.png"), cv2.IMREAD_UNCHANGED
            )
            if img_channels == 1:
                img = np.reshape(img, img.shape + (1,))

            # cut to square resolution on center
            img = img[
                :,
                img.shape[1] // 2
                - img.shape[0] // 2 : img.shape[1] // 2
                + img.shape[0] // 2,
            ]
            map = tiff.imread(os.path.join(data_path, "maps", f"{id}_map.tiff"))
            if n_classes == 1:
                map = np.reshape(map, map.shape + (1,))
            map = map[
                :,
                map.shape[1] // 2
                - map.shape[0] // 2 : map.shape[1] // 2
                + map.shape[0] // 2,
            ]
            images[i] = img
            maps[i] = map

            # save data as numpy arrays
            np.save(os.path.join(data_path, "images.npy"), images)
            np.save(os.path.join(data_path, "maps.npy"), maps)

    return images, maps


def custom_rotate(image):
    """
    Custom augmentation function to rotate image 0/90/180/270 degrees on random
    basis.

    Parameters
    ----------
    image : numpy.ndarray

    Returns
    -------
    image : numpy.ndarray
    """

    mode = np.random.randint(0, 4)
    return np.rot90(image, mode)


def image_map_generator(image_generator, map_generator):
    """
    Combines image and map generators to generate same augmentations.

    Parameters
    ----------
    image_generator : generator
    map_generator : generator

    Yields
    -------
    (img, map) : (numpy.ndarray, numpy.ndarray)
    """

    train_generator = zip(image_generator, map_generator)
    for (img, map) in train_generator:
        yield (img, map)
