import cv2
import numpy as np
import os

def bury_pipe_stream():
    src_path = '/Users/davidkotnik/.gemini/antigravity/brain/07019d04-a214-43ab-9565-86f4e8f17e5b/uploaded_media_1769607894587.jpg'
    
    print(f"Loading {src_path}")
    img = cv2.imread(src_path)
    if img is None: return

    # 1. GrabCut to isolate object from BG
    mask = np.zeros(img.shape[:2], np.uint8)
    bgdModel = np.zeros((1,65),np.float64)
    fgdModel = np.zeros((1,65),np.float64)
    h, w = img.shape[:2]
    cv2.grabCut(img, mask, (10, 10, w-20, h-20), bgdModel, fgdModel, 5, cv2.GC_INIT_WITH_RECT)
    mask_final = np.where((mask==2)|(mask==0),0,1).astype('uint8') 
    
    # 2. REMOVE OUTER WALLS (The "Burying" Step)
    # The walls are the darkest parts of the ground structure.
    # The "Top Rim" (ground level) is lighter brown.
    # The water is murky brown/gray.
    
    # Convert to HSV to separate by brightness/color
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    
    # Define "Wall" Color range (Dark Brown / Black shadow)
    # H: Any (Browns are 0-20 or 160-180)
    # S: Low to Medium
    # V: VERY LOW (Shadows)
    
    # Let's visualize: brightness < 40?
    v_channel = hsv[:,:,2]
    
    # Create a mask of "Very Dark" pixels
    wall_mask = (v_channel < 60)
    
    # But we want to keep outlines (which are thin dark lines).
    # Walls are LARGE blocks of dark.
    # We can use morphological opening to select only LARGE dark areas (walls) and ignore thin lines (outlines).
    wall_mask_uint8 = wall_mask.astype(np.uint8)
    kernel = np.ones((5,5), np.uint8)
    # Erode then Dilate -> Removes small noise (lines), keeps big blobs (walls)
    walls_only = cv2.morphologyEx(wall_mask_uint8, cv2.MORPH_OPEN, kernel, iterations=2)
    
    # Now, subtract 'walls_only' from our main mask.
    # We want mask_final to be 0 where walls_only is 1.
    mask_buried = mask_final.copy()
    mask_buried[walls_only == 1] = 0
    
    # 3. FIX: This might delete the interior of the pipe (shadow) or the drain hole!
    # We must PROTECT the pipe/drain area.
    # The Drain is at the top right. Let's define a Region of Interest (ROI) that we DO NOT touch.
    # Pipe is roughly in top right quadrant.
    # Heuristic: Only apply wall removal to the bottom half/rim?
    # Or better: The walls are on the OUTSIDE boundary.
    
    # Let's perform the subtraction.
    
    # 4. FIX PIPE HOLE (Make opaque again if GrabCut lost it)
    # (Re-using logic from previous step if needed, but GrabCut usually keeps it).
    
    # Combine
    b, g, r = cv2.split(img)
    alpha = mask_buried * 255
    
    img_rgba = cv2.merge([b, g, r, alpha])
    
    # Crop
    coords = cv2.findNonZero(alpha)
    if coords is not None:
        x, y, cw, ch = cv2.boundingRect(coords)
        img_rgba = img_rgba[y:y+ch, x:x+cw]
        
    targets = [
        '/Users/davidkotnik/repos/novafarma/main/assets/stream_pipe.png',
        '/Users/davidkotnik/repos/novafarma/assets/DEMO_FAZA1/Environment/stream_pipe.png'
    ]
    for t in targets:
        cv2.imwrite(t, img_rgba)
        print(f"Saved {t}")

if __name__ == "__main__":
    bury_pipe_stream()