import gradio as gr
from PIL import Image
import numpy as np

def combine_opacity_layers(top_img, bg_img, invert_opacity=False):
    if top_img is None or bg_img is None:
        return None

    # Convert both to RGB and resize
    top = top_img.convert("RGB")
    bg = bg_img.convert("RGBA").resize(top.size)
    top_np = np.array(top).astype(np.float32)

    # Compute brightness from top image
    brightness = np.dot(top_np[..., :3], [0.299, 0.587, 0.114])
    if invert_opacity:
        brightness = 255 - brightness
    alpha_top = (brightness / 255.0).clip(0, 1)  # shape (H, W)

    # Get alpha from background image
    bg_np = np.array(bg).astype(np.float32) / 255.0
    alpha_bg = bg_np[..., 3]  # shape (H, W)

    # Multiply alpha channels
    combined_alpha = (alpha_top * alpha_bg).clip(0, 1)

    # Create output RGBA: white RGB + combined alpha
    height, width = combined_alpha.shape
    white_rgb = np.ones((height, width, 3), dtype=np.float32)  # all white
    out_rgba = np.dstack((white_rgb, combined_alpha)) * 255
    out_img = Image.fromarray(out_rgba.astype(np.uint8), mode="RGBA")
    return out_img

# Gradio UI
iface = gr.Interface(
    fn=combine_opacity_layers,
    inputs=[
        gr.Image(type="pil", label="Top Brightness-Based Image"),
        gr.Image(type="pil", label="Base Opacity Image (with alpha)"),
        gr.Checkbox(label="Invert Top Brightness", value=False)
    ],
    outputs=gr.Image(type="pil", label="Combined Opacity Image"),
    title="🔁 Combine Two Opacity Images (Multiply Alphas)",
    description=(
        "Takes two images and multiplies their opacity (alpha). Output is solid white "
        "with combined alpha values. Useful for stacking brightness masks."
    )
)

if __name__ == "__main__":
    iface.launch()