##############################################

#                                            #

#           Custom data generator            #

#                                            #

# Author: Amine Neggazi                      #

# Email: neggazimedlamine@gmail/com          #

# Nick: nemo256                              #

#                                            #

# Please read bc-count/LICENSE               #

#                                            #

##############################################

import os

import json

import cv2

import numpy as np

import tensorflow as tf

from tensorflow import keras

# custom imports

from config import *

def load_image_list(img_files, gray=False):

    '''

    This is the load image list function, which loads an enumerate

    of images (param: img_files)

    :param img_files --> the input image files which we want to read

    :return imgs --> the images that we read

    '''

    imgs = []

    if gray:

        for image_file in img_files:

            img = cv2.imread(image_file, cv2.IMREAD_GRAYSCALE)

            img = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)[1]

            imgs += [img]

    else:

        for image_file in img_files:

            imgs += [cv2.imread(image_file)]

    return imgs

def clahe_images(img_list):

    '''

    This is the clahe images function, which applies a clahe threshold

    the input image list.

    :param img_files --> the input image files which we want to read

    :return img_list --> the output images

    '''

    for i, img in enumerate(img_list):

        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))

        lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)

        lab[..., 0] = clahe.apply(lab[..., 0])

        img_list[i] = cv2.cvtColor(lab, cv2.COLOR_LAB2BGR)

    return img_list

def preprocess_image(imgs, padding=padding[1]):

    '''

    This is the preprocess data function, which adds a padding to 

    the input images, masks and edges if there are any.

    :param imgs --> the input list of images.

    :param padding --> the input padding which is going to be applied.

    :return imgs --> output images with added padding.

    '''

    imgs = [np.pad(img, ((padding, padding),

                         (padding, padding), (0, 0)), mode='constant') for img in imgs]

    return imgs

def preprocess_data(imgs, mask, edge=None, padding=padding[1]):

    '''

    This is the preprocess data function, which adds a padding to 

    the input images, masks and edges if there are any.

    :param imgs --> the input list of images.

    :param mask --> the input list of masks.

    :param edge --> the input list of edges.

    :param padding --> the input padding which is going to be applied.

    :return tuple(imgs, mask, edge if exists) --> output images, masks and edges with padding added.

    '''

    imgs = [np.pad(img, ((padding, padding),

                         (padding, padding), (0, 0)), mode='constant') for img in imgs]

    mask = [np.pad(mask, ((padding, padding),

                          (padding, padding)), mode='constant') for mask in mask]

    if edge is not None:

        edge = [np.pad(edge, ((padding, padding),

                              (padding, padding)), mode='constant') for edge in edge]

    if edge is not None:

        return imgs, mask, edge

    return imgs, mask

def load_data(img_list, mask_list, edge_list=None, padding=padding[1]):

    '''

    This is the load data function, which will handle image loading and preprocessing.

    :param img_list --> list of input images

    :param mask_list --> list of input masks

    :param edge_list --> list of input edges

    :param padding --> padding to be applied on preprocessing

    :return tuple(imgs, masks and edges if exists) --> the output preprocessed imgs, masks and edges.

    '''

    imgs = load_image_list(img_list)

    imgs = clahe_images(imgs)

    mask = load_image_list(mask_list, gray=True)

    if edge_list:

        edge = load_image_list(edge_list, gray=True)

    else:

        edge = None

    return preprocess_data(imgs, mask, edge, padding=padding)

def load_image(img_list, padding=padding[1]):

    '''

    This is the load data function, which will handle image loading and preprocessing.

    :param img_list --> list of input images

    :param padding --> padding to be applied on preprocessing

    :return imgs --> the output preprocessed imgs.

    '''

    imgs = load_image_list(img_list)

    imgs = clahe_images(imgs)

    return preprocess_image(imgs, padding=padding)

def aug_lum(image, factor=None):

    '''

    This is the augment luminosity function, which we apply to

    augment the luminosity of an input image.

    :param image --> the input image we want to augment

    :param factor --> the factor of luminosity augment (default is 0.5 * random number)

    :return image --> the output luminosity augmented image

    '''

    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    hsv = hsv.astype(np.float64)

    if factor is None:

        lum_offset = 0.5 + np.random.uniform()

    else:

        lum_offset = factor

    hsv[..., 2] = hsv[..., 2] * lum_offset

    hsv[..., 2][hsv[..., 2] > 255] = 255

    hsv = hsv.astype(np.uint8)

    return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)

def aug_img(image):

    '''

    This is the augment colors function, which we apply to

    augment the colors of an given image.

    :param image --> the input image we want to augment

    :return image --> the output colors augmented image

    '''

    hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

    hsv = hsv.astype(np.float64)

    hue_offset = 0.8 + 0.4*np.random.uniform()

    sat_offset = 0.5 + np.random.uniform()

    lum_offset = 0.5 + np.random.uniform()

    hsv[..., 0] = hsv[..., 0] * hue_offset

    hsv[..., 1] = hsv[..., 1] * sat_offset

    hsv[..., 2] = hsv[..., 2] * lum_offset

    hsv[..., 0][hsv[..., 0] > 255] = 255

    hsv[..., 1][hsv[..., 1] > 255] = 255

    hsv[..., 2][hsv[..., 2] > 255] = 255

    hsv = hsv.astype(np.uint8)

    return cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)

def train_generator(imgs, mask, edge=None,

                    scale_range=None,

                    padding=padding[1],

                    input_size=input_shape[0],

                    output_size=output_shape[0],

                    skip_empty=False):

    '''

    This is the train generator function, which generates the train dataset.

    :param imgs --> the input images

    :param mask --> the input masks

    :param edge --> the input edges if there are any (red blood cells only)

    :param scale_range --> the factor (i, j) of rescaling.

    :param padding --> the padding which will be applied to each image

    :param input_size --> the input shape

    :param output_size --> the output shape

    :param skip_empty --> skip empty chips (random if not set)

    :return chips --> yields an image, mask and edge chip each time it gets executed (called)

    '''

    if scale_range is not None:

        scale_range = [1 - scale_range, 1 + scale_range]

    while True:

        # select which type of cell to return

        chip_type = np.random.choice([True, False])

        while True:

            # pick random image

            i = np.random.randint(len(imgs))

            # pick random central location in the image (200 + 196/2)

            center_offset = padding + (output_size / 2)

            x = np.random.randint(center_offset, imgs[i].shape[0] - center_offset)

            y = np.random.randint(center_offset, imgs[i].shape[1] - center_offset)

            # scale the box randomly from x0.8 - 1.2x original size

            scale = 1

            if scale_range is not None:

                scale = scale_range[0] + ((scale_range[0] - scale_range[0]) * np.random.random())

            # find the edges of a box around the image chip and the mask chip

            chip_x_l = int(x - ((input_size / 2) * scale))

            chip_x_r = int(x + ((input_size / 2) * scale))

            chip_y_l = int(y - ((input_size / 2) * scale))

            chip_y_r = int(y + ((input_size / 2) * scale))

            mask_x_l = int(x - ((output_size / 2) * scale))

            mask_x_r = int(x + ((output_size / 2) * scale))

            mask_y_l = int(y - ((output_size / 2) * scale))

            mask_y_r = int(y + ((output_size / 2) * scale))

            # take a slice of the image and mask accordingly

            temp_chip = imgs[i][chip_x_l:chip_x_r, chip_y_l:chip_y_r]

            temp_mask = mask[i][mask_x_l:mask_x_r, mask_y_l:mask_y_r]

            if edge is not None:

                temp_edge = edge[i][mask_x_l:mask_x_r, mask_y_l:mask_y_r]

            if skip_empty:

                if ((temp_mask > 0).sum() > 5) is chip_type:

                    continue

            # resize the image chip back to 380 and the mask chip to 196

            temp_chip = cv2.resize(temp_chip,

                                   (input_size, input_size),

                                   interpolation=cv2.INTER_CUBIC)

            temp_mask = cv2.resize(temp_mask,

                                   (output_size, output_size),

                                   interpolation=cv2.INTER_NEAREST)

            if edge is not None:

                temp_edge = cv2.resize(temp_edge,

                                       (output_size, output_size),

                                       interpolation=cv2.INTER_NEAREST)

            # randomly rotate (like below)

            rot = np.random.randint(4)

            temp_chip = np.rot90(temp_chip, k=rot, axes=(0, 1))

            temp_mask = np.rot90(temp_mask, k=rot, axes=(0, 1))

            if edge is not None:

                temp_edge = np.rot90(temp_edge, k=rot, axes=(0, 1))

            # randomly flip

            if np.random.random() > 0.5:

                temp_chip = np.flip(temp_chip, axis=1)

                temp_mask = np.flip(temp_mask, axis=1)

                if edge is not None:

                    temp_edge = np.flip(temp_edge, axis=1)

            # randomly luminosity augment

            temp_chip = aug_lum(temp_chip)

            # randomly augment chip

            temp_chip = aug_img(temp_chip)

            # rescale the image

            temp_chip = temp_chip.astype(np.float32) * 2

            temp_chip /= 255

            temp_chip -= 1

            # later on ... randomly adjust colours

            if edge is not None:

                yield temp_chip, ((temp_mask > 0).astype(float)[..., np.newaxis], 

                                  (temp_edge > 0).astype(float)[..., np.newaxis])

            else:

                yield temp_chip, ((temp_mask > 0).astype(float)[..., np.newaxis])

            break

def test_chips(imgs, mask,

               edge=None,

               padding=padding[1],

               input_size=input_shape[0],

               output_size=output_shape[0]):

    '''

    This is the test chips function, which generates the test dataset.

    :param imgs --> the input images

    :param mask --> the input masks

    :param edge --> the input edges if there are any (red blood cells only)

    :param padding --> the padding which will be applied to each image

    :param input_size --> the input shape

    :param output_size --> the output shape

    :return chips --> yields an image, mask and edge chip each time it gets executed (called)

    '''

    center_offset = padding + (output_size / 2)

    for i, _ in enumerate(imgs):

        for x in np.arange(center_offset, imgs[i].shape[0] - input_size / 2, output_size):

            for y in np.arange(center_offset, imgs[i].shape[1] - input_size / 2, output_size):

                chip_x_l = int(x - (input_size / 2))

                chip_x_r = int(x + (input_size / 2))

                chip_y_l = int(y - (input_size / 2))

                chip_y_r = int(y + (input_size / 2))

                mask_x_l = int(x - (output_size / 2))

                mask_x_r = int(x + (output_size / 2))

                mask_y_l = int(y - (output_size / 2))

                mask_y_r = int(y + (output_size / 2))

                temp_chip = imgs[i][chip_x_l:chip_x_r, chip_y_l:chip_y_r]

                temp_mask = mask[i][mask_x_l:mask_x_r, mask_y_l:mask_y_r]

                if edge is not None:

                    temp_edge = edge[i][mask_x_l:mask_x_r, mask_y_l:mask_y_r]

                temp_chip = temp_chip.astype(np.float32) * 2

                temp_chip /= 255

                temp_chip -= 1

                if edge is not None:

                    yield temp_chip, ((temp_mask > 0).astype(float)[..., np.newaxis], 

                                      (temp_edge > 0).astype(float)[..., np.newaxis])

                else:

                    yield temp_chip, ((temp_mask > 0).astype(float)[..., np.newaxis])

                break

def slice_image(imgs,

                padding=padding[1],

                input_size=input_shape[0],

                output_size=output_shape[0]):

    '''

    This is the slice function, which slices each image into image chips.

    :param imgs --> the input images

    :param padding --> the padding which will be applied to each image

    :param input_size --> the input shape

    :param output_size --> the output shape

    :return list tuple (list, list, list) --> the tuple list of output (image, mask and edge chips)

    '''

    img_chips = []

    center_offset = padding + (output_size / 2)

    for i, _ in enumerate(imgs):

        for x in np.arange(center_offset, imgs[i].shape[0] - input_size / 2, output_size):

            for y in np.arange(center_offset, imgs[i].shape[1] - input_size / 2, output_size):

                chip_x_l = int(x - (input_size / 2))

                chip_x_r = int(x + (input_size / 2))

                chip_y_l = int(y - (input_size / 2))

                chip_y_r = int(y + (input_size / 2))

                temp_chip = imgs[i][chip_x_l:chip_x_r, chip_y_l:chip_y_r]

                temp_chip = temp_chip.astype(np.float32) * 2

                temp_chip /= 255

                temp_chip -= 1

                img_chips += [temp_chip]

    return np.array(img_chips)

def slice(imgs, mask,

          edge=None,

          padding=padding[1],

          input_size=input_shape[0],

          output_size=output_shape[0]):

    '''

    This is the slice function, which slices each image into image chips.

    :param imgs --> the input images

    :param mask --> the input masks

    :param edge --> the input edges if there are any (red blood cells only)

    :param padding --> the padding which will be applied to each image

    :param input_size --> the input shape

    :param output_size --> the output shape

    :return list tuple (list, list, list) --> the tuple list of output (image, mask and edge chips)

    '''

    img_chips = []

    mask_chips = []

    if edge is not None:

        edge_chips = []

    center_offset = padding + (output_size / 2)

    for i, _ in enumerate(imgs):

        for x in np.arange(center_offset, imgs[i].shape[0] - input_size / 2, output_size):

            for y in np.arange(center_offset, imgs[i].shape[1] - input_size / 2, output_size):

                chip_x_l = int(x - (input_size / 2))

                chip_x_r = int(x + (input_size / 2))

                chip_y_l = int(y - (input_size / 2))

                chip_y_r = int(y + (input_size / 2))

                mask_x_l = int(x - (output_size / 2))

                mask_x_r = int(x + (output_size / 2))

                mask_y_l = int(y - (output_size / 2))

                mask_y_r = int(y + (output_size / 2))

                temp_chip = imgs[i][chip_x_l:chip_x_r, chip_y_l:chip_y_r]

                temp_mask = mask[i][mask_x_l:mask_x_r, mask_y_l:mask_y_r]

                if edge is not None:

                    temp_edge = edge[i][mask_x_l:mask_x_r, mask_y_l:mask_y_r]

                temp_chip = temp_chip.astype(np.float32) * 2

                temp_chip /= 255

                temp_chip -= 1

                img_chips += [temp_chip]

                mask_chips += [(temp_mask > 0).astype(float)[..., np.newaxis]]

                if edge is not None:

                    edge_chips += [(temp_edge > 0).astype(float)[..., np.newaxis]]

    img_chips = np.array(img_chips)

    mask_chips = np.array(mask_chips)

    if edge is not None:

        edge_chips = np.array(edge_chips)

    if edge is not None:

        return img_chips, mask_chips, edge_chips

    return img_chips, mask_chips

def generator(img_list, mask_list, edge_list=None, type='train'):

    '''

    This is the generator function, which provides the list of image, mask and edge lists to the train generator and test chips functions.

    :param img_list --> the input list of images

    :param mask_list --> the input list of masks

    :param edge_list --> the input list of edges if there are any

    :param type --> can be either train or test, used to determine which generator function is to be called

    :return tensorflow dataset --> the output generated functions fed to tensorflow

    '''

    if cell_type == 'rbc':

        img, mask, edge = load_data(img_list, mask_list, edge_list)

    elif cell_type == 'wbc' or cell_type == 'plt':

        img, mask = load_data(img_list, mask_list)

        edge = None

    def gen():

        if type == 'train':

            return train_generator(img, mask, edge,

                                   padding=padding[0],

                                   input_size=input_shape[0],

                                   output_size=output_shape[0])

        elif type == 'test':

            return test_chips(img, mask, edge,

                              padding=padding[0],

                              input_size=input_shape[0],

                              output_size=output_shape[0])

    # load train dataset to tensorflow for training

    if cell_type == 'rbc':

        return tf.data.Dataset.from_generator(

            gen,

            (tf.float64, ((tf.float64), (tf.float64))),

            (input_shape, (output_shape, output_shape))

        )

    elif cell_type == 'wbc' or cell_type == 'plt':

        return tf.data.Dataset.from_generator(

            gen,

            (tf.float64, (tf.float64)),

            (input_shape, (output_shape))

        )