#!/usr/bin/env python3
"""
Composite the FF symbol onto astronaut suit as realistic worn marking.
- Extracts only the black chevron paths (white = transparent)
- Uses multiply blend so it picks up ALL existing lighting and texture
- Adds grain, blur, slight warp for fabric feel
"""
from PIL import Image, ImageFilter, ImageChops
import numpy as np

def composite_patch(astronaut_path, symbol_path, output_path,
                    chest_x_frac=0.38, chest_y_frac=0.46,
                    patch_w_frac=0.11,
                    rotation_deg=-3,
                    opacity=0.72):

    astro = Image.open(astronaut_path).convert("RGB")
    aw, ah = astro.size

    symbol = Image.open(symbol_path).convert("RGBA")

    # ── Extract only the black/dark pixels as alpha mask ──────────
    # White background → transparent, black chevrons → opaque
    sym_rgb = np.array(symbol.convert("RGB")).astype(float)
    # Darkness of each pixel becomes its alpha (0=white=transparent, 255=black=opaque)
    darkness = 255 - (sym_rgb[..., 0] * 0.299 + sym_rgb[..., 1] * 0.587 + sym_rgb[..., 2] * 0.114)
    darkness = np.clip(darkness, 0, 255)

    # Only keep pixels that are clearly dark
    alpha_mask = np.where(darkness > 80, darkness, 0).astype(np.uint8)

    # Create RGBA symbol with extracted alpha
    sym_clean = Image.fromarray(
        np.dstack([np.zeros_like(alpha_mask), np.zeros_like(alpha_mask),
                   np.zeros_like(alpha_mask), alpha_mask]).astype(np.uint8)
    )

    # ── Size and transform ─────────────────────────────────────────
    patch_w = int(aw * patch_w_frac)
    patch_h = int(patch_w * (symbol.height / symbol.width))
    sym_clean = sym_clean.resize((patch_w, patch_h), Image.LANCZOS)

    # Slight rotation
    sym_clean = sym_clean.rotate(rotation_deg, expand=True, resample=Image.BICUBIC)

    # Subtle shear to simulate suit curvature
    sw, sh = sym_clean.size
    sym_clean = sym_clean.transform(
        (sw, sh), Image.AFFINE,
        (1, 0.025, -sw * 0.012, 0.005, 1, -sh * 0.005),
        resample=Image.BICUBIC
    )

    # Soften edges so it blends into fabric weave
    r_ch, g_ch, b_ch, a_ch = sym_clean.split()
    a_soft = a_ch.filter(ImageFilter.GaussianBlur(radius=1.0))
    sym_clean = Image.merge("RGBA", (r_ch, g_ch, b_ch, a_soft))

    # ── Add grain to match 70mm film texture ──────────────────────
    sw2, sh2 = sym_clean.size
    a_arr = np.array(a_soft).astype(float)
    grain = np.random.normal(0, 12, (sh2, sw2))
    a_arr = np.clip(a_arr + grain, 0, 255)
    # Also fade slightly to look worn
    a_arr = a_arr * 0.80
    a_worn = Image.fromarray(a_arr.astype(np.uint8))
    sym_clean = Image.merge("RGBA", (r_ch, g_ch, b_ch, a_worn))

    # ── Position on chest ─────────────────────────────────────────
    px = int(aw * chest_x_frac) - sym_clean.width // 2
    py = int(ah * chest_y_frac) - sym_clean.height // 2

    # ── Composite using multiply-style blend ──────────────────────
    # Multiply: result = base * (1 - alpha * opacity) + base * 0 * alpha * opacity
    # Since symbol is black (0), multiplying makes the suit darker in those areas
    # = base * (1 - alpha_normalized * opacity)  — this darkens base by symbol shape

    astro_arr = np.array(astro).astype(float)
    alpha_full = np.zeros((ah, aw), dtype=float)

    # Place the alpha into the full canvas
    sym_a_arr = np.array(sym_clean.split()[3]).astype(float) / 255.0 * opacity
    y0 = max(0, py); x0 = max(0, px)
    y1 = min(ah, py + sym_a_arr.shape[0])
    x1 = min(aw, px + sym_a_arr.shape[1])
    sy0 = y0 - py; sy1 = sy0 + (y1 - y0)
    sx0 = x0 - px; sx1 = sx0 + (x1 - x0)
    alpha_full[y0:y1, x0:x1] = sym_a_arr[sy0:sy1, sx0:sx1]

    # Darken the suit by the symbol shape (multiply with near-zero)
    # This makes the symbol look like dark embroidery on the suit
    darken_factor = 1.0 - alpha_full[..., np.newaxis] * 0.75
    result_arr = np.clip(astro_arr * darken_factor, 0, 255).astype(np.uint8)

    result = Image.fromarray(result_arr)
    result.save(output_path, "PNG")
    print(f"Saved: {output_path} ({aw}x{ah})")


if __name__ == "__main__":
    BASE = "/root/.openclaw/workspace/work/pitches/future-forward/generator"
    composite_patch(
        f"{BASE}/clean-chest-01.png",
        f"{BASE}/ff-symbol.png",
        f"{BASE}/patch-10-composited.png",
        chest_x_frac=0.38,
        chest_y_frac=0.46,
        patch_w_frac=0.11,
        rotation_deg=-3,
        opacity=0.75
    )