import os
from PIL import Image
import numpy as np

def clean_checkerboard_and_green(image_path):
    try:
        if not os.path.exists(image_path):
            print(f"Skipping (not found): {image_path}")
            return

        img = Image.open(image_path).convert("RGBA")
        data = np.array(img)

        # RED, GREEN, BLUE, ALPHA
        r, g, b, a = data[:,:,0], data[:,:,1], data[:,:,2], data[:,:,3]

        # 1. Target Checkerboard White (approx 255, 255, 255)
        # Allow slight compression noise
        mask_white = (r > 240) & (g > 240) & (b > 240)

        # 2. Target Checkerboard Grey
        # Usually exact matches like (204, 204, 204) or (192, 192, 192)
        # We target a narrow grey range common in these gens
        mask_grey = (r > 190) & (r < 225) & \
                    (g > 190) & (g < 225) & \
                    (b > 190) & (b < 225) & \
                    (np.abs(r.astype(int) - g.astype(int)) < 15) & \
                    (np.abs(r.astype(int) - b.astype(int)) < 15)

        # 3. Target Green Screen (Pure Green #00FF00)
        # e.g: High Green, Low Red/Blue
        mask_green = (g > 200) & (r < 100) & (b < 100)
        
        # Combine masks
        mask_transparent = mask_white | mask_grey | mask_green

        # Set Alpha to 0 where mask is True
        data[mask_transparent, 3] = 0

        # Save back
        new_img = Image.fromarray(data)
        new_img.save(image_path)
        print(f"cleaned: {image_path}")

    except Exception as e:
        print(f"Error processing {image_path}: {e}")

# PATHS
paths = [
    # UI
    "/Users/davidkotnik/repos/novafarma/assets/slike/NOVE_SLIKE/UI/okvir_zarjavel.png",
    "/Users/davidkotnik/repos/novafarma/assets/slike/NOVE_SLIKE/UI/merilec_zdravja.png",
    "/Users/davidkotnik/repos/novafarma/assets/slike/NOVE_SLIKE/UI/amnezija_maska.png",
    "/Users/davidkotnik/repos/novafarma/assets/slike/NOVE_SLIKE/UI/gumb_start.png",
    
    # CHARACTERS
    "/Users/davidkotnik/repos/novafarma/assets/slike/NOVE_SLIKE/Characters/starsa/Ghost/ghost_otac_cyan.png",
    "/Users/davidkotnik/repos/novafarma/assets/slike/NOVE_SLIKE/Characters/starsa/Ghost/MOJE_SLIKE_KONCNA_ostalo_parents_transparent_ghosts_clean.png"
]

for p in paths:
    clean_checkerboard_and_green(p)