from PIL import Image, ImageSequence, ImageDraw
import numpy as np
import cv2

from mediaug.utils import convert_array_to_poly
from mediaug.variables import COLOR_CYTO_MASK, COLOR_NUC_MASK


def np_to_pil(img):
	""" Converts a PIL format image to numpy
    Args:
	  new_img (np.array): converted image
	Returns:
	  img (PIL.Image): input image
	"""
	return Image.fromarray(img)


def pil_to_np(img):
	""" Converts a PIL format image to numpy
	Args:
	  img (PIL.Image): input image
    Returns:
	  new_img (np.array): converted image
	"""
	return np.array(img)


def read_tiff(path):
    """ Reads an image with a .tff extension, used in the Unet example.
    Args:
      path (str): The path of the image
    Returns:
      ans (np.array): Numpy array of the image values
    """
    return np.array([np.array(p) for p in ImageSequence.Iterator(Image.open(path))])


def read_bmp(img_path):
    """ Reads an image with a .bmp extension Returns np.array 
    Args:
      path (str): The path of the image
    Returns:
      ans (np.array): Numpy array of the image values
    """
    return cv2.imread(img_path)


def read_png(img_path):
    """ Reads an image with a .png extension. Returns np.array
    Args:
      path (str): The path of the image
    Returns:
      ans (np.array): Numpy array of the image values
    """
    return cv2.imread(img_path)


def read_dat_file(path):
    """ Reads an .dat file with polygons, the data from SIPaKMeD. Returns np.array
    Args:
      path (str): The path of the .dat file
    Returns:
      ans (np.array): The [n,2] array of a polygon
    """
    return np.loadtxt(path, delimiter=',')


def read_img(img_path):
    """ Reads an image
    Args:
      path (str): The path of the image
    Returns:
      ans (np.array): Numpy array of the image values
    """
    return cv2.imread(img_path)


def save_img(img, path):
    """ Save an img to given path
    Args:
      img (np.array): The image numpy array
      path (str): The path to save the image to
    """
    cv2.imwrite(path, img)
    return path


def rotate(image, angle):
    """ Rotates an image by angle in degrees, increases 
    the dimension of the imageas necessary
    Args:
      image (np.array): Image array
      angle (int): Degree of rotation clockwise in degrees
    Returns:
      rotated (np.array): The new rotated image array
    """
    (h, w) = image.shape[:2]
    (cX, cY) = (w // 2, h // 2)
 
    # grab the rotation matrix (applying the negative of the
    # angle to rotate clockwise), then grab the sine and cosine
    # (i.e., the rotation components of the matrix)
    M = cv2.getRotationMatrix2D((cX, cY), -angle, 1.0)
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])
 
    # compute the new bounding dimensions of the image
    nW = int((h * sin) + (w * cos))
    nH = int((h * cos) + (w * sin))
 
    # adjust the rotation matrix to take into account translation
    M[0, 2] += (nW / 2) - cX
    M[1, 2] += (nH / 2) - cY
    return cv2.warpAffine(image, M, (nW, nH))


def soften_mask(mask, amount=5):
    """ Softens the edges of a mask by dialating it and then Gaussian blurring
    Args:
      mask (np.array): The mask array
      amount (int): The number of times to apply the dilation, should be arround 5
    Returns:
      ans (np.array): The transformed mask
    """
    kernel = np.ones((5,5), np.uint8) 
    mask_dilation = cv2.dilate(mask, kernel, iterations=amount)
    blur = cv2.GaussianBlur(mask_dilation, (5,5), 0)
    return cv2.max(mask, blur)


def get_blank_mask(img, greyscale=False):
    """ Return a blank mask the same size as the given image
    Args:
      img (np.array): The input image array
      greyscale (Boolean): Make it a single channel mask
    Returns:
      mask (np.array): Black mask the size of img
    """
    if greyscale:
        return np.zeros(img.shape[:2], np.uint8)
    return np.zeros(img.shape,np.uint8)


def image_on_image_alpha(bg, fg, fg_mask, center):
    """ Place an image on an image with a backround mask
    Blends them using the alpha mask.
    Args:
      bg (np.array): Background image array
      fg (np.array): Foreground image array
      fg_mask (np.array): Foreground alpha mask array
      center (tuple[int,int]): The postion on bg to place fg
    Returns:
      ans (np.array): The merged image
    """
    alpha = np.zeros(bg.shape[:2], dtype=np.uint8)
    alpha = place_img_on_img(alpha, fg_mask, center)
    cell_img = place_img_on_img(bg.copy(), fg, center)

    alpha = alpha.astype(float)/255
    alpha = np.repeat(alpha[:, :, np.newaxis], 3, axis=2)

    fg = cv2.multiply(alpha, cell_img.astype(float))
    bg = cv2.multiply(1.0 - alpha, bg.astype(float))

    return cv2.add(fg, bg).astype(np.uint8)


def place_img_on_img(bg, fg, center):
    """ Place an image on top of another image
    Args:
      bg (np.array): Backgourn image
      fg (np.array): Foreground image
      center (x, y): Where to put CENTER of fg image
    Returns:
      ans (np.array): The new image
    """
    ch, cw = center
    fg_h, fg_w = fg.shape[:2]
    bg_h, bg_w = bg.shape[:2]

    # check offset in bg
    if cw < 0 or cw > bg_w or ch < 0 or ch > bg_h:
        raise ValueError('Center not in backgound bounds')

    # find top left corner of fg in respect to bg
    left_w = cw - (fg_w // 2)
    left_h = ch - (fg_h // 2)
    end_w = left_w + fg_w
    end_h = left_h + fg_h

    # for if goes over bg boundries
    abs_left_w = max(left_w, 0)
    abs_left_h = max(left_h, 0)
    abs_end_w = min(end_w, bg_w)
    abs_end_h = min(end_h, bg_h)
    
    # for fg boundries
    fg_left_w = abs(min(left_w, 0))
    fg_left_h = abs(min(left_h, 0))
    diff_w = bg_w - end_w
    diff_h = bg_h - end_h
    if diff_w >= 0:
        fg_end_w = fg_w
    else:
        fg_end_w = diff_w
    if diff_h >= 0:
        fg_end_h = fg_h
    else:
        fg_end_h = diff_h

    ans = np.copy(bg)
    ans[abs_left_h:abs_end_h, abs_left_w:abs_end_w] = fg[fg_left_h:fg_end_h, fg_left_w:fg_end_w]
    return ans


def generate_cell_mask(img, cyto, nuc):
    """ Generate a mask for a labelled cell
    Args:
      img (np.array): The image array
      cyto (np.array): A [n,2] numpy array representing the polygon for the cytoplasm of a cell
      nuc (np.array): A [n,2] numpy array representing the polygon for the nucleus of a cell
    Returns:
      mask (np.array): The mask of parts of a cell
    """
    w, h, _ = img.shape
    mask = Image.new(mode="RGB", size=(h, w))
    ImageDraw.Draw(mask).polygon(convert_array_to_poly(cyto), outline=None, fill=COLOR_CYTO_MASK)
    ImageDraw.Draw(mask).polygon(convert_array_to_poly(nuc), outline=None, fill=COLOR_NUC_MASK)
    return np.array(mask)


def generate_cell_mask_list(img, cytos, nucs):
    """ Generate a masks for a list of labelled cell in slide
    Args:
      img (np.array): The image array
      cytos (list[np.array]): List of[n,2] numpy array representing the polygon for the cytoplasm of a cell
      nucs (list[np.array]): List of [n,2] numpy array representing the polygon for the nucleus of a cell
    Returns:
      mask (np.array): The mask of cells in slide
    """
    h, w, _ = img.shape
    mask = Image.new(mode="RGB", size=(w, h))
    for poly in cytos:
        ImageDraw.Draw(mask).polygon(convert_array_to_poly(poly), outline=None, fill=COLOR_CYTO_MASK)
    for poly in nucs:
        ImageDraw.Draw(mask).polygon(convert_array_to_poly(poly), outline=None, fill=COLOR_NUC_MASK)
    return np.array(mask)


def is_greyscale(img):
    """ Checks if an image array is greyscale
    Args:
      img (np.array): The image array
    Returns:
      is_greyscaled (Boolean): Is the image a greyscale image
    """
    if len(img.shape) < 3:
        return True
    return False