#!/usr/bin/env python3
"""
AGGRESSIVE Background Removal - Match Green Screen Reference Style

Target: Remove ALL white/gray/light pixels around the object
Result: Clean edges like the zombie bunny reference
"""

import os
import shutil
from pathlib import Path
from PIL import Image
import numpy as np

# Test images
TEST_IMAGES = [
    "assets/PHASE_PACKS/1_FAZA_1/tools/wood/watering_can.png",
    "assets/PHASE_PACKS/1_FAZA_1/animals/horse.png", 
    "assets/PHASE_PACKS/1_FAZA_1/infrastructure/farm_elements/manure_pile.png",
    "assets/PHASE_PACKS/1_FAZA_1/tools/iron/pickaxe.png",
    "assets/PHASE_PACKS/1_FAZA_1/animals/sheep/walk.png",
]

OUTPUT_DIR = "test_transparency"

def aggressive_bg_removal(input_path, output_path):
    """
    Aggressively remove ALL light/white backgrounds.
    Also removes white shadows and glows.
    """
    img = Image.open(input_path)
    
    if img.mode != 'RGBA':
        img = img.convert('RGBA')
    
    data = np.array(img).astype(np.float32)
    r, g, b, a = data[:,:,0], data[:,:,1], data[:,:,2], data[:,:,3]
    
    # Calculate brightness (0-255)
    brightness = (r + g + b) / 3
    
    # Calculate saturation (how colorful)
    max_rgb = np.maximum(np.maximum(r, g), b)
    min_rgb = np.minimum(np.minimum(r, g), b)
    saturation = max_rgb - min_rgb
    
    # ===== AGGRESSIVE RULES =====
    
    # 1. Remove pure white (threshold 250+)
    pure_white = (r > 250) & (g > 250) & (b > 250)
    
    # 2. Remove pure black (threshold 5-)
    pure_black = (r < 5) & (g < 5) & (b < 5)
    
    # 3. Remove near-white with low saturation (the shadows!)
    # This catches white glows and shadows
    near_white_low_sat = (brightness > 230) & (saturation < 30)
    
    # 4. Remove light gray (common in anti-aliased edges to white bg)
    light_gray = (brightness > 200) & (saturation < 20)
    
    # 5. Remove checkered pattern pixels (if present)
    # Checkered = alternating gray shades
    checkered_light = (brightness > 180) & (brightness < 220) & (saturation < 15)
    checkered_dark = (brightness > 100) & (brightness < 160) & (saturation < 15)
    
    # Combine all background criteria
    is_background = (
        pure_white | 
        pure_black | 
        near_white_low_sat | 
        light_gray |
        checkered_light |
        checkered_dark
    )
    
    # ===== PROTECT COLORED PIXELS =====
    # Don't remove pixels with high saturation (actual colored parts)
    is_colored = saturation > 40
    is_background = is_background & ~is_colored
    
    # ===== APPLY TRANSPARENCY =====
    new_alpha = np.where(is_background, 0, a)
    
    # For semi-background (transitional), use graduated alpha
    # based on how "background-like" the pixel is
    semi_bg = (brightness > 180) & (saturation < 50) & ~is_colored
    fade_factor = np.clip((brightness - 180) / 70, 0, 1)  # 0-1 scale
    new_alpha = np.where(semi_bg, new_alpha * (1 - fade_factor * 0.8), new_alpha)
    
    # Apply new alpha
    data[:,:,3] = np.clip(new_alpha, 0, 255)
    
    # Clean up transparent pixels (set RGB to 0)
    fully_transparent = data[:,:,3] < 10
    data[fully_transparent, 0] = 0
    data[fully_transparent, 1] = 0
    data[fully_transparent, 2] = 0
    data[fully_transparent, 3] = 0
    
    result = Image.fromarray(data.astype(np.uint8))
    result.save(output_path, 'PNG')
    print(f"  ✅ Processed: {os.path.basename(output_path)}")

def main():
    print("🔥 AGGRESSIVE BG REMOVAL - Match Green Screen Style")
    print("=" * 55)
    print("\nTarget: Clean edges like the zombie bunny reference!")
    print("Removing: white, shadows, glows, checkered patterns\n")
    
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    
    processed = []
    
    for img_path in TEST_IMAGES:
        if not os.path.exists(img_path):
            print(f"  ❌ Not found: {img_path}")
            continue
        
        name = os.path.basename(img_path)
        name_no_ext = os.path.splitext(name)[0]
        
        # Copy original
        orig_dest = os.path.join(OUTPUT_DIR, f"{name_no_ext}_ORIGINAL.png")
        shutil.copy(img_path, orig_dest)
        print(f"  📋 Original: {name}")
        
        # Process with aggressive removal
        proc_dest = os.path.join(OUTPUT_DIR, f"{name_no_ext}_CLEAN.png")
        aggressive_bg_removal(img_path, proc_dest)
        
        processed.append({
            'name': name_no_ext,
            'original': f"{name_no_ext}_ORIGINAL.png",
            'processed': f"{name_no_ext}_CLEAN.png"
        })
    
    # Generate comparison HTML
    html = '''<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <title>🔥 Aggressive BG Removal Test</title>
    <style>
        body {
            background: #1a1a2e;
            color: white;
            font-family: Arial, sans-serif;
            padding: 30px;
        }
        h1 { color: #ff4444; text-align: center; }
        .ref-note {
            background: #00ff00;
            color: black;
            padding: 20px;
            border-radius: 10px;
            text-align: center;
            margin-bottom: 30px;
            font-weight: bold;
        }
        .comparison {
            display: flex;
            gap: 40px;
            margin: 30px 0;
            padding: 20px;
            background: rgba(0,0,0,0.3);
            border-radius: 15px;
            align-items: center;
        }
        .image-box {
            flex: 1;
            text-align: center;
        }
        .image-box h3 { color: #ff6666; margin-bottom: 15px; }
        .image-box img {
            max-width: 300px;
            max-height: 300px;
            border-radius: 10px;
        }
        /* Multiple background options to test */
        .bg-checker {
            background: linear-gradient(45deg, #444 25%, #666 25%, #666 50%, #444 50%, #444 75%, #666 75%);
            background-size: 20px 20px;
            padding: 20px;
            border-radius: 10px;
        }
        .bg-green {
            background: #00ff00;
            padding: 20px;
            border-radius: 10px;
        }
        .bg-game {
            background: linear-gradient(to bottom, #2d5016, #1a3009);
            padding: 20px;
            border-radius: 10px;
        }
        .label { margin-top: 10px; font-size: 0.85em; color: #aaa; }
        .success { color: #44ff44; }
    </style>
</head>
<body>
    <h1>🔥 Aggressive Background Removal Test</h1>
    <div class="ref-note">
        🎯 TARGET: Clean edges like the zombie bunny reference (no white shadows!)
    </div>
'''
    
    for item in processed:
        html += f'''
    <div class="comparison">
        <div class="image-box">
            <h3>📷 ORIGINAL</h3>
            <div class="bg-checker">
                <img src="{item['original']}" alt="Original">
            </div>
            <div class="label">{item['name']}</div>
        </div>
        <div class="image-box">
            <h3 class="success">✅ CLEANED (Checker)</h3>
            <div class="bg-checker">
                <img src="{item['processed']}" alt="Cleaned">
            </div>
        </div>
        <div class="image-box">
            <h3 class="success">✅ CLEANED (Green)</h3>
            <div class="bg-green">
                <img src="{item['processed']}" alt="On Green">
            </div>
        </div>
        <div class="image-box">
            <h3 class="success">✅ CLEANED (Game)</h3>
            <div class="bg-game">
                <img src="{item['processed']}" alt="In Game">
            </div>
        </div>
    </div>
'''
    
    html += '''
</body>
</html>
'''
    
    html_path = os.path.join(OUTPUT_DIR, "comparison_v2.html")
    with open(html_path, 'w') as f:
        f.write(html)
    
    print("\n" + "=" * 55)
    print(f"✅ DONE! Open: {html_path}")

if __name__ == '__main__':
    main()