import numpy as np
from PIL import Image

def create_cylindrical_depth_map(image_path, output_path):
    # Load image
    img = Image.open(image_path).convert("RGBA")
    width, height = img.size
    
    # Create a numpy array for the depth map
    # We want a gradient that is white in the center horizontal axis and black at the edges (cylindrical)
    # 0 = far (black), 255 = near (white)
    
    # Generate X coordinates (0 to width)
    x = np.linspace(-1, 1, width)
    # Compute cylindrical depth: sqrt(1 - x^2) for a perfect cylinder, or just a cosine/parabolic falloff
    # Let's use cosine for smooth roundness: cos(x * pi / 2)
    depth_profile = np.cos(x * np.pi / 2) # Center (0) is 1, Edges (-1, 1) are 0
    
    # Normalize to 0-255
    depth_profile = (depth_profile * 255).astype(np.uint8)
    
    # Tile vertically to create the full map
    depth_map = np.tile(depth_profile, (height, 1))
    
    # Create Image from array
    depth_img = Image.fromarray(depth_map, mode='L')
    
    # MASKING: We only want the vase to have depth, the background (transparent) should be flat/far.
    # Use the alpha channel of the original image as a mask
    alpha = np.array(img.split()[-1])
    
    # Where alpha is 0 (background), set depth to 0 (flat/far)
    depth_array = np.array(depth_img)
    depth_array[alpha < 10] = 0 # Threshold for transparency
    
    # Save
    final_depth = Image.fromarray(depth_array)
    final_depth.save(output_path)
    print(f"Generated depth map: {output_path}")

if __name__ == "__main__":
    create_cylindrical_depth_map("pottery-vase.png", "pottery-vase_depth.png")