import cv2
import numpy as np
from PIL import Image

# marker choices
COLORS = [(246, 195, 203), (112, 221, 208)]
MARKERS = [1, 5]

def blend_seg(input_image, segmented_image, dot_locations, dot_labels, contour_color=(63, 126, 174)):
    input_image = np.array(input_image)
    segmented_image = np.array(segmented_image)
    
    # Create a mask for the foreground (non-transparent) pixels
    foreground_mask = segmented_image[:, :, 3] > 0

    # Create a mask for the background (transparent) pixels
    background_mask = ~foreground_mask

    # Darken the background pixels
    darkened_background = input_image.copy()
    darkened_background[background_mask] = darkened_background[background_mask] * 0.52  # Adjust the multiplier as needed to control darkness

    # Create an empty mask for the boundary
    boundary_mask = np.zeros_like(segmented_image[:, :, 3], dtype=np.uint8)
    solid_boundary_mask = np.zeros_like(segmented_image[:, :, 3], dtype=np.uint8)

    # Find the contour of the segmented region
    contours, _ = cv2.findContours(
        segmented_image[:, :, 3], cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE
    )

    # Draw the boundary on the boundary mask
    cv2.drawContours(boundary_mask, contours, -1, (255), thickness=6)
    cv2.drawContours(solid_boundary_mask, contours, -1, (255), thickness=2)

    blur_mask = cv2.GaussianBlur(boundary_mask, (0, 0), sigmaX=4)

    # Create a mask for the contour region
    contour_region_mask = (blur_mask > 0)

    # Blend the contour color with the existing pixel colors
    result_image = darkened_background.copy()
    mask_weight = 0.9 * blur_mask[contour_region_mask, None]/255
    result_image[contour_region_mask] = (
        darkened_background[contour_region_mask] * (1-mask_weight) + np.array(contour_color) * mask_weight
    ).astype(np.uint8)

    # Overlay the contour on the result image without blending
    result_image[solid_boundary_mask > 0] = contour_color  # Set contour pixels to blue

    # Draw dots at the specified locations
    dot_radius = 6
    for location, label in zip(dot_locations, dot_labels):
        if label:
            cv2.circle(result_image, location, dot_radius, COLORS[label], -1)
        else:
            cv2.drawMarker(result_image, location, COLORS[label], markerType=MARKERS[label],
                        markerSize=6, thickness=3)

    return Image.fromarray(result_image)


def blend_seg_pure(input_image, segmented_image, dot_locations, dot_labels):
    input_image = np.array(input_image)
    segmented_image = np.array(segmented_image)
    
    # Create a mask for the foreground (non-transparent) pixels
    foreground_mask = segmented_image[:, :, 3] > 0

    # Blend the foreground
    red_foreground = input_image.copy()
    blend_weight = 0.8
    red_foreground[foreground_mask] = red_foreground[foreground_mask] *(1-blend_weight) + np.array((255,0,0)) * blend_weight

    result_image = red_foreground

    # Draw dots at the specified locations
    dot_radius = 6
    for location, label in zip(dot_locations, dot_labels):
        if label:
            cv2.circle(result_image, location, dot_radius, COLORS[label], -1)
        else:
            cv2.drawMarker(result_image, location, COLORS[label], markerType=MARKERS[label],
                        markerSize=6, thickness=3)

    return Image.fromarray(result_image)