import glob
import numpy as np
import cv2
from PIL import Image
from utils import config


def write(filename, content, class_num=2, color_map=True):
    """
    Save image array to a specified path.
    The image will be automatically recolored via the class number.
    :param filename: The specified path.
    :param content: Numpy array containing the image.
    :param class_num: Total class number.
    :param color_map: Whether change the probability into gray grade.
    """
    if class_num <= 1:
        raise Exception('ERROR: Class number should be >= 2.')
    color_stage = 255. / (class_num - 1) if color_map else 1.0
    new_image = Image.fromarray(np.uint8(content * color_stage))
    new_image.save(filename, "PNG")


def generate_effective_regions(size):
    """
    This function is used to generate effective regions for inference according to the given slide size.
    :param size: Given slide size, should be in the form of [w, h].
    """
    width = size[0]
    height = size[1]
    x_step = int(width / config.CENTER_SIZE)
    y_step = int(height / config.CENTER_SIZE)
    regions = []
    for x in range(0, x_step):
        for y in range(0, y_step):
            regions.append([x * config.CENTER_SIZE, y * config.CENTER_SIZE, 0, 0,
                            config.CENTER_SIZE - 1, config.CENTER_SIZE - 1])
    if not height % config.CENTER_SIZE == 0:
        for x in range(0, x_step):
            regions.append([x * config.CENTER_SIZE, height - config.CENTER_SIZE,
                                0, (y_step + 1) * config.CENTER_SIZE - height,
                            config.CENTER_SIZE - 1, config.CENTER_SIZE - 1])
    if not width % config.CENTER_SIZE == 0:
        for y in range(0, y_step):
            regions.append([width - config.CENTER_SIZE, y * config.CENTER_SIZE,
                                (x_step + 1) * config.CENTER_SIZE - width, 0,
                            config.CENTER_SIZE - 1, config.CENTER_SIZE - 1])
    if not (height % config.CENTER_SIZE == 0 or width % config.CENTER_SIZE == 0):
        regions.append([width - config.CENTER_SIZE, height - config.CENTER_SIZE,
                                (x_step + 1) * config.CENTER_SIZE - width, (y_step + 1) * config.CENTER_SIZE - height,
                        config.CENTER_SIZE - 1, config.CENTER_SIZE - 1])
    return regions


def generate_overlap_tile(region, dimensions):
    """
    This function is used to process border patches.
    """
    shifted_region_x = region[0] - config.BORDER_SIZE
    shifted_region_y = region[1] - config.BORDER_SIZE
    clip_region_x = config.BORDER_SIZE
    clip_region_y = config.BORDER_SIZE
    if region[0] == 0:
        shifted_region_x = shifted_region_x + config.BORDER_SIZE
        clip_region_x = 0
    if region[1] == 0:
        shifted_region_y = shifted_region_y + config.BORDER_SIZE
        clip_region_y = 0
    if region[0] == dimensions[0] - config.CENTER_SIZE:
        shifted_region_x = shifted_region_x - config.BORDER_SIZE
        clip_region_x = 2 * config.BORDER_SIZE
    if region[1] == dimensions[1] - config.CENTER_SIZE:
        shifted_region_y = shifted_region_y - config.BORDER_SIZE
        clip_region_y = 2 * config.BORDER_SIZE
    return [shifted_region_x, shifted_region_y], [clip_region_x, clip_region_y]


def image_to_array(input_image):
    """
    Loads image into numpy array.
    """
    im_array = np.array(input_image.getdata(), dtype=np.uint8)
    im_array = im_array.reshape((input_image.size[0], input_image.size[1]))
    return im_array


def post_processing(image_patch):
    """
    Remove small noisy points.
    """
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (config.FILTER_KERNEL, config.FILTER_KERNEL))
    open_patch = cv2.morphologyEx(image_patch, cv2.MORPH_OPEN, kernel)
    close_patch = cv2.morphologyEx(open_patch, cv2.MORPH_CLOSE, kernel)
    return close_patch


def concat_patches(temp_dir, image_name):
    """
    Concatenate the predicted patches into a thumbnail result.
    """
    prediction_list = glob.glob(temp_dir + image_name + '*_prediction.png')
    patch_list = []
    for prediction_image in prediction_list:
        name_parts = prediction_image.split('/')[-1].split('_')
        pos_x, pos_y = int(name_parts[-3]), int(name_parts[-2])
        patch_list.append([pos_x, pos_y])
    image_patches = []
    patch_list.sort()
    last_x = -1
    row_patch = []
    for position in patch_list:
        pos_x = position[0]
        pos_y = position[1]
        image = Image.open(temp_dir + '_'.join([image_name, str(pos_x), str(pos_y), 'prediction']) + '.png')
        original_width, original_height = image.size
        if original_width < config.THUMBNAIL_RATIO or original_height < config.THUMBNAIL_RATIO:
            continue
        image = image.resize(
            (int(original_width / config.THUMBNAIL_RATIO),
             int(original_height / config.THUMBNAIL_RATIO)), Image.NEAREST)
        image_patch = image_to_array(image)
        if not pos_x == last_x:
            last_x = pos_x
            if len(row_patch) == 0:
                row_patch = image_patch
            else:
                if not len(image_patches) == 0:
                    image_patches = np.column_stack((image_patches, row_patch))
                else:
                    image_patches = row_patch
                row_patch = image_patch
        else:
            row_patch = np.row_stack((row_patch, image_patch))
    prediction = np.column_stack((image_patches, row_patch))
    return prediction