from abc import ABC, abstractmethod

import numpy as np
from keras.preprocessing.image import ImageDataGenerator
from matplotlib.path import Path
from skimage.measure import block_reduce


def minmax(x):
    ax = (1, 2, 3)
    mi = np.min(x, ax, keepdims=1)
    ma = np.max(x, ax, keepdims=1)

    return (x - mi) / (ma - mi)


def downsample_image(image, by=2):
    """Downsample image by half by applying mean pooling"""
    return block_reduce(image, block_size=(by, by), func=np.mean)


def get_one_hot(targets, nb_classes):
    """Returns the targets as one-hot encoded"""
    res = np.eye(nb_classes)[np.array(targets).reshape(-1)]
    return res.reshape(list(targets.shape) + [nb_classes])


def create_masks(x, y):
    """
    Create masks from a dataset and its labels

    parameters:
        x: dataset of images
        y: array of labels of 2D coordinates
    returns:
        An array containing multiclass images segmented as follows:
            - Value 0 represents background
            - Value 1 represents the media
            - Value 2 represents the lumen
    """

    masks = []

    for i in range(y["lum"].shape[0]):
        # Retrieve path from the label
        path_lum = Path(y["lum"][i])  # TODO: why use mpl?
        path_med = Path(y["med"][i])
        img = x[i][2]

        # Create a matrix of coordinates to check which points are inside the closed path
        img_meshgrid = np.array(
            np.meshgrid(*(np.arange(s) for s in img.shape[:2]), indexing="xy")
        )
        img_coordinates = img_meshgrid.transpose([1, 2, 0]).reshape(
            img.shape[0] * img.shape[1], 2
        )
        mask_lum = path_lum.contains_points(img_coordinates).reshape(
            img.shape[0], img.shape[1]
        )
        mask_med = path_med.contains_points(img_coordinates).reshape(
            img.shape[0], img.shape[1]
        )

        final_mask = mask_lum.astype(int) + mask_med.astype(int)

        masks.append(final_mask)

    return np.array(masks)


def downsample(x, y, num_classes):
    x_ds = np.array([[downsample_image(image)
                    for image in batch] for batch in x])

    y_label_ds = [dict(lum=label["lum"] / 2.0, med=label["med"] / 2.0)
                  for label in y]
    masks = np.concatenate(
        [create_masks(x_ds, observer_labels) for observer_labels in y_label_ds]
    )

    y_ds = get_one_hot(masks, num_classes)
    x_ds = np.concatenate(
        [x_ds] * len(y_label_ds)
    )  # TODO: check if can be done with numpy
    return (x_ds, y_ds)


def augment(x_ds, y_ds):
    images = (
        x_ds[:, 2:5, :, :].transpose([0, 2, 3, 1]) / 255
    )  # Now the shape is (109, 384, 384, 3) (for dataset B)
    masks = y_ds  # Now the shape is (109, 384, 384, channels) (for dataset B)
    return images, masks


class Preprocessor(ABC):
    @abstractmethod
    def call(self, x):
        raise NotImplementedError()

    @abstractmethod
    def input_shape(self):
        raise NotImplementedError()

    def __call__(self, x):
        return self.call(x)


# from ivus_segmentation.signal_transforms import MonogenicFilter


class IdentityPreprocessor(Preprocessor):
    def __init__(self, input_shape) -> None:
        self._input_shape = input_shape

    def input_shape(self):
        return self._input_shape

    def call(self, x):
        return x


# class MonogenicSignalTransform(Preprocessor):
#     def __init__(self, wave_length, input_shape):
#         self.ysize, self.xsize, *rest = input_shape
#         self.cw = [wave_length]
#         print(input_shape, self.input_shape())

#         self.mfilter = MonogenicFilter(
#             self.ysize, self.xsize, self.cw, 'lg', 0.55)

#     def input_shape(self):
#         return self.ysize, self.xsize, 1 + 2 * len(self.cw)

#     def call(self, x):
#         m1, m2, m3 = map(np.squeeze,
#                          self.mfilter.monogenic_signal(x[..., 0]))
#         arr = np.stack([x[..., 0],
#                         m2,        m3], -1)
#         return arr


def get_image_generator(
    images,
    masks,
    batch_size,
    augment=False,
    seed=None,
    data_gen_args=dict(data_format="channels_last", validation_split=0.15),
):
    # we create two instances with the same arguments
    val_data_gen_args = dict(
        data_format=data_gen_args["data_format"],
        validation_split=data_gen_args["validation_split"],
    )

    image_datagen = ImageDataGenerator(**data_gen_args)
    mask_datagen = ImageDataGenerator(**data_gen_args)

    val_datagen = ImageDataGenerator(**val_data_gen_args)

    # TODO: set to false so it doesnt kill my memory
    image_datagen.fit(images, augment=augment, seed=seed)
    mask_datagen.fit(masks, augment=augment, seed=seed)

    # Provide the same seed and keyword arguments to the fit and flow methods
    image_generator = image_datagen.flow(
        images, seed=seed, batch_size=batch_size, subset="training"
    )
    mask_generator = mask_datagen.flow(
        masks, seed=seed, batch_size=batch_size, subset="training"
    )

    image_val_generator = val_datagen.flow(
        images, seed=seed, batch_size=batch_size, subset="validation"
    )
    mask_val_generator = val_datagen.flow(
        masks, seed=seed, batch_size=batch_size, subset="validation"
    )

    # combine generators into one which yields image and masks
    train_generator = zip(image_generator, mask_generator)
    val_generator = zip(image_val_generator, mask_val_generator)

    return train_generator, val_generator


def gen_dtype(gen):
    while True:
        for i, j in gen:
            yield (i.astype("float32"), j.astype("float32"))
