import cv2
import numpy as np

def apply_gaussian_blur(mask, kernel_size=5):
    """Apply Gaussian blur to smooth the mask edges."""
    return cv2.GaussianBlur(mask, (kernel_size, kernel_size), 0)

def apply_threshold(mask, threshold=127):
    """Apply binary threshold to sharpen the mask."""
    _, binary_mask = cv2.threshold(mask, threshold, 255, cv2.THRESH_BINARY)
    return binary_mask

def refine_edges(mask, kernel_size=3):
    """Refine edges using morphological operations."""
    kernel = np.ones((kernel_size, kernel_size), np.uint8)
    eroded = cv2.erode(mask, kernel, iterations=1)
    dilated = cv2.dilate(mask, kernel, iterations=1)
    refined = dilated - eroded
    return cv2.bitwise_or(eroded, refined)

def apply_contour_smoothing(mask):
    """Smooth contours of the mask."""
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    smooth_mask = np.zeros_like(mask)
    for contour in contours:
        epsilon = 0.02 * cv2.arcLength(contour, True)
        approx = cv2.approxPolyDP(contour, epsilon, True)
        cv2.drawContours(smooth_mask, [approx], 0, 255, -1)
    return smooth_mask

def refine_mask(mask, blur_kernel=5, edge_kernel=3, threshold_value=127):
    """Apply a series of refinement operations to the mask."""
    mask = apply_gaussian_blur(mask, blur_kernel)
    mask = apply_threshold(mask, threshold_value)
    mask = refine_edges(mask, edge_kernel)
    mask = apply_contour_smoothing(mask)
    return mask
    

def apply_morphology(mask, kernel_size=3):
    """Apply morphological operations to clean up the mask."""
    kernel = np.ones((kernel_size, kernel_size), np.uint8)
    
    # Opening operation to remove small noise
    mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
    
    # Closing operation to fill small holes
    mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
    
    return mask

def remove_small_objects(mask, min_size=100):
    """Remove small objects from the mask based on area."""
    # Ensure mask is binary and single channel
    if len(mask.shape) > 2:
        mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
    _, binary_mask = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)
    
    num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(binary_mask, connectivity=8)
    
    for i in range(1, num_labels):  # Start from 1 to skip the background
        if stats[i, cv2.CC_STAT_AREA] < min_size:
            binary_mask[labels == i] = 0
    
    return binary_mask

def clean_mask(mask, morph_kernel_size=3, min_object_size=100):
    """Apply both morphological operations and small object removal."""
    mask = apply_morphology(mask, kernel_size=morph_kernel_size)
    mask = remove_small_objects(mask, min_size=min_object_size)
    return mask