import cv2
import numpy as np
import os

def process_hay():
    # Paths
    input_path = 'assets/references/proizvodnja_sena.png'
    output_dir = 'assets/DEMO_FAZA1/Items/Hay'
    
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    # Load image
    img = cv2.imread(input_path)
    if img is None:
        print(f"Error: Could not load {input_path}")
        return

    # Convert to RGBA
    img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)

    # Define Green Key (assuming bright green background)
    # Range for Chroma Key Green in HSV is best, but BGR is fine for simple flat green
    # Pure Green in BGR is (0, 255, 0)
    
    # Using HSV for better masking
    hsv = cv2.cvtColor(img[:,:,:3], cv2.COLOR_BGR2HSV)
    
    # Mask for Green (Hue 35-85 approx for standard green screen)
    lower_green = np.array([35, 50, 50])
    upper_green = np.array([85, 255, 255])
    
    mask = cv2.inRange(hsv, lower_green, upper_green)
    
    # Invert mask (we want non-green parts)
    mask_inv = cv2.bitwise_not(mask)
    
    # Set Alpha channel based on mask
    # Everything green becomes transparent
    img[:, :, 3] = mask_inv

    # Find contours of objects
    # Use the mask_inv (where 255 is the object)
    contours, _ = cv2.findContours(mask_inv, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    print(f"Found {len(contours)} objects.")

    count = 0
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        
        # Filter noise
        if w < 10 or h < 10:
            continue
            
        # Crop
        crop = img[y:y+h, x:x+w]
        
        # Determine if it's a small pile or big bale based on size
        # This is a heuristic. We'll label them by size.
        area = w * h
        label = "hay_piece"
        
        # Save
        output_filename = f"{output_dir}/hay_drop_{count}.png"
        cv2.imwrite(output_filename, crop)
        print(f"Saved {output_filename} (Size: {w}x{h})")
        count += 1

if __name__ == "__main__":
    process_hay()