# should work for all apporaches (Deep learning or not)
# should work on segmentation masks
# should include a smoothing and an max-area threshold
# should try to work with the tortina circle (or maybe better estimate it again)
# maybe even to apply before thresholding, directly on anomaly maps/other outputs to have better smoothing
# output should be again a segmentation mask

import numpy as np
import cv2





def get_points_in_circle(circle):
    x0, y0, radius = circle.astype(np.int32)
    x_ = np.arange(x0 - radius - 1, x0 + radius + 1, dtype=np.int32)
    y_ = np.arange(y0 - radius - 1, y0 + radius + 1, dtype=np.int32)
    x, y = np.where((x_[:,np.newaxis] - x0)**2 + (y_ - y0)**2 <= radius**2)
    # for x, y in zip(x_[x], y_[y]):    #     yield x, y   
    return (x_[x], y_[y])




def find_single_tortina_circle(image):
    height, width, num_channels = image.shape
    
    #find the biggest circle (tortina) in the image
    #init with default circle
    selected_circle = np.array((width//2, height//2, int(height/2*0.9)))
    try:
        #circles = find_tortinas(image, 1, force_num_tortinas=True)
        height, width = image.shape[:2]
        if (image.ndim == 2) or (image.shape[-1] == 1):
            gray = image
        else:
            gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
        blur = cv2.blur(gray, (7,7))
        circles = cv2.HoughCircles(blur, cv2.HOUGH_GRADIENT, 3, minDist=width//2)
        if circles is None:
            circles = []
        else:
            circles = circles[0].astype(np.uint32)

    except:
        circles = []
    
    if len(circles) > 0:
        max_r = 0
        for x,y,r in circles:
            if r > max_r:
                max_r = r
                selected_circle = np.array((x,y,r))
    
    return selected_circle



def postprocessing(image, segmask, fat_bloom_id=1):
    
    if segmask.shape != image.shape[:2]:
        raise ValueError(
            """segmask argument should represent a segmentation mask with 2 dimensions (height, width)!
            This means that its values should already be thresholded and (in case of rgb) reduced to
            a single channel.
            """)
    
    height, width = image.shape[:2]
    selected_circle = find_single_tortina_circle(image)

    #debugging
    #plot_image_with_circle(image, selected_circle)

    # smooth the segmask
    binary = (segmask == fat_bloom_id).astype(np.uint8)
    kernel = np.ones((5,5),np.uint8)
    closing = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel)
    opening = cv2.morphologyEx(closing, cv2.MORPH_OPEN, kernel)

    #the opening values should overwrite the fat_bloom of the input
    #segmask. Where we had fat bloom in the input but not anymore in the
    #opening we need to determine whether it is backgorund or tortina. We
    #can do so using the estimated circle.

    #initialize everything to background. background = 0
    new_segmask = np.zeros_like(segmask)
    tortina_indices = get_points_in_circle(selected_circle)
    #correction
    out_of_width_0 = (tortina_indices[0] < 0)
    out_of_width_1 = (tortina_indices[0] >= width)
    out_of_height_0 = (tortina_indices[1] < 0)
    out_of_height_1 = (tortina_indices[1] >= height)
    tortina_indices[0][out_of_width_0] = 0
    tortina_indices[0][out_of_width_1] = width - 1
    tortina_indices[1][out_of_height_0] = 0
    tortina_indices[1][out_of_height_1] = height - 1
    # remove_indices = (tortina_indices[0] < 0) | (tortina_indices[0] >= height) | \
    #     (tortina_indices[1] < 0) | (tortina_indices[1] >= width)
    # keep_indices = np.setdiff1d(np.arange(len(tortina_indices[0])), remove_indices)
    # tortina_indices = (tortina_indices[0][keep_indices], tortina_indices[1][keep_indices])
    

    is_tortina = np.zeros_like(segmask, dtype=bool)
    is_tortina[tortina_indices] = True
    #tortina = 1
    new_segmask[is_tortina] = 1
    #fat_bloom = 2
    new_segmask[(opening == 1) & is_tortina] = 2

    return new_segmask


def final_prediction(anomaly_map, segmask):
    
    fat_bloom_area = np.count_nonzero(segmask == 2)
    tortina_area = np.count_nonzero(segmask == 1) + fat_bloom_area

    relative_area = fat_bloom_area / tortina_area
    
    area = 0
    if (relative_area > 0) and (relative_area <= 0.25):
        area = 1
    elif (relative_area > 0.25) and (relative_area <= 0.5):
        area = 2
    elif (relative_area > 0.5) and (relative_area <= 0.75):
        area = 3
    elif relative_area > 0.75:
        area = 4
    
    relative_intensity = 0.0
    if fat_bloom_area > 0.0:
        relative_intensity = anomaly_map[segmask==2].mean()

    #TODO: intensity should be depending on colour of underlying tortina
    # i.e. for a darker tortina intesity is automatically higher,
    # but that should be relativiert.
    intensity = 0
    if (relative_intensity > 0) and (relative_intensity <= 0.25):
        intensity = 1
    elif (relative_intensity > 0.25) and (relative_intensity <= 0.5):
        intensity = 2
    elif (relative_intensity > 0.5) and (relative_intensity <= 0.75):
        intensity = 3
    elif relative_intensity > 0.75:
        intensity = 4

    return area, intensity