from abc import ABC, abstractmethod
from typing import Literal, Optional

import numpy as np  # linear algebra
import tensorflow as tf
import tensorflow_addons as tfa
from scipy.ndimage import gaussian_filter
from scipy.ndimage.filters import gaussian_filter


class Augment(ABC):
    """Abstract class for data augmentation"""

    @abstractmethod
    def call(self, x):
        raise NotImplemented()

    def __call__(self, x, y):
        t = tf.concat([x, y], axis=-1)
        t = self.call(t)
        x, y = tf.split(t, 2, axis=-1)
        y = tf.round(y)
        bg = 1 - tf.maximum(y[..., 1], y[..., 2])  # bias to bg
        y = tf.stack([bg, y[..., 1], y[..., 2]], -1)
        return x, y


class RandomFlip(Augment):
    """Extend `Augment` to apply a random flip to the image (both axes)"""

    def call(self, x):
        return tf.image.random_flip_left_right(x)


class RandomRotation(Augment):
    """Extend `Augment` to apply a random rotation in the range [0, 2π).

    Parameters
    ----------
    range : float
        fraction of the rotation range [0, 2π)
    """

    def __init__(self, range) -> None:
        self.range = range
        self.layer = tf.keras.layers.RandomRotation(self.range,
                                                    fill_mode='constant')

    def call(self, x):
        return self.layer(x)
        # tf.keras.preprocessing.image.random_rotation(x, self.range)


class RandomZoom(Augment):
    """Extend `Augment` to apply a random zoom tot he image

    Parameters
    ----------
    range : float
        represents the maximum zoom distance in terms of the fraction of the
        size of the image.
    """

    def __init__(self, range) -> None:
        self.range = range  # (range, range)
        self.layer = tf.keras.layers.RandomZoom(self.range,
                                                fill_mode='constant')

    def call(self, x):
        return self.layer(x)
        # tf.keras.preprocessing.image.random_zoom(x, self.range)


class ElasticDeformation(Augment):
    """Extend `Augment` to apply an elastic deformation

    Parameters
    ----------
    a: float
        Scale [TODO: too vague]
    b: float
        Standard deviation for the Gaussian kernel.
    c: float
        TODO
    """

    def __init__(self, alpha, sigma) -> None:
        super().__init__()
        self.alpha = alpha = float(alpha)
        self.sigma = sigma = float(sigma)

    def call(self, x):
        return tf.numpy_function(elastic_distortion,
                                 [x,
                                  self.alpha * x.shape[1],
                                  self.sigma * x.shape[1]],
                                 x.dtype)


class ElasticDeformationOnlyMasks(Augment):
    """Extend `Augment` to apply an elastic deformation

    Parameters
    ----------
    a: float
        Scale [TODO: too vague]
    b: float
        Standard deviation for the Gaussian kernel.
    c: float
        TODO
    """

    def __init__(self, alpha, sigma) -> None:
        super().__init__()
        self.alpha = alpha = float(alpha)
        self.sigma = sigma = float(sigma)

    def call(self, x):
        x, y = tf.split(x, 2, axis=-1)
        y = tf.numpy_function(elastic_distortion,
                              [y,
                               self.alpha * y.shape[1],
                               self.sigma * y.shape[1]],
                              y.dtype)

        t = tf.concat([x, y], axis=-1)
        return t


def perform_augments(x, y, batch_size, *args) -> tf.data.Dataset:
    """Given a set of transformations, create a dataset with the applied
    transformations.
    """
    ds = tf.data.Dataset.from_tensor_slices((x, y))

    if len(args) == 1:
        # iterable
        try:
            augments = list(args[0])
        except TypeError:
            augments = [args[0]]
    else:
        augments = args
    for augment in augments:
        ds = ds.map(augment)
    ds = ds.batch(batch_size)
    ds = ds.prefetch(tf.data.AUTOTUNE)

    return ds


def gaussian_kernel(shape=(3, 3), sigma=0.5, epsilon=1e7):
    """create a `shape` sized gaussian kernel
    """
    m, n = [(ss-1.)/2. for ss in shape]
    y, x = np.ogrid[-m:m+1, -n:n+1]
    y = np.maximum(y - epsilon, 0) + epsilon
    x = np.maximum(x - epsilon, 0) + epsilon

    h = np.exp(-1 * (x * x + y * y) / 2.0 * sigma * sigma)

    return h / (np.sum(h) + epsilon)


def elastic_distortion(src: np.ndarray, alpha: np.float32, sigma: np.float32,
                       dst: Optional[np.ndarray] = None,
                       gen: Optional[np.random.Generator] = None,
                       mode: Literal['zero'] = 'zero'):
    """Performs elastic deformation as described in the report"""
    assert len(src.shape) == 3
    assert mode in ['zero'], "allowing padding modes are: ['zero']"

    rows, cols = src.shape[:-1]

    gen = gen or np.random.default_rng()
    dst = np.zeros_like(src)

    ksize = np.int64(sigma * 3) * 2 + 1
    ksize = (ksize, ksize)

    dx = gen.uniform(-1, 1, size=(rows, cols))
    dx = np.multiply(dx, alpha, out=dx)
    dx = gaussian_filter(dx, sigma)
    # cv2.GaussianBlur(dx, ksize, sigma, dst=dx)

    dy = gen.uniform(-1, 1, size=(rows, cols))
    dy = gaussian_filter(dy, sigma)
    # dy = cv2.GaussianBlur(dy, ksize, sigma, dst=dy)
    dy = np.multiply(dy, alpha, out=dy)

    i, j = np.mgrid[:rows, :cols]
    oy, ox = np.int32(i - dy[i, j]), \
        np.int32(j - dx[i, j])

    t = (0 <= oy) & (oy < rows) & (0 <= ox) & (ox < cols)

    i = i[t]
    oy = oy[t]

    j = j[t]
    ox = ox[t]

    dst[i, j] = src[oy, ox]

    return dst
