# PortraitBlurrer.py
import cv2
import numpy as np
from PIL import Image

class PortraitBlurrer:
    def __init__(self, max_blur=31, depth_threshold=120,
                 feather_strength=3, sharpen_strength=1):
        self.max_blur = max_blur
        # Ensure max_blur is odd and positive
        if self.max_blur % 2 == 0:
            self.max_blur += 1
        if self.max_blur <= 0:
            self.max_blur = 3 # Default odd positive

        self.depth_threshold = depth_threshold
        self.feather_strength = feather_strength
        self.sharpen_strength = sharpen_strength

    def refine_depth_map(self, depth_map):
        # Apply a bilateral filter to smooth depth while preserving edges
        refined_depth = cv2.bilateralFilter(depth_map, 9, 75, 75)
        return refined_depth

    def create_subject_mask(self, depth_map):
        _, mask = cv2.threshold(depth_map, self.depth_threshold, 255, cv2.THRESH_BINARY)
        kernel = np.ones((5, 5), np.uint8)
        mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)

        ksize = self.feather_strength
        if ksize % 2 == 0:
            ksize += 1
        if ksize <= 0:
            ksize = 3

        mask = cv2.GaussianBlur(mask, (ksize, ksize), 0)
        return mask.astype(np.float32) / 255.0

    def sharpen_image(self, image):
        # Ensure sharpen_strength is not zero to avoid division issues later if needed
        strength = max(0.1, self.sharpen_strength) # Prevent zero strength

        # Simple sharpening kernel
        kernel = np.array([[-1, -1, -1],
                           [-1,  9, -1],
                           [-1, -1, -1]])

        # Apply the kernel - adjust strength application if needed
        # A common way is to blend the sharpened with original based on strength
        sharpened = cv2.filter2D(image, -1, kernel)

        # Blend sharpened and original based on strength
        # strength=1 means mostly sharpened, strength close to 0 means mostly original
        if strength != 1.0: # Avoid unnecessary work if strength is 1
           blended = cv2.addWeighted(image, 1.0 - (strength - 1.0) if strength > 1.0 else 1.0 ,
                                     sharpened, strength if strength <= 1.0 else 1.0, 0)
           # Basic clipping if values go out of range due to sharpening
           return np.clip(blended, 0, 255).astype(np.uint8)
        else:
           # Basic clipping if values go out of range due to sharpening
           return np.clip(sharpened, 0, 255).astype(np.uint8)


    def apply_blur(self, original_bgr, depth_map_array):
        # Resize depth map to match image dimensions
        depth_resized = cv2.resize(depth_map_array, (original_bgr.shape[1], original_bgr.shape[0]),
                                   interpolation=cv2.INTER_LINEAR)

        refined_depth = self.refine_depth_map(depth_resized)
        mask = self.create_subject_mask(refined_depth) # Float mask [0, 1]

        blurred = cv2.GaussianBlur(original_bgr, (self.max_blur, self.max_blur), 0)

        # Only sharpen if strength is significant
        if self.sharpen_strength > 0.05: # Threshold to avoid unnecessary computation
             sharpened_original = self.sharpen_image(original_bgr)
             # Blend sharpened subject with original based on mask
             foreground = sharpened_original * mask[:, :, np.newaxis] + \
                          original_bgr * (1 - mask[:, :, np.newaxis])
        else:
             foreground = original_bgr # Use original if no sharpening

        # Blend the (potentially sharpened) foreground with the blurred background
        background = blurred * (1 - mask[:, :, np.newaxis])
        # Combine the foreground (where mask is 1) and background (where mask is 0)
        # Note: Foreground already contains the original where it wasn't sharpened
        # A potentially better blend:
        result = original_bgr * mask[:, :, np.newaxis] + blurred * (1 - mask[:, :, np.newaxis])
        if self.sharpen_strength > 0.05:
            sharpened_subject_only = self.sharpen_image(original_bgr)
            # Apply sharpening only where the mask is high
            result = sharpened_subject_only * mask[:, :, np.newaxis] + result * (1 - mask[:, :, np.newaxis])

        # Ensure result is uint8
        final_result = np.clip(result, 0, 255).astype(np.uint8)

        # Return the final blurred image as a NumPy array (BGR)
        # Also return the refined depth map and the mask for potential display
        return final_result, refined_depth, (mask * 255).astype(np.uint8)


    def process_image(self, original_bgr_np, depth_image_pil):
        depth_map_array = np.array(depth_image_pil)

        if len(depth_map_array.shape) > 2:
            # Assuming input PIL depth map might be RGB, convert to grayscale
            depth_map_array = cv2.cvtColor(depth_map_array, cv2.COLOR_RGB2GRAY)
        elif len(depth_map_array.shape) == 2:
             # Already grayscale, ensure it's uint8 if necessary (though pipeline likely outputs it correctly)
             if depth_map_array.dtype != np.uint8:
                 # Normalize if it's float or other types before potential processing
                 if depth_map_array.max() > 1.0: # Basic check if it might be 0-255
                     depth_map_array = depth_map_array.astype(np.uint8)
                 else: # Assume 0-1 float, scale to 0-255
                      depth_map_array = (depth_map_array * 255).astype(np.uint8)


        # apply_blur now returns the result, depth map, and mask
        blurred_image_np, refined_depth_np, mask_np = self.apply_blur(original_bgr_np, depth_map_array)

        # Return the blurred image, the refined depth map (grayscale), and the mask (grayscale)
        return blurred_image_np, refined_depth_np, mask_np