import os
import numpy as np
from PIL import Image

def remove_checkerboard(image_path, sensitivity=20):
    try:
        if not os.path.exists(image_path):
            print(f"File not found: {image_path}")
            return

        print(f"Processing: {image_path}")
        img = Image.open(image_path).convert("RGBA")
        data = np.array(img)
        
        # Get dimensions
        height, width = data.shape[:2]
        
        # 1. SAMPLE BACKGROUND COLORS FROM CORNERS
        # We assume corners are background. We'll take a 10x10 patch from top-left.
        corner_patch = data[0:10, 0:10] 
        # Find unique colors in the corner (likely the two checkerboard colors)
        unique_colors = np.unique(corner_patch.reshape(-1, 4), axis=0)
        
        # We also manually add standard checkerboard colors just in case corners are weird
        # Light grey, Dark grey, White
        standard_bg_colors = [
            [255, 255, 255], # White
            [204, 204, 204], # Light Grey
            [238, 238, 238], # Lighter Grey
            [192, 192, 192], # Classic Grey
            [128, 128, 128], # Darker Grey
            [0, 0, 0]        # Black (if user mentioned black earlier)
        ]
        
        # Create a boolean mask for transparency (Starts as False = Keep)
        to_remove = np.zeros((height, width), dtype=bool)
        
        r, g, b = data[:,:,0], data[:,:,1], data[:,:,2]
        
        # Helper to match color with tolerance
        def match_color(target_rgb, tol):
            return (np.abs(r - target_rgb[0]) < tol) & \
                   (np.abs(g - target_rgb[1]) < tol) & \
                   (np.abs(b - target_rgb[2]) < tol)

        # Mark pixels matching sampled corner colors
        for color in unique_colors:
            # Ignore if alpha is already 0
            if color[3] == 0: continue
            to_remove |= match_color(color[:3], sensitivity)
            
        # Mark pixels matching standard checkerboard colors
        for bg_c in standard_bg_colors:
            to_remove |= match_color(bg_c, sensitivity)

        # APPLY MASK
        # Set alpha to 0 where mask is True
        data[to_remove, 3] = 0
        
        # SPECIAL CASE: AMNESIA MASK
        # If it's the mask/vignette, we might need to inverted logic or specific handling?
        # User said "through the holes in the rusty frame".
        # If the frame has a solid checkerboard INSIDE the frame, we need to nuke that too.
        # The logic above replaces ALL instances of those colors. 
        # This is risky if the object itself uses grey/white. 
        # BUT for "Rusty" (brown/orange) and "Health Gauge" (brown/red), grey/white is safe to remove.
        # "Start Button" has white text... "match_color" might kill the text!
        
        # REFINEMENT: MASKING
        # We only remove if it touches the "outside" or is part of the large background blocks.
        # But implementing full floodfill in numpy is hard.
        # Let's hope the "Rusty" colors distinct enough from "Pure Grey/White".
        
        # Save
        result = Image.fromarray(data)
        result.save(image_path)
        print(f"Saved cleaned: {image_path}")

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

# PATHS TO FIX
files_to_fix = [
    "/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",
     # Also cleaning the ghosts as requested
    "/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 f in files_to_fix:
    remove_checkerboard(f, sensitivity=25)