"""Methods/Shortcuts for modifying and handling image data."""
import warnings

import cv2
import numpy as np

from .goudapath import GoudaPath

__author__ = "Sean Mullan"
__copyright__ = "Sean Mullan"
__license__ = "mit"


def imread(path, as_RGB=False, as_greyscale=False, unchanged=True):
    """SHORTCUT: Load an image from a path using OpenCV modified for RGB.

    Parameters
    ----------
    path: string
        Path to image file
    as_RGB: bool
        If true, assumes an RGB image and flips the channels. (OpenCV uses BGR by default) (the default is True)
    as_greyscale: bool
        If true, reads the image as greyscale (the default is False)
    unchanged: bool
        If true, reads in the image without constraining to 3-channel uint8 format. (the default is True)

    Note
    ----
    * Priority for options is: as_greyscale > as_RGB > unchanged > default
    * Greyscale transforms input image based on perceived color. See [https://docs.opencv.org/2.4/modules/imgproc/doc/miscellaneous_transformations.html#cvtcolor]
    """
    if isinstance(path, GoudaPath):
        path = path.path
    if as_greyscale:
        return cv2.imread(path, 0)
    elif as_RGB:
        return cv2.imread(path)[:, :, ::-1]
    elif unchanged:
        return cv2.imread(path, -1)
    else:
        return cv2.imread(path)


def imwrite(path, image, as_RGB=True):
    """SHORTCUT: Write an image to a path using OpenCV modified for RGB.

    Parameters
    ----------
    path: string
        Path to save image file to
    image: image data
        image data to save - must be uint8/uint16 and with shape [x, y], [x, y, 1], or [x, y, 3]
    as_RGB: bool
        If true, flips the channels before saving (OpenCV assumes BGR image by default)
        """
    if isinstance(path, GoudaPath):
        path = path.path
    if image.ndim == 2:
        image = image[:, :, np.newaxis]
        as_RGB = False
    elif image.ndim == 3 and image.shape[2] == 1:
        as_RGB = False
    elif image.ndim == 3 and image.shape[2] == 3:
        pass
    else:
        raise ValueError("Image must be of shape [x, y, 1], [x, y, 3], or [x, y], not {}".format(image.shape))
    if as_RGB:
        cv2.imwrite(path, image[:, :, ::-1])
    else:
        cv2.imwrite(path, image)


def to_uint8(x):
    """Convert an image to a uint8 type with range [0, 255] based on inferred normalization type."""
    if x.dtype == np.uint8:
        return x
    if x.max() > 1 and x.max() <= 255 and x.min() >= 0:  # input range [0, 255]
        pass
    elif x.min() < 0 and x.min() >= -1 and x.max() <= 1:  # input range [-1, 1]
        x = ((x * 127.5) + 127.5)
    elif x.min() >= 0 and x.max() <= 1:  # input range [0, 1]
        x = (x * 255.0)
    else:
        warnings.warn("Cannot determine input range. Rescaling to [0, 1]")
        x = rescale(x, min_val=0, max_val=255)
    return x.astype(np.uint8)


def rescale(data, column_wise=False, max_val=1, min_val=0, return_type=np.float):
    """Scale either an image or rows of data to [0, 1].

        NOTE
        ----
        If the range of values in a column/image is 0, data will scale to 0
    """
    data = data.astype(np.float)
    if column_wise:
        range = data.max(axis=0) - data.min(axis=0)
        range[range == 0] = 1
        rescaled = (data - data.min(axis=0)) / range
        rescaled[range == 0] = 0
    else:
        range = data.max() - data.min() + 0.0
        if range == 0:
            rescaled = np.zeros_like(data)
        else:
            rescaled = (data - data.min()) / range
    return ((rescaled * (max_val - min_val)) + min_val).astype(return_type)


def stack_label(label, label_channel=0, as_uint8=True):
    """Convert 2d label to 3d.

    Parameters
    ----------
    label : np.ndarray
        2D label to stack
    label_channel : int
        The color channel for the label, or -1 to set all channels to the label (the default is 0)
    as_uint8 : bool
        Whether to convert the label to a uint8
    """
    label = np.squeeze(label)
    if as_uint8:
        label = to_uint8(label)
    if label_channel < 0:
        return np.dstack([label, label, label])
    elif label_channel < 3:
        to_stack = [np.zeros_like(label), np.zeros_like(label), np.zeros_like(label)]
        to_stack[label_channel] = label
        return np.dstack(to_stack)
    else:
        raise ValueError("Not a valid color channel index: {}".format(label_channel))


def laplacian_var(image):
    """Return the laplacian variance of an image"""
    # Laplacian is the rate of change of pixel intensity (2nd order derivative)
    blur = cv2.GaussianBlur(image, (3, 3), 0, 0)
    grey = cv2.cvtColor(blur, cv2.COLOR_BGR2GRAY)
    return cv2.Laplacian(grey, cv2.CV_16S).var()


def sobel_var(image):
    """Return the sobal variance of an image"""
    # Sobel is the gradient of pixel intensity (1st order derivative)
    blur = cv2.GaussianBlur(image, (3, 3), 0, 0)
    grey = cv2.cvtColor(blur, cv2.COLOR_BGR2GRAY)
    grad_x = cv2.Sobel(grey, cv2.CV_16S, 1, 0, 3)
    grad_x = cv2.convertScaleAbs(grad_x)
    grad_y = cv2.Sobel(grey, cv2.CV_16S, 0, 1, 3)
    grad_y = cv2.convertScaleAbs(grad_y)
    return cv2.addWeighted(grad_x, 0.5, grad_y, 0.5, 0).var()


def split_signs(mask):
    """Split a single channel image mask with +/- values into green/red color channels"""
    mask = np.squeeze(mask)
    if mask.ndim != 2:
        raise ValueError("Single channel mask required for split_signs")
    negative = np.clip(mask, -np.inf, 0) * -1
    positive = np.clip(mask, 0, np.inf)
    new_mask = np.dstack([negative, positive, np.zeros_like(positive)])
    return new_mask


def add_overlay(image, mask, label_channel=0, separate_signs=False):
    """Return image with a mask overlay.

    Parameters
    ----------
        image: numpy.ndarray
            Image array with shape (x, y, channels)
        mask: numpy.ndarray
            np array with shape (x, y) and range [0, 1]
        label_channel : int
            The color channel to use for the overlay if separate_signs is False (the default is 0)
        separate_signs : bool
            Whether to separate +/- values of the mask into green/red channels
    """
    output = np.squeeze(image)
    mask = np.squeeze(mask)

    if mask.shape[:2] != image.shape[:2]:
        raise ValueError('Mask width/height does not match image width/height: {} != {}'.format(mask.shape[:2], image.shape[:2]))
    if mask.ndim == 2:
        if separate_signs:
            mask = split_signs(mask)
        else:
            mask = stack_label(mask, label_channel, as_uint8=False)
    elif mask.ndim == 3 and mask.shape[2] == 3:
        pass
    else:
        raise ValueError("Mask must have shape [x, y] or [x, y, z], not {}".format(mask.shape))
    if image.dtype == np.uint8:
        mask = to_uint8(mask)

    if output.ndim == 2:
        output = np.dstack([output, output, output])

    overlay = cv2.addWeighted(output, 0.5, mask, 0.5, 0)
    output = np.where(mask.sum(axis=2, keepdims=True) != 0, overlay, output)
    return output


def masked_lineup(image, label):
    """Return a list of image, masked_image, mask.

    Note
    ----
        Image will not be changed, but the masked_image and mask will be uint8 arrays with range [0, 255].

    Parameters
    ----------
        image: np array
            Image array with shape (x, y, channels)
        label: np array
            np array with shape (x, y) and range [0, 1]
        norm_type: int | None
            Type of normalization to undo when masking the image. If None, it will
            be inferred based on the value range of the array.
    """
    return [
        image,
        add_overlay(image, label),
        np.dstack([label * 255,
                   np.zeros_like(label),
                   np.zeros_like(label)]).astype(np.uint8)
    ]


def grabCut(image, labels, use_tresholds=True, thresholds=(0.2, 0.6, 0.8), iterations=2, clean=False):
    """Use the predicted mask and OpenCV's GrabCut algorithm to mask an image.

    Parameters
    ----------
    image : numpy.ndarray
        input image to mask
    labels : numpy.ndarray
        labels for background/foreground.
    use_thresholds: bool
        If true, thresholds the labels to find FG/BG components. Otherwise assumes that the values are [0, 2, 3, 1] defined by OpenCV for [BG, PR_BG, PR_FG, FG]
    thresholds : (int, int, int)
        Thesholds to divide [BG, PR_BG, PR_FG, FG] in the input labels
    iterations : int
        Number of GrabCut iterations to perform (the default is 2).
    clean : bool
        Should the GrabCut results be smoothed using :func:`~gouda.image.clean_grabCut_mask` (the default is False).

    Returns
    -------
    The GrabCut masked image and the generated mask : numpy.ndarray, numpy.ndarray

    """
    if labels.max() <= 1:
        mask = np.zeros_like(labels, dtype=np.uint8)
        mask += cv2.GC_BGD
        mask[labels > thresholds[0]] = cv2.GC_PR_BGD
        mask[labels > thresholds[1]] = cv2.GC_PR_FGD
        mask[labels > thresholds[2]] = cv2.GC_FGD
    else:
        mask = labels.astype(np.uint8)

    if cv2.GC_FGD not in mask and cv2.GC_PR_FGD not in mask:
        raise ValueError("Labels cannot all be background")
    elif cv2.GC_BGD not in mask and cv2.GC_PR_BGD not in mask:
        raise ValueError("Labels cannot all be foreground")

    bg, fg = np.zeros((1, 65), np.float64), np.zeros((1, 65), np.float64)
    cv2.grabCut(image, mask, None, bg, fg, iterations, cv2.GC_INIT_WITH_MASK)

    if clean:
        mask = clean_grabCut_mask(mask)
    mask_out = np.where((mask == 1) + (mask == 3), 1, 0).astype(np.uint8)
    masked_img = cv2.bitwise_and(image, image, mask=mask_out)
    return masked_img, mask_out


def clean_grabCut_mask(mask):
    """Apply light smoothing to grabCut labels"""
    temp = np.copy(mask)
    change = np.where((temp == 2) + (temp == 3), 1, 0)
    remapped = np.zeros_like(temp)
    remapped[temp == 0] = 0
    remapped[temp == 1] = 3
    remapped[temp == 2] = 1
    remapped[temp == 3] = 2
    up = cv2.pyrUp(remapped)
    for i in range(15):
        up = cv2.medianBlur(up, 7)
    down = cv2.pyrDown(up)
    down = np.round(down)
    if temp.ndim == 3:
        down = down[:, :, np.newaxis]
    remapped = np.where(change == 1, down, remapped)
    output = np.zeros_like(temp)
    output[remapped == 0] = 0
    output[remapped == 3] = 1
    output[remapped == 1] = 2
    output[remapped == 2] = 3
    return output.reshape(mask.shape).astype(mask.dtype)


def crop_to_mask(image, mask, with_label=False, smoothing=True):
    """Crop input image to only be size of input mask.

    Parameters
    ----------
        image: numpy uint8 array
            Image to mask/crop
        mask: numpy bool array
            Masking array with same x and y shape as image
        with_label: bool
            Whether or not to return a version of the mask that is cropped to the new boundaries
        smoothing: bool
            Whether or not to smooth the border of the mask

    Returns
    -------
        numpy uint8 array
            Masked image with dimensions that are only as large
            as the positive region of the mask
    """
    vert = np.any(mask, axis=0)
    y_range = np.where(vert == True)  # noqa: E712
    horiz = np.any(mask, axis=1)
    x_range = np.where(horiz == True)  # noqa: E712
    x0, x1 = x_range[0][0], x_range[0][-1]
    y0, y1 = y_range[0][0], y_range[0][-1]
    if smoothing:
        smooth_mask = cv2.GaussianBlur(
            mask.astype(np.uint8) * 255.0, (5, 5), 0)
        # mask[mask<0.8] = 0
        # mask[mask>=0.8] = 1
    else:
        smooth_mask = mask.astype(np.uint8) * 255.0
    masked_image = cv2.bitwise_and(image,
                                   image,
                                   mask=smooth_mask.astype(np.uint8))
    if with_label:
        return masked_image[x0:x1, y0:y1], mask[x0:x1, y0:y1]
    return masked_image[x0:x1, y0:y1]


def get_bounds(mask):
    """Return the bounding box corners of the positive content for the mask."""
    vert = np.any(mask, axis=0)
    horiz = np.any(mask, axis=1)
    y_range = np.where(vert == True)  # noqa: E712
    x_range = np.where(horiz == True)  # noqa: E712
    x0, x1 = x_range[0][0], x_range[0][-1]
    y0, y1 = y_range[0][0], y_range[0][-1]
    return (x0, y0), (x1, y1)


def crop_to_content(image):
    """Crop image to only be as large as the contained image excluding black space."""
    vert = np.mean(image, axis=(-1, 0))
    y_range = np.where(vert > 0)
    horiz = np.mean(image, axis=(-1, 1))
    x_range = np.where(horiz > 0)
    x0, x1 = x_range[0][0], x_range[0][-1]
    y0, y1 = y_range[0][0], y_range[0][-1]
    return image[x0:x1, y0:y1]


def rotate(img, degrees=90, allow_resize=True):
    """Rotate image clock-wise using OpenCV.

    Parameters
    ----------
    img: image (as required by OpenCV getRotationMatrix2D)
        image/path to rotate
    degrees: int
        Degrees to rotate the image clockwise
    allow_resize: bool
        If yes, the image boundaries will change to fit the rotated image. Otherwise the output rotated image is cropped to fit the original boundaries.
    """
    (h, w) = img.shape[:2]
    (cX, cY) = (w / 2, h / 2)
    M = cv2.getRotationMatrix2D((cX, cY), -degrees, 1.0)
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])
    if allow_resize:
        nW = int((h * sin) + (w * cos))
        nH = int((h * cos) + (w * sin))
    else:
        nW, nH = w, h
    M[0, 2] += (nW / 2) - cX
    M[1, 2] += (nH / 2) - cY
    # Fixes a 1-pixel offset courtesy of: https://github.com/opencv/opencv/issues/4585#issuecomment-397895187
    M[0, 2] += (M[0, 0] + M[0, 1] - 1) / 2
    M[1, 2] += (M[1, 0] + M[1, 1] - 1) / 2
    return cv2.warpAffine(img, M, (nW, nH)).astype(img.dtype)


def padded_resize(image, size=[960, 540], allow_rotate=True):
    """Resize input image to given size, only padding as needed for aspect ratio.

    Parameters
    ----------
    image: image/path
        image/path to image to resize
    size: [int, int]
        output size
    allow_rotate: bool
        Whether the image can be rotated to minimize required padding or if orientation should be preserved
    Input image number of channels does not matter as long as the first two dimensions are x and y.
    """
    if type(image) == str:
        image = imread(image, as_RGB=True)
    data_type = image.dtype
    if image.ndim == 2:
        x, y = image.shape
        c = 0
        image = image[:, :, np.newaxis]
    else:
        x, y, c = image.shape
    if ((size[0] < size[1] and x > y) or (size[0] > size[1] and x < y)) and allow_rotate:
        image = rotate(image, 90)
        x, y = image.shape[:2]
    padx = int(y * (float(size[0]) / size[1])) - x
    pady = int(x * (float(size[1]) / size[0])) - y
    if abs(padx) > 10 or abs(pady) > 10:
        padx = 10000 if padx < 0 else padx
        pady = 10000 if pady < 0 else pady
        if padx < pady:
            new_shape = [x + padx, y, max(c, 1)]
            padded_image = np.zeros(new_shape, dtype=data_type)
            padded_image[int(padx // 2):-int(padx - padx // 2), :] = image
        else:
            new_shape = [x, y + pady, max(c, 1)]
            padded_image = np.zeros(new_shape, dtype=data_type)
            padded_image[:, int(pady // 2):-int(pady - pady // 2)] = image
    else:
        padded_image = image
    padded_image = cv2.resize(padded_image, (size[1], size[0])).astype(data_type)
    if c == 1:
        # cv2.resize auto-squeezes images
        return padded_image[:, :, np.newaxis]
    return padded_image


def horizontal_flip(image):
    """This is for convenience, but it is ~2e-7s faster to just copy the source and do this in-line"""
    return image[:, ::-1, :]


def vertical_flip(image):
    """This is for convenience, but it is ~1e-7s faster to just copy the source and do this in-line"""
    return image[::-1, ...]


def adjust_gamma(image, gamma=1.0):
    """Adjust the gamma of the image."""
    invGamma = 1.0 / gamma
    table = np.array([((i / 255.0)**invGamma) * 255
                      for i in np.arange(0, 256)]).astype(np.uint8)
    return cv2.LUT(image, table)


def polar_to_cartesian(image, output_shape=None):
    """Convert a square image with a polar object (circle/tube) to cartesian (unroll it).

    NOTE: output_shape uses numpy shape: [rows, columns]
    """
    if output_shape is None:
        output_shape = image.shape

    degrees = np.linspace(0, 360, output_shape[1], endpoint=False, dtype=np.float32)
    degrees = np.radians(degrees)
    center = image.shape[0] // 2
    radius = np.linspace(0, center, output_shape[0], endpoint=False, dtype=np.float32)
    d, r = np.meshgrid(degrees, radius)
    newx = r * np.cos(d) + center
    newy = r * np.sin(d) + center

    output = cv2.remap(image, newx, newy, cv2.INTER_LINEAR)
    return output
