import cv2
import numpy as np
from PIL import Image

def segment_image_by_depth(original_img_path, depth_map_path, num_segments=3):
    # Load original image (JPG) and depth map
    original_img = Image.open(original_img_path).convert("RGB")  # Convert to RGB
    depth_map = cv2.imread(depth_map_path, cv2.IMREAD_GRAYSCALE)
    
    # Normalize depth map
    depth_map = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)

    # Calculate depth ranges based on the number of segments
    segment_ranges = np.linspace(0, 255, num_segments + 1).astype(int)

    segmented_images = []

    for i in range(num_segments):
        # Generate mask for the current depth range
        lower, upper = int(segment_ranges[i]), int(segment_ranges[i + 1])
        mask = cv2.inRange(depth_map, lower, upper)
        
        # Convert mask to 3 channels to match the original image
        mask_rgb = cv2.merge([mask, mask, mask])
        
        # Apply mask to the original image
        segmented_img = cv2.bitwise_and(np.array(original_img), mask_rgb)

        # Convert back to PIL Image and add an alpha channel
        segmented_img_pil = Image.fromarray(segmented_img).convert("RGBA")
        alpha_channel = Image.fromarray(mask).convert("L")  # Use the mask as the alpha channel
        segmented_img_pil.putalpha(alpha_channel)  # Add transparency based on depth mask

        segmented_images.append(segmented_img_pil)

    # Save each segmented image as PNG
    output_paths = []
    for idx, seg_img in enumerate(segmented_images):
        output_path = f"outputs/segment_{idx + 1}.png"
        seg_img.save(output_path)
        output_paths.append(output_path)
    
    return output_paths