import os.path
import sys

import cv2
import numpy as np
import scipy.ndimage as ndi

from deepliif.postprocessing import overlay, refine, remove_cell_noise, remove_background_noise, \
    remove_small_objects_from_image


def align_seg_on_image(input_image, input_mask, output_image, thresh=100, noise_objects_size=100):
    seg_image = cv2.cvtColor(cv2.imread(input_mask), cv2.COLOR_BGR2RGB)
    orig_image = cv2.cvtColor(cv2.imread(input_image), cv2.COLOR_BGR2RGB)

    final_mask = orig_image.copy()
    processed_mask = np.zeros_like(orig_image)

    red = seg_image[:, :, 0]
    blue = seg_image[:, :, 2]
    boundary = seg_image[:, :, 1]

    boundary[boundary < thresh] = 0

    positive_cells = np.zeros((seg_image.shape[0], seg_image.shape[1]), dtype=np.uint8)
    negative_cells = np.zeros((seg_image.shape[0], seg_image.shape[1]), dtype=np.uint8)

    positive_cells[red > thresh] = 255
    positive_cells[boundary > thresh] = 0
    negative_cells[blue > thresh] = 255
    negative_cells[boundary > thresh] = 0

    negative_cells[red >= blue] = 0
    positive_cells[blue > red] = 0

    positive_cells = cv2.morphologyEx(positive_cells, cv2.MORPH_DILATE, kernel=np.ones((2, 2)))
    negative_cells = cv2.morphologyEx(negative_cells, cv2.MORPH_DILATE, kernel=np.ones((2, 2)))

    negative_cells = remove_background_noise(negative_cells, boundary)
    positive_cells = remove_background_noise(positive_cells, boundary)

    negative_cells, positive_cells = remove_cell_noise(negative_cells, positive_cells)
    positive_cells, negative_cells = remove_cell_noise(positive_cells, negative_cells)

    negative_cells = remove_small_objects_from_image(negative_cells, noise_objects_size)
    negative_cells = ndi.binary_fill_holes(negative_cells).astype(np.uint8)

    positive_cells = remove_small_objects_from_image(positive_cells, noise_objects_size)
    positive_cells = ndi.binary_fill_holes(positive_cells).astype(np.uint8)

    contours, hierarchy = cv2.findContours(positive_cells,
                                           cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    cv2.drawContours(final_mask, contours, -1, (255, 0, 0), 2)

    contours, hierarchy = cv2.findContours(negative_cells,
                                           cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    cv2.drawContours(final_mask, contours, -1, (0, 0, 255), 2)

    processed_mask[positive_cells > 0] = (0, 0, 255)
    processed_mask[negative_cells > 0] = (255, 0, 0)

    contours, hierarchy = cv2.findContours(positive_cells,
                                           cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    cv2.drawContours(processed_mask, contours, -1, (0, 255, 0), 2)

    contours, hierarchy = cv2.findContours(negative_cells,
                                           cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    cv2.drawContours(processed_mask, contours, -1, (0, 255, 0), 2)

    cv2.imwrite(output_image, cv2.cvtColor(final_mask, cv2.COLOR_BGR2RGB))
    cv2.imwrite(output_image.replace('Overlaid', 'Refined'), processed_mask)


def align_seg_on_image2(input_image, input_mask, output_image, thresh=100, noise_objects_size=20):
    seg_image = cv2.cvtColor(cv2.imread(input_mask), cv2.COLOR_BGR2RGB)
    orig_image = cv2.cvtColor(cv2.imread(input_image), cv2.COLOR_BGR2RGB)

    overlaid_mask = overlay(orig_image, seg_image, thresh, noise_objects_size)
    cv2.imwrite(output_image, overlaid_mask)

    refined_mask = refine(orig_image, seg_image, thresh, noise_objects_size)
    cv2.imwrite(output_image.replace('Overlaid', 'Refined'), refined_mask)


def post_process_segmentation_mask(input_dir, seg_thresh=100, noise_object_size=100):
    images = os.listdir(input_dir)
    image_extensions = ['.png', '.jpg', '.tif', '.tiff']
    for img in images:
        if '_fake_B_5.png' in img:
            align_seg_on_image2(os.path.join(input_dir, img.replace('_fake_B_5', '_real_A')),
                                os.path.join(input_dir, img),
                                os.path.join(input_dir, img.replace('_fake_B_5', '_Seg_Overlaid_')),
                                thresh=seg_thresh, noise_objects_size=noise_object_size)
        elif '_Seg.png' in img:
            orig_img_ext = '.png'
            for img_ext in image_extensions:
                if os.path.exists(os.path.join(input_dir, img.replace('_Seg.png', img_ext))):
                    orig_img_ext = img_ext
                    break
            align_seg_on_image2(os.path.join(input_dir, img.replace('_Seg.png', orig_img_ext)),
                                os.path.join(input_dir, img),
                                os.path.join(input_dir, img.replace('_Seg', '_SegOverlaid')),
                                thresh=seg_thresh, noise_objects_size=noise_object_size)


if __name__ == '__main__':
    base_dir = sys.argv[1]
    segmentation_thresh = 100
    noise_obj_size = 20
    if len(sys.argv) > 2:
        segmentation_thresh = int(sys.argv[2])
    if len(sys.argv) > 3:
        noise_obj_size = int(sys.argv[3])

    post_process_segmentation_mask(base_dir, segmentation_thresh, noise_obj_size)
