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

def create_color_ghost(cover, reveal):
    if cover is None or reveal is None:
        return None

    # 1. Ensure images are RGB
    img1_raw = cover.convert('RGB')
    img2_raw = reveal.convert('RGB')

    # 2. Resize reveal to match cover
    if img1_raw.size != img2_raw.size:
        img2_raw = img2_raw.resize(img1_raw.size, Image.LANCZOS)

    # 3. Convert to Float Arrays (0.0 - 1.0)
    # This time we keep the 3 color channels (Height, Width, 3)
    rgb1 = np.array(img1_raw, dtype=float) / 255.0
    rgb2 = np.array(img2_raw, dtype=float) / 255.0

    # 4. Create Grayscale copies just to calculate the Alpha Channel
    # We use the luminance of the pixels to decide transparency
    gray1 = np.mean(rgb1, axis=2) 
    gray2 = np.mean(rgb2, axis=2)

    # 5. Squash the contrast (Math Trick)
    # We squash the target values so they don't clip when we combine them
    rgb1 = rgb1 * 0.5
    rgb2 = 0.5 + (rgb2 * 0.5)
    
    # We also squash the grayscale versions for the Alpha calculation
    gray1 = gray1 * 0.5
    gray2 = 0.5 + (gray2 * 0.5)

    # 6. Calculate the Alpha Channel based on Luminance
    # Formula: Alpha = 1 - Light_Image + Dark_Image
    alpha = 1.0 - gray2 + gray1
    
    # Clip alpha to avoid errors
    alpha = np.clip(alpha, 0.001, 0.999)

    # 7. Reconstruct the Color Output
    # We define the pixel color based on the "Dark Background" version (rgb1)
    # P = Target / Alpha
    
    # We have to reshape alpha to (H, W, 1) so we can divide the RGB array (H, W, 3) by it
    alpha_3d = alpha[:, :, np.newaxis]
    
    out_rgb = np.divide(rgb1, alpha_3d, out=np.zeros_like(rgb1), where=alpha_3d > 0.001)

    # 8. Convert back to 0-255 Integers
    final_rgb_int = (out_rgb * 255).clip(0, 255).astype(np.uint8)
    final_alpha_int = (alpha * 255).clip(0, 255).astype(np.uint8)

    # 9. Stack them into an RGBA Image
    # Concatenate the Color channels with the Alpha channel
    rgba_data = np.dstack((final_rgb_int, final_alpha_int))

    return Image.fromarray(rgba_data, 'RGBA')

# --- Gradio UI ---
with gr.Blocks(title="Color Ghost Generator") as demo:
    gr.Markdown("# 👻 Color Ghost Image Generator")
    gr.Markdown("Create a magic image that changes when clicked (Supports Color).")
    
    with gr.Row():
        with gr.Column():
            img_dark = gr.Image(type="pil", label="Image seen in DARK Mode (Cover)")
            img_light = gr.Image(type="pil", label="Image seen in LIGHT Mode (Reveal)")
            btn = gr.Button("Generate Ghost", variant="primary")
        
        with gr.Column():
            output_img = gr.Image(label="Result (Download to test)", type="pil")

    btn.click(fn=create_color_ghost, inputs=[img_dark, img_light], outputs=output_img)

demo.launch()